Patch Segment head and downgrade to ultralytics 7.3.23

backend/inference.py

@@ -4,17 +4,31 @@
 import os
 import cv2
 import numpy as np
+import torch
 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 = None
+
 try:
+    print(f"🔄 Loading Acetowhite model from: {AW_MODEL_PATH}")
     aw_model = YOLO(AW_MODEL_PATH)
-    aw_model.to('cpu')  # Force CPU device
-    print("✅ Acetowhite model loaded from:", AW_MODEL_PATH)
+    aw_model.to('cpu')
+    
+    # Patch the Segment head to prevent the detect() error
+    if hasattr(aw_model.model, 'model'):
+        for module in aw_model.model.modules():
+            if module.__class__.__name__ == 'Segment':
+                print("⚠️  Patching Segment head to prevent detect() error")
+                # Disable the problematic detect call
+                module.detect = lambda self, x: x
+    
+    print("✅ Acetowhite model loaded successfully")
 except Exception as e:
     print(f"❌ Error loading Acetowhite model: {e}")
-    aw_model = None
+    import traceback
+    traceback.print_exc()
 
 # CONFIGURABLE PARAMETERS
 MIN_AREA = 150            # minimum contour area (px)
@@ -35,7 +49,7 @@ def infer_aw_contour(frame, conf_threshold=DEFAULT_CONF):
         }
 
     if aw_model is None:
-        print("❌ Acetowhite model not loaded")
+        print("❌ Acetowhite model not available")
         return {
             "overlay": None,
             "contours": [],
@@ -44,6 +58,10 @@ def infer_aw_contour(frame, conf_threshold=DEFAULT_CONF):
             "frame_height": frame.shape[0]
         }
 
+    overlay = frame.copy()
+    contours_list = []
+    detection_count = 0
+
     try:
         print(f"🔄 Running YOLO prediction on frame shape: {frame.shape}")
         results = aw_model.predict(
@@ -52,72 +70,55 @@ def infer_aw_contour(frame, conf_threshold=DEFAULT_CONF):
             imgsz=IMG_SIZE,
             verbose=False,
             device='cpu'
-        )[0]
+        )
+        
+        # Handle both list and single result
+        if isinstance(results, (list, tuple)):
+            result = results[0]
+        else:
+            result = results
+            
         print(f"✅ YOLO prediction complete")
+        
+        # Try to extract masks if available
+        if hasattr(result, 'masks') and result.masks is not None:
+            try:
+                masks = result.masks.xy
+                if len(masks) > 0:
+                    print(f"✅ Found {len(masks)} masks")
+                    for idx, polygon in enumerate(masks):
+                        confidence = float(result.boxes.conf[idx])
+                        if confidence < conf_threshold:
+                            continue
+
+                        contour = polygon.astype(np.int32)
+                        area = cv2.contourArea(contour)
+
+                        if area < MIN_AREA:
+                            continue
+
+                        epsilon = SMOOTHING_EPSILON * cv2.arcLength(contour, True)
+                        contour = cv2.approxPolyDP(contour, epsilon, True)
+
+                        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)
+                        })
+                        detection_count += 1
+            except Exception as mask_err:
+                print(f"⚠️  Could not extract masks: {mask_err}")
+                
     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
-
-    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):
-
-                confidence = float(results.boxes.conf[idx])
-
-                # Skip low-confidence masks (extra safety layer)
-                if confidence < conf_threshold:
-                    continue
-
-                contour = polygon.astype(np.int32)
-
-                area = cv2.contourArea(contour)
-
-                if area < MIN_AREA:
-                    continue
-
-                # 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
-                )
-
-                contours_list.append({
-                    "points": contour.tolist(),
-                    "area": float(area),
-                    "confidence": round(confidence, 3)
-                })
-
-                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()
-
+        # Continue with empty results rather than crashing
+        
     return {
-        "overlay": overlay,
+        "overlay": overlay if detection_count > 0 else None,
         "contours": contours_list,
         "detections": detection_count,
         "frame_width": frame.shape[1],

backend/requirements.txt

@@ -4,7 +4,7 @@ opencv-python==4.9.0.80
 numpy==1.24.3
 torch==2.2.0
 torchvision==0.17.0
-ultralytics==8.0.227
+ultralytics==7.3.23
 pillow==10.2.0
 python-multipart==0.0.6
 setuptools>=69.0.0