Update app.py

app.py

@@ -1,4 +1,3 @@
-import streamlit as st
 import re
 
 def generate_script_v8(dataset_code, task, model_size, epochs, batch_size):
@@ -16,6 +15,9 @@ def generate_script_v8(dataset_code, task, model_size, epochs, batch_size):
     project_name = project_name_match.group(1)
     version_number = int(version_number_match.group(1))
 
+    # Determine the model type based on the selected task
+    model_type = "seg" if task == "Segmentation" else "cls"
+
     # Generate the script
     script = f"""
 import yaml
@@ -53,16 +55,13 @@ def auto_train():
         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}-{model_type}.pt"
 
         # Load and train the model
@@ -79,7 +78,6 @@ def auto_train():
         # 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
@@ -92,14 +90,12 @@ def auto_train():
         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)
 
@@ -115,7 +111,7 @@ def generate_script_v9(dataset_code, task, model_size, epochs, batch_size):
     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):
+    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)
@@ -132,13 +128,11 @@ def generate_script_v9(dataset_code, task, model_size, epochs, batch_size):
     # 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 \\
@@ -148,7 +142,8 @@ dataset = version.download("yolov9")
     """
     return script
 
-# Streamlit interface
+import streamlit as st
+
 st.title("Auto Train Script Generator")
 st.write("Generate a YOLO training script using a Roboflow dataset")