Fix ultralytics model loading with CPU device and error handling

backend/inference.py

@@ -8,9 +8,13 @@ from ultralytics import YOLO
 
 # MODEL LOAD (Safe Backend Path)
 AW_MODEL_PATH = os.path.join(os.path.dirname(__file__), "models", "AW_yolo.pt")
-aw_model = YOLO(AW_MODEL_PATH)
-
-print("Acetowhite model loaded from:", AW_MODEL_PATH)
+try:
+    aw_model = YOLO(AW_MODEL_PATH)
+    aw_model.to('cpu')  # Force CPU device
+    print("✅ Acetowhite model loaded from:", AW_MODEL_PATH)
+except Exception as e:
+    print(f"❌ Error loading Acetowhite model: {e}")
+    aw_model = None
 
 # CONFIGURABLE PARAMETERS
 MIN_AREA = 150            # minimum contour area (px)
@@ -30,54 +34,87 @@ def infer_aw_contour(frame, conf_threshold=DEFAULT_CONF):
             "frame_height": 0
         }
 
-    results = aw_model.predict(
-        frame,
-        conf=conf_threshold,
-        imgsz=IMG_SIZE,
-        verbose=False
-    )[0]
+    if aw_model is None:
+        print("❌ Acetowhite model not loaded")
+        return {
+            "overlay": None,
+            "contours": [],
+            "detections": 0,
+            "frame_width": frame.shape[1],
+            "frame_height": frame.shape[0]
+        }
+
+    try:
+        print(f"🔄 Running YOLO prediction on frame shape: {frame.shape}")
+        results = aw_model.predict(
+            frame,
+            conf=conf_threshold,
+            imgsz=IMG_SIZE,
+            verbose=False,
+            device='cpu'
+        )[0]
+        print(f"✅ YOLO prediction complete")
+    except Exception as e:
+        print(f"❌ YOLO prediction error: {e}")
+        import traceback
+        traceback.print_exc()
+        return {
+            "overlay": None,
+            "contours": [],
+            "detections": 0,
+            "frame_width": frame.shape[1],
+            "frame_height": frame.shape[0]
+        }
 
     overlay = frame.copy()
     contours_list = []
     detection_count = 0
 
-    if results.masks is not None and len(results.masks.xy) > 0:
+    try:
+        if results.masks is not None and len(results.masks.xy) > 0:
+            print(f"✅ Found masks: {len(results.masks.xy)}")
 
-        for idx, polygon in enumerate(results.masks.xy):
+            for idx, polygon in enumerate(results.masks.xy):
 
-            confidence = float(results.boxes.conf[idx])
+                confidence = float(results.boxes.conf[idx])
 
-            # Skip low-confidence masks (extra safety layer)
-            if confidence < conf_threshold:
-                continue
+                # Skip low-confidence masks (extra safety layer)
+                if confidence < conf_threshold:
+                    continue
 
-            contour = polygon.astype(np.int32)
+                contour = polygon.astype(np.int32)
 
-            area = cv2.contourArea(contour)
+                area = cv2.contourArea(contour)
 
-            if area < MIN_AREA:
-                continue
+                if area < MIN_AREA:
+                    continue
 
-            # Optional smoothing
-            epsilon = SMOOTHING_EPSILON * cv2.arcLength(contour, True)
-            contour = cv2.approxPolyDP(contour, epsilon, True)
+                # Optional smoothing
+                epsilon = SMOOTHING_EPSILON * cv2.arcLength(contour, True)
+                contour = cv2.approxPolyDP(contour, epsilon, True)
 
-            # Draw clean contour
-            cv2.polylines(
-                overlay,
-                [contour],
-                isClosed=True,
-                color=(0, 255, 0),
-                thickness=2
-            )
+                # Draw clean contour
+                cv2.polylines(
+                    overlay,
+                    [contour],
+                    isClosed=True,
+                    color=(0, 255, 0),
+                    thickness=2
+                )
 
-            contours_list.append({
-                "points": contour.tolist(),
-                "area": float(area),
-                "confidence": round(confidence, 3)
-            })
+                contours_list.append({
+                    "points": contour.tolist(),
+                    "area": float(area),
+                    "confidence": round(confidence, 3)
+                })
 
-            detection_count += 1
+                detection_count += 1
+        else:
+            print("ℹ️  No masks found in results")
+    except Exception as e:
+        print(f"❌ Error processing masks: {e}")
+        import traceback
+        traceback.print_exc()
 
     return {
         "overlay": overlay,
@@ -96,10 +133,14 @@ from ultralytics import YOLO
 from collections import deque
 
 cervix_MODEL_PATH = os.path.join(os.path.dirname(__file__), "models", "cervix_yolo.pt")
-cervix_model = YOLO(cervix_MODEL_PATH)
-
-print("Cervix model loaded from:", cervix_MODEL_PATH)
-print("Classes:", cervix_model.model.names)
+try:
+    cervix_model = YOLO(cervix_MODEL_PATH)
+    cervix_model.to('cpu')  # Force CPU device
+    print(f"✅ Cervix model loaded from: {cervix_MODEL_PATH}")
+    print(f"✅ Classes: {cervix_model.model.names}")
+except Exception as e:
+    print(f"❌ Error loading Cervix model: {e}")
+    cervix_model = None
 
 # Stability buffer for video
 detect_history = deque(maxlen=10)
@@ -142,14 +183,35 @@ def analyze_frame(frame, conf_threshold=0.3):
             "quality_percent": 0
         }
 
+    if cervix_model is None:
+        print("❌ Cervix model not loaded")
+        return {
+            "detected": False,
+            "detection_confidence": 0.0,
+            "quality_score": 0.0,
+            "quality_percent": 0
+        }
+
     gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
 
-    results = cervix_model.predict(
-        frame,
-        conf=conf_threshold,
-        imgsz=640,
-        verbose=False
-    )
+    try:
+        results = cervix_model.predict(
+            frame,
+            conf=conf_threshold,
+            imgsz=640,
+            verbose=False,
+            device='cpu'
+        )
+    except Exception as e:
+        print(f"❌ Cervix model prediction error: {e}")
+        import traceback
+        traceback.print_exc()
+        return {
+            "detected": False,
+            "detection_confidence": 0.0,
+            "quality_score": 0.0,
+            "quality_percent": 0
+        }
 
     r = results[0]