Update app.py

app.py

@@ -1,6 +1,7 @@
 import streamlit as st
+import re
 
-def generate_script(dataset_code, task, model_size, epochs, batch_size):
+def generate_script_v8(dataset_code, task, model_size, epochs, batch_size):
     # Extract the necessary information from the dataset code
     api_key_match = re.search(r'api_key="(.*?)"', dataset_code)
     workspace_match = re.search(r'workspace\("([^"]+)"\)', dataset_code)
@@ -107,16 +108,72 @@ if __name__ == "__main__":
     """
     return script
 
-# Streamlit interface
-st.title("Auto Train Script Generator")
-st.write("Generate a YOLOv8 training script using a Roboflow dataset")
+def generate_script_v9(dataset_code, task, model_size, epochs, batch_size):
+    # Extract the necessary information from the dataset code
+    api_key_match = re.search(r'api_key="(.*?)"', dataset_code)
+    workspace_match = re.search(r'workspace\("([^"]+)"\)', dataset_code)
+    project_name_match = re.search(r'project\("([^"]+)"\)', dataset_code)
+    version_number_match = re.search(r'version\((\d+)\)', dataset_code)
 
-dataset_code = st.text_input("Roboflow Dataset Code", placeholder="Paste your Roboflow dataset code here")
-task = st.selectbox("Task", ["Object Detection", "Segmentation"], index=0)
-model_size = st.selectbox("Model Size", ["n", "s", "m", "l", "x"], index=0)
-epochs = st.selectbox("Epochs", [50, 100, 200, 300, 400, 500], index=3)
-batch_size = st.selectbox("Batch Size", [1, 2, 4, 8, 16, 32], index=0)
+    if not (api_key_match and workspace_match and project_name_match and version_number_match):
+        return "Error: Could not extract necessary information from the dataset code."
 
-if st.button("Generate Script"):
-    script = generate_script(dataset_code, task, model_size, epochs, batch_size)
-    st.code(script, language="python")
+    api_key = api_key_match.group(1)
+    workspace = workspace_match.group(1)
+    project_name = project_name_match.group(1)
+    version_number = int(version_number_match.group(1))
+
+    # Determine the model name based on the selected size and task
+    if task == "Segmentation":
+        model_name = f"gelan-c-seg.pt" if model_size == "c" else f"yolov9-{model_size}-seg.pt"
+    else:
+        model_name = f"yolov9-{model_size}.pt"
+
+    # Generate the script
+    script = f"""
+!pip install roboflow
+
+from roboflow import Roboflow
+rf = Roboflow(api_key="{api_key}")
+project = rf.workspace("{workspace}").project("{project_name}")
+version = project.version({version_number})
+dataset = version.download("yolov9")
+
+!python train.py \\
+--batch {batch_size} --epochs {epochs} --img 640 --device 0 --min-items 0 --close-mosaic 15 \\
+--data {{dataset.location}}/data.yaml \\
+--weights {{HOME}}/weights/{model_name} \\
+--cfg models/detect/{model_name.split('.')[0]}.yaml \\
+--hyp hyp.scratch-high.yaml
+    """
+    return script
+
+# Streamlit interface
+st.title("Auto Train Script Generator")
+st.write("Generate a YOLO training script using a Roboflow dataset")
+
+tab1, tab2 = st.tabs(["YOLOv8", "YOLOv9"])
+
+with tab1:
+    st.subheader("YOLOv8 Script Generator")
+    dataset_code = st.text_input("Roboflow Dataset Code", key="dataset_code_v8", placeholder="Paste your Roboflow dataset code here")
+    task = st.selectbox("Task", ["Object Detection", "Segmentation"], index=0, key="task_v8")
+    model_size = st.selectbox("Model Size", ["n", "s", "m", "l", "x"], index=0, key="model_size_v8")
+    epochs = st.selectbox("Epochs", [50, 100, 200, 300, 400, 500], index=3, key="epochs_v8")
+    batch_size = st.selectbox("Batch Size", [1, 2, 4, 8, 16, 32], index=0, key="batch_size_v8")
+
+    if st.button("Generate YOLOv8 Script"):
+        script = generate_script_v8(dataset_code, task, model_size, epochs, batch_size)
+        st.code(script, language="python")
+
+with tab2:
+    st.subheader("YOLOv9 Script Generator")
+    dataset_code = st.text_input("Roboflow Dataset Code", key="dataset_code_v9", placeholder="Paste your Roboflow dataset code here")
+    task = st.selectbox("Task", ["Object Detection", "Segmentation"], index=0, key="task_v9")
+    model_size = st.selectbox("Model Size", ["t", "s", "m", "c", "e"], index=0, key="model_size_v9")
+    epochs = st.selectbox("Epochs", [50, 100, 200, 300, 400, 500], index=3, key="epochs_v9")
+    batch_size = st.selectbox("Batch Size", [1, 2, 4, 8, 16, 32], index=0, key="batch_size_v9")
+
+    if st.button("Generate YOLOv9 Script"):
+        script = generate_script_v9(dataset_code, task, model_size, epochs, batch_size)
+        st.code(script, language="python")