Upload train_ppe_fixed.py

train_ppe_fixed.py

@@ -1,13 +1,24 @@
 #!/usr/bin/env python3
 """
 PPE Compliance Detection Training - FIXED VERSION
+- Swaps opencv-python for opencv-python-headless before importing ultralytics
 - Downloads keremberke dataset as ZIP files (script-based datasets no longer supported)
 - Uses model.train() return value correctly (no results.best)
 - Pushes best.pt to HuggingFace Hub after training
 """
 
-import os
+import subprocess
 import sys
+import os
+
+# FIX: opencv-python needs libGL which is missing in container; use headless instead
+print("[0/5] Swapping opencv-python for opencv-python-headless...")
+subprocess.run([sys.executable, "-m", "pip", "uninstall", "-y", "opencv-python"],
+               capture_output=True)
+subprocess.run([sys.executable, "-m", "pip", "install", "--quiet", "opencv-python-headless"],
+               capture_output=True)
+print("    Done")
+
 import zipfile
 import shutil
 import json
@@ -179,7 +190,7 @@ def merge_datasets(ppe_extract_dir: Path, keremberke_dir: Path, output_dir: Path
 
     if ppe_dir is None:
         print("    ERROR: Could not find PPE dataset structure")
-        sys.exit(1)
+        os._exit(1)
 
     print(f"    Found PPE dataset at: {ppe_dir}")
 
@@ -244,8 +255,6 @@ def train_model(data_yaml_path: Path):
 
     model = YOLO("yolov8s.pt")
 
-    # model.train() returns metrics, NOT the model object
-    # The trained weights are auto-saved to /app/runs/ppe_improved/weights/best.pt
     model.train(
         data=str(data_yaml_path),
         epochs=EPOCHS,