Update main.py

main.py

@@ -16,13 +16,13 @@ import uvicorn
 # OOM PREVENTION 1: Force PyTorch to use minimal memory overhead
 torch.set_num_threads(1)
 
-# Import your helpers (assuming cv_helpers.py is in the same folder)
+# Import your helpers
 from cv_helpers import blend_mask_overlays, stem_tip_tangent_deg
 
 # --- CONFIGURATION ---
 MODEL_PATH = "best.pt"
 PIXELS_TO_CM = 1.0  
-MAX_IMAGE_SIZE = 1024  # OOM PREVENTION 2: Max pixels on the longest side
+MAX_IMAGE_SIZE = 1024  
 
 @dataclass
 class ProcessResult:
@@ -52,6 +52,7 @@ class WatermelonProcessor:
     def get_stable_perimeter_data(rind_mask, flesh_mask):
         cnts, _ = cv2.findContours(rind_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
         if not cnts: return None
+        
         best_cnt = None
         max_overlap = -1
         for cnt in cnts:
@@ -61,6 +62,7 @@ class WatermelonProcessor:
             if overlap_area > max_overlap:
                 max_overlap = overlap_area
                 best_cnt = cnt
+                
         if best_cnt is None: best_cnt = max(cnts, key=cv2.contourArea)
         moments = cv2.moments(best_cnt)
         if moments["m00"] == 0: return None
@@ -84,8 +86,8 @@ class WatermelonProcessor:
         return final_theta, final_r, (cx, cy), best_cnt
 
     @staticmethod
-    def get_ray_scan_midline(flesh_mask, rind_cnt, predicted_cnt, cx, cy):
-        h, w = flesh_mask.shape
+    def get_dual_mask_midline(f_left, f_right, rind_cnt, predicted_cnt, cx, cy):
+        h, w = f_left.shape
         if len(rind_cnt) > 5:
             _, (ma, Ma), angle = cv2.fitEllipse(rind_cnt)
             rot_angle = angle if ma < Ma else angle + 90
@@ -93,32 +95,50 @@ class WatermelonProcessor:
 
         m_rot = cv2.getRotationMatrix2D((cx, cy), rot_angle, 1.0)
         m_inv = cv2.getRotationMatrix2D((cx, cy), -rot_angle, 1.0)
-        f_rot = cv2.warpAffine(flesh_mask, m_rot, (w, h))
+        
+        l_rot = cv2.warpAffine(f_left, m_rot, (w, h))
+        r_rot = cv2.warpAffine(f_right, m_rot, (w, h))
+
+        # Failsafe: Swap left/right if YOLO got labels crossed
+        l_idx = np.where(l_rot > 0)[1]
+        r_idx = np.where(r_rot > 0)[1]
+        if len(l_idx) > 0 and len(r_idx) > 0:
+            if np.mean(l_idx) > np.mean(r_idx):
+                l_rot, r_rot = r_rot, l_rot 
 
         gap_points =[]
-        y_indices, _ = np.where(f_rot > 0)
-        if len(y_indices) > 0:
-            for y in range(np.min(y_indices), np.max(y_indices)):
-                row = f_rot[y, :]
-                white_px = np.where(row > 0)[0]
-                if len(white_px) >= 2:
-                    blanks = np.where(row[white_px[0]:white_px[-1]] == 0)[0] + white_px[0]
-                    if len(blanks) > 0: gap_points.append([y, np.median(blanks)])
+        y_l = np.where(l_rot > 0)[0]
+        y_r = np.where(r_rot > 0)[0]
+        
+        if len(y_l) > 0 and len(y_r) > 0:
+            y_min = max(np.min(y_l), np.min(y_r))
+            y_max = min(np.max(y_l), np.max(y_r))
+            for y in range(y_min, y_max):
+                row_l = np.where(l_rot[y, :] > 0)[0]
+                row_r = np.where(r_rot[y, :] > 0)[0]
+                if len(row_l) > 0 and len(row_r) > 0:
+                    edge_l = row_l[-1]
+                    edge_r = row_r[0]
+                    gap_points.append([y, (edge_l + edge_r) / 2.0])
 
         gap_points = np.array(gap_points)
         if len(gap_points) > 10:
-            y_min, y_max = np.min(gap_points[:, 0]), np.max(gap_points[:, 0])
-            y_span, y_mean = max(y_max - y_min, 1), (y_max + y_min) / 2.0
+            y_min_g, y_max_g = np.min(gap_points[:, 0]), np.max(gap_points[:, 0])
+            y_span = max(y_max_g - y_min_g, 1)
+            y_mean = (y_max_g + y_min_g) / 2.0
             y_norm = (gap_points[:, 0] - y_mean) / y_span
             x_data = gap_points[:, 1]
 
             def parabola(y_n, a, b, c): return a * (y_n**2) + b * y_n + c
+            max_bend = w * 0.08 
             try:
-                popt_mid, _ = curve_fit(parabola, y_norm, x_data, bounds=([-w*0.08, -np.inf, -np.inf],[w*0.08, np.inf, np.inf]))
-            except: popt_mid =[0.0, 0.0, cx]
+                popt_mid, _ = curve_fit(parabola, y_norm, x_data, bounds=([-max_bend, -np.inf, -np.inf],[max_bend, np.inf, np.inf]))
+            except:
+                popt_mid =[0.0, 0.0, cx]
 
             ys_extrap = np.linspace(0, h, 500)
-            xs_extrap = parabola((ys_extrap - y_mean) / y_span, *popt_mid)
+            ys_extrap_norm = (ys_extrap - y_mean) / y_span
+            xs_extrap = parabola(ys_extrap_norm, *popt_mid)
             pts_rot = np.vstack([xs_extrap, ys_extrap, np.ones_like(xs_extrap)])
         else:
             ys_extrap = np.linspace(0, h, 500)
@@ -132,22 +152,35 @@ class WatermelonProcessor:
         if image is None: return ProcessResult(success=False, message="Could not decode image.")
         h, w = image.shape[:2]
         
-        results = self.model(image, conf=0.25, verbose=False)
-        rind_mask, flesh_mask = np.zeros((h, w), dtype=np.uint8), np.zeros((h, w), dtype=np.uint8)
+        # retina_masks=True removes the plateau artifacts during inference!
+        results = self.model(image, conf=0.25, retina_masks=True, verbose=False)
+        
+        rind_mask = np.zeros((h, w), dtype=np.uint8)
+        flesh_l = np.zeros((h, w), dtype=np.uint8)
+        flesh_r = np.zeros((h, w), dtype=np.uint8)
 
         if results[0].masks is None:
             return ProcessResult(success=False, message="No masks detected.")
 
+        # --- THE FIX: Load 0 (Whole), 1 (Left), and 2 (Right) ---
         for mask_data, cls in zip(results[0].masks.xy, results[0].boxes.cls):
             contour = np.array(mask_data, dtype=np.int32)
-            if int(cls) == 0: cv2.drawContours(rind_mask, [contour], -1, 255, -1)
-            elif int(cls) == 1: cv2.drawContours(flesh_mask, [contour], -1, 255, -1)
+            c_id = int(cls)
+            if c_id == 0:
+                cv2.drawContours(rind_mask, [contour], -1, 255, -1)
+            elif c_id == 1:
+                cv2.drawContours(flesh_l, [contour], -1, 255, -1)
+            elif c_id == 2:
+                cv2.drawContours(flesh_r, [contour], -1, 255, -1)
+
+        flesh_combined = cv2.bitwise_or(flesh_l, flesh_r)
 
-        perimeter_data = self.get_stable_perimeter_data(rind_mask, flesh_mask)
+        perimeter_data = self.get_stable_perimeter_data(rind_mask, flesh_combined)
         if perimeter_data is None: return ProcessResult(success=False, message="No stable perimeter.")
 
         t_data, r_raw, (cx, cy), rind_cnt = perimeter_data
         scale = np.mean(r_raw)
+        if scale <= 0: return ProcessResult(success=False, message="Invalid perimeter scale.")
         
         try:
             popt, _ = curve_fit(
@@ -157,35 +190,37 @@ class WatermelonProcessor:
             )
         except Exception as exc: return ProcessResult(success=False, message=f"Fit failed: {exc}")
 
-        r2 = 1 - (np.sum((r_raw / scale - self.watermelon_model(t_data, *popt)) ** 2) / np.sum((r_raw / scale - 1) ** 2))
-        
+        denom = np.sum((r_raw / scale - 1) ** 2)
+        if denom == 0: return ProcessResult(success=False, message="R2 denominator became zero.")
+
+        r2 = 1 - (np.sum((r_raw / scale - self.watermelon_model(t_data, *popt)) ** 2) / denom)
         t_fit = np.linspace(-np.pi, np.pi, 500)
         r_fit = self.watermelon_model(t_fit, *popt) * scale
         fit_pts = np.array([[r * np.cos(t) + cx, cy - r * np.sin(t)] for t, r in zip(t_fit, r_fit)])
         
-        # --- FEATURE EXTRACTION (WITH TRUE-SIZE CORRECTION) ---
+        orig_scale = 1.0 / scale_ratio
         width_px = float(np.max(fit_pts[:, 0]) - np.min(fit_pts[:, 0]))
         height_px = float(np.max(fit_pts[:, 1]) - np.min(fit_pts[:, 1]))
         diffs = np.diff(fit_pts, axis=0)
         perimeter_px = float(np.sum(np.linalg.norm(diffs, axis=1)) + np.linalg.norm(fit_pts[-1] - fit_pts[0]))
 
-        # We divide by scale_ratio to perfectly undo the downscaling for measurements!
-        orig_scale = 1.0 / scale_ratio
         width_val = width_px * PIXELS_TO_CM * orig_scale
         height_val = height_px * PIXELS_TO_CM * orig_scale
         perimeter_val = perimeter_px * PIXELS_TO_CM * orig_scale
 
-        # --- DRAWING ---
-        midline = self.get_ray_scan_midline(flesh_mask, rind_cnt, fit_pts, cx, cy)
-        output = blend_mask_overlays(image, rind_mask, flesh_mask)
-        if len(midline) > 1: cv2.polylines(output,[midline.astype(np.int32)], False, (0, 255, 255), 3)
-        cv2.polylines(output,[fit_pts.astype(np.int32)], True, (0, 255, 0), 3)
+        # --- THE FIX: Dual-Mask Midline ---
+        midline = self.get_dual_mask_midline(flesh_l, flesh_r, rind_cnt, fit_pts, cx, cy)
+        
+        output = blend_mask_overlays(image, rind_mask, flesh_combined)
+        if len(midline) > 1:
+            cv2.polylines(output, [midline.astype(np.int32)], False, (0, 255, 255), 3)
+        cv2.polylines(output, [fit_pts.astype(np.int32)], True, (0, 255, 0), 3)
 
         stem = stem_tip_tangent_deg(rind_cnt, (cx, cy))
         if stem is not None:
             tx, ty, tdeg = stem
-            L = min(w, h) * 0.08
             rad = np.deg2rad(tdeg)
+            L = min(w, h) * 0.08
             p1 = (int(round(tx)), int(round(ty)))
             p2 = (int(round(tx + L * np.cos(rad))), int(round(ty + L * np.sin(rad))))
             cv2.circle(output, p1, 6, (255, 0, 255), -1)
@@ -224,7 +259,6 @@ async def process_single(file: UploadFile = File(...)):
     nparr = np.frombuffer(contents, np.uint8)
     img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
     
-    # OOM PREVENTION 3: Resize image if it's massive
     h, w = img.shape[:2]
     scale_ratio = 1.0
     if max(h, w) > MAX_IMAGE_SIZE:
@@ -233,8 +267,10 @@ async def process_single(file: UploadFile = File(...)):
 
     res = processor.process_image(img, file.filename, scale_ratio)
     
-    # OOM PREVENTION 4: Force garbage collection immediately after processing
     del img, nparr, contents
     gc.collect()
 
-    return res.__dict__
+    return res.__dict__
+
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=7860)