Create train_yolo.py

safety_analyzer/train_yolo.py

@@ -0,0 +1,50 @@
+```python
+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():
+    try:
+        # Load pretrained YOLOv8 model
+        model = YOLO("yolov8n.pt")
+        logger.info("Loaded pretrained model: yolov8n.pt")
+
+        # 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
+        )
+        logger.info("Training completed. Model saved in runs/train/safety_model/weights/")
+
+        # Validate the model
+        model.val()
+        logger.info("Validation completed.")
+
+        # 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()
+```