Create app.py

app.py

@@ -0,0 +1,122 @@
+import streamlit as st
+
+def generate_script(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)
+
+    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."
+
+    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))
+
+    # Generate the script
+    script = f"""
+import yaml
+from ultralytics import YOLO
+from roboflow import Roboflow
+import logging
+import re
+import threading
+import time
+from io import StringIO
+
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+def auto_train():
+    log_stream = StringIO()
+    log_handler = logging.StreamHandler(log_stream)
+    log_handler.setLevel(logging.INFO)
+    logger.addHandler(log_handler)
+
+    try:
+        api_key = "{api_key}"
+        workspace = "{workspace}"
+        project_name = "{project_name}"
+        version_number = {version_number}
+
+        # Load the Roboflow dataset
+        rf = Roboflow(api_key=api_key)
+        project = rf.workspace(workspace).project(project_name)
+        version = project.version(version_number)
+        dataset = version.download("yolov8")
+
+        # Modify the data structure
+        yaml_file_path = f'{{dataset.location}}/data.yaml'
+        with open(yaml_file_path, 'r') as file:
+            data = yaml.safe_load(file)
+
+        data['val'] = '../valid/images'
+        data['test'] = '../test/images'
+        data['train'] = '../train/images'
+
+        with open(yaml_file_path, 'w') as file:
+            yaml.safe_dump(data, file)
+
+        # Determine the model name based on the selected size and task
+        model_type = "seg" if task == "Segmentation" else "cls"
+        model_name = f"yolov8{model_size[0]}-{model_type}.pt"
+
+        # Load and train the model
+        model = YOLO(model_name)
+        model.info()
+
+        # Function to read logs in real-time and update the Streamlit textbox
+        def update_logs():
+            while getattr(threading.currentThread(), "do_run", True):
+                time.sleep(1)
+                log_stream.seek(0)
+                print(log_stream.read())
+
+        # Start a thread to update logs in real-time
+        log_thread = threading.Thread(target=update_logs)
+        log_thread.start()
+
+        results = model.train(data=yaml_file_path, epochs={epochs}, imgsz=640, batch={batch_size})
+
+        # Stop the log update thread
+        logger.removeHandler(log_handler)
+        log_thread.do_run = False
+        log_thread.join()
+
+        # Return the result path and logs
+        log_stream.seek(0)
+        log_output = log_stream.read()
+        print("Results Directory:", results.results_dir)
+        print("Final Training Logs:", log_output)
+
+    except Exception as e:
+        logger.error(f"An error occurred: {{e}}")
+        log_stream.seek(0)
+        log_output = log_stream.read()
+        print(f"Error: {{e}}")
+        print(log_output)
+
+    finally:
+        logger.removeHandler(log_handler)
+
+if __name__ == "__main__":
+    auto_train()
+    """
+    return script
+
+# Streamlit interface
+st.title("Auto Train Script Generator")
+st.write("Generate a YOLOv8 training script using a Roboflow dataset")
+
+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 st.button("Generate Script"):
+    script = generate_script(dataset_code, task, model_size, epochs, batch_size)
+    st.code(script, language="python")