Update main.py

main.py

@@ -17,7 +17,7 @@ from skimage import color
 # OOM PREVENTION
 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 ---
@@ -194,7 +194,6 @@ class WatermelonProcessor:
                 popt_mid, _ = curve_fit(parabola, (gap_points[:,0]-y_mean)/y_span, gap_points[:,1], bounds=([-w*0.08, -np.inf, -np.inf],[w*0.08, np.inf, np.inf]))
             except: popt_mid = [0.0, 0.0, cx]
             
-            # FIXED TYPO HERE (ys_extrap)
             ys_extrap = np.linspace(0, h, 500)
             xs_extrap = parabola((ys_extrap - y_mean)/y_span, *popt_mid)
             pts_rot = np.vstack([xs_extrap, ys_extrap, np.ones_like(ys_extrap)])
@@ -210,6 +209,7 @@ class WatermelonProcessor:
         if image is None: return ProcessResult(success=False, message="Could not decode image.")
         h, w = image.shape[:2]
 
+        # --- 1. CALIBRATION & SCALING ---
         dE_initial, dE_final, cm_per_px, checker_corners = None, None, None, None
         
         try:
@@ -229,8 +229,9 @@ class WatermelonProcessor:
         except Exception as e:
             print(f"Calibration skipped for {source_name}: {e}")
         
+        # --- 2. YOLO INFERENCE ---
         results = self.model(image, conf=0.25, retina_masks=True, verbose=False)
-        rind_mask, flesh_l, flesh_r = np.zeros((h, w), dtype=np.uint8), np.zeros((h, w), dtype=np.uint8), np.zeros((h, w), dtype=np.uint8)
+        rind_mask, flesh_contours = np.zeros((h, w), dtype=np.uint8), []
 
         if results[0].masks is None:
             return ProcessResult(success=False, message="No masks detected.")
@@ -239,10 +240,18 @@ class WatermelonProcessor:
             contour = np.array(mask_data, dtype=np.int32)
             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)
+            elif c_id == 1: flesh_contours.append(contour)
 
-        flesh_combined = cv2.bitwise_or(flesh_l, flesh_r)
+        flesh_contours.sort(key=lambda cnt: cv2.moments(cnt)['m10'] / (cv2.moments(cnt)['m00'] + 1e-5))
+        flesh_l_m, flesh_r_m = np.zeros((h, w), dtype=np.uint8), np.zeros((h, w), dtype=np.uint8)
+        
+        if len(flesh_contours) >= 2:
+            cv2.drawContours(flesh_l_m, [flesh_contours[0]], -1, 255, -1)
+            cv2.drawContours(flesh_r_m,[flesh_contours[1]], -1, 255, -1)
+        elif len(flesh_contours) == 1:
+            cv2.drawContours(flesh_l_m,[flesh_contours[0]], -1, 255, -1)
+
+        flesh_combined = cv2.bitwise_or(flesh_l_m, flesh_r_m)
         perimeter_data = self.get_stable_perimeter_data(rind_mask, flesh_combined)
         if perimeter_data is None: return ProcessResult(success=False, message="No stable perimeter.")
 
@@ -263,6 +272,7 @@ class WatermelonProcessor:
         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)])
         
+        # --- 3. DIMENSIONAL MATH (IN CM) ---
         orig_scale = 1.0 / scale_ratio
         if cm_per_px is None: cm_per_px = 1.0 
         
@@ -270,24 +280,26 @@ class WatermelonProcessor:
         height_px = float(np.max(fit_pts[:, 1]) - np.min(fit_pts[:, 1]))
         perimeter_px = float(np.sum(np.linalg.norm(np.diff(fit_pts, axis=0), axis=1)) + np.linalg.norm(fit_pts[-1] - fit_pts[0]))
 
-        width_val = width_px * cm_per_px * orig_scale
-        height_val = height_px * cm_per_px * orig_scale
-        perimeter_val = perimeter_px * cm_per_px * orig_scale
+        # CRITICAL FIX: Wrapped in float() to prevent Numpy JSON serialization errors
+        width_val = float(width_px * cm_per_px * orig_scale)
+        height_val = float(height_px * cm_per_px * orig_scale)
+        perimeter_val = float(perimeter_px * cm_per_px * orig_scale)
 
-        midline = self.get_dual_mask_midline(flesh_l, flesh_r, rind_cnt, fit_pts, cx, cy)
+        # --- 4. DRAWING ---
+        midline = self.get_dual_mask_midline(flesh_l_m, flesh_r_m, rind_cnt, fit_pts, cx, cy)
         output = blend_mask_overlays(image, rind_mask, flesh_combined)
         
         if checker_corners is not None:
-            cv2.polylines(output, [np.int32(checker_corners)], True, (0, 165, 255), 4)
+            cv2.polylines(output,[np.int32(checker_corners)], True, (0, 165, 255), 4)
 
-        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)
+        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)
@@ -317,18 +329,17 @@ def read_root(): return {"status": "Watermelon API is awake!"}
 @app.post("/process_single")
 async def process_single(file: UploadFile = File(...)):
     contents = await file.read()
-    nparr = np.frombuffer(contents, np.uint8)
-    img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
+    img = cv2.imdecode(np.frombuffer(contents, np.uint8), cv2.IMREAD_COLOR)
     
-    h, w = img.shape[:2]
     scale_ratio = 1.0
+    h, w = img.shape[:2]
     if max(h, w) > MAX_IMAGE_SIZE:
         scale_ratio = MAX_IMAGE_SIZE / float(max(h, w))
         img = cv2.resize(img, (int(w * scale_ratio), int(h * scale_ratio)), interpolation=cv2.INTER_AREA)
 
     res = processor.process_image(img, file.filename, scale_ratio)
     
-    del img, nparr, contents
+    del img, contents
     gc.collect()
 
     return res.__dict__