Update safety_analyzer/train_yolo.py

safety_analyzer/train_yolo.py

@@ -1,48 +1,38 @@
 from ultralytics import YOLO
 import logging
-import torch
-import shutil
-import os
 
 # Setup logging
 logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 logger = logging.getLogger(__name__)
 
-def train_model():
+def train_yolov8():
     try:
-        # Load pretrained YOLOv8 model
+        # Load a pretrained YOLOv8 model
         model = YOLO("yolov8n.pt")
-        logger.info("Loaded pretrained model: yolov8n.pt")
+        logger.info("Loaded pretrained YOLOv8n model")
 
         # Train the model
         model.train(
-            data="data.yaml",
-            epochs=50,
-            imgsz=640,
-            batch=16,
-            device=0 if torch.cuda.is_available() else "cpu",
-            name="safety_model",
-            project="runs/train",
-            patience=10
+            data="path/to/data.yaml",  # Path to your data.yaml file
+            epochs=100,                # Number of epochs
+            imgsz=640,                 # Image size
+            batch=16,                  # Batch size
+            device=0,                  # Use GPU (0) if available
+            patience=50,               # Early stopping patience
+            project="runs/train",      # Output directory
+            name="safety_model",       # Experiment name
+            exist_ok=True
         )
-        logger.info("Training completed. Model saved in runs/train/safety_model/weights/")
+        logger.info("Model training completed")
 
-        # Validate the model
-        model.val()
-        logger.info("Validation completed.")
+        # Save the trained model
+        model.save("yolov8_safety.pt")
+        logger.info("Saved trained model as yolov8_safety.pt")
 
-        # Copy the best model to the project root as yolov8_safety.pt
-        best_model_path = "runs/train/safety_model/weights/best.pt"
-        if os.path.exists(best_model_path):
-            shutil.copy(best_model_path, "yolov8_safety.pt")
-            logger.info("Model copied to yolov8_safety.pt")
-        else:
-            logger.error("Best model not found at expected path")
-            raise FileNotFoundError("Best model not found")
     except Exception as e:
         logger.error(f"Error during training: {e}")
         raise
 
 if __name__ == "__main__":
     logger.info("Starting YOLOv8 model training...")
-    train_model()
+    train_yolov8()