Update main.py

main.py

@@ -17,7 +17,7 @@ from skimage import color
 # OOM PREVENTION
 torch.set_num_threads(1)
 
-# Import your helpers
+# Import your helpers (assuming cv_helpers.py is in the same folder)
 from cv_helpers import blend_mask_overlays, stem_tip_tangent_deg
 
 # --- CONFIGURATION ---
@@ -26,7 +26,7 @@ MAX_IMAGE_SIZE = 2048
 CHECKER_WIDTH_CM = 6.3 
 
 # ==============================================================================
-# --- COLOR CALIBRATION LOGIC ---
+# --- COLOR CALIBRATION LOGIC (Embedded) ---
 # ==============================================================================
 def to_linear_srgb(u8_bgr):
     rgb = cv2.cvtColor(u8_bgr, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
@@ -177,6 +177,12 @@ class WatermelonProcessor:
         m_inv = cv2.getRotationMatrix2D((cx, cy), -rot_angle, 1.0)
         l_rot, r_rot = cv2.warpAffine(f_left, m_rot, (w, h)), cv2.warpAffine(f_right, m_rot, (w, h))
 
+        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_l, y_r = np.where(l_rot > 0)[0], np.where(r_rot > 0)[0]
         if len(y_l) > 0 and len(y_r) > 0:
@@ -187,8 +193,8 @@ class WatermelonProcessor:
 
         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, y_mean = max(y_max_g - y_min_g, 1), (y_max_g + y_min_g) / 2.0
             def parabola(y_n, a, b, c): return a*(y_n**2) + b*y_n + c
             try:
                 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]))
@@ -229,9 +235,12 @@ class WatermelonProcessor:
         except Exception as e:
             print(f"Calibration skipped for {source_name}: {e}")
         
-        # --- 2. YOLO INFERENCE ---
+        # --- 2. YOLO INFERENCE (3 CLASSES FIXED) ---
         results = self.model(image, conf=0.25, retina_masks=True, verbose=False)
-        rind_mask, flesh_contours = np.zeros((h, w), dtype=np.uint8), []
+        rind_mask = np.zeros((h, w), dtype=np.uint8)
+        
+        flesh_l_contours =[]
+        flesh_r_contours = []
 
         if results[0].masks is None:
             return ProcessResult(success=False, message="No masks detected.")
@@ -239,19 +248,32 @@ class WatermelonProcessor:
         for mask_data, cls in zip(results[0].masks.xy, results[0].boxes.cls):
             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: flesh_contours.append(contour)
-
-        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 c_id == 0: 
+                cv2.drawContours(rind_mask, [contour], -1, 255, -1)
+            elif c_id == 1: 
+                flesh_l_contours.append(contour)
+            elif c_id == 2: 
+                flesh_r_contours.append(contour)
+
+        # Failsafe: If YOLO predicted multiple class 1s and 0 class 2s (or vice versa), split them up by X-coordinate
+        if len(flesh_l_contours) >= 2 and len(flesh_r_contours) == 0:
+            flesh_l_contours.sort(key=lambda cnt: cv2.moments(cnt)['m10'] / (cv2.moments(cnt)['m00'] + 1e-5))
+            flesh_r_contours.append(flesh_l_contours.pop())
+        elif len(flesh_r_contours) >= 2 and len(flesh_l_contours) == 0:
+            flesh_r_contours.sort(key=lambda cnt: cv2.moments(cnt)['m10'] / (cv2.moments(cnt)['m00'] + 1e-5))
+            flesh_l_contours.append(flesh_r_contours.pop(0))
+
+        flesh_l = np.zeros((h, w), dtype=np.uint8)
+        flesh_r = 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)
+        for cnt in flesh_l_contours:
+            cv2.drawContours(flesh_l, [cnt], -1, 255, -1)
+        for cnt in flesh_r_contours:
+            cv2.drawContours(flesh_r, [cnt], -1, 255, -1)
 
-        flesh_combined = cv2.bitwise_or(flesh_l_m, flesh_r_m)
+        flesh_combined = cv2.bitwise_or(flesh_l, flesh_r)
+
+        # --- 3. FIT & EXTRACTION ---
         perimeter_data = self.get_stable_perimeter_data(rind_mask, flesh_combined)
         if perimeter_data is None: return ProcessResult(success=False, message="No stable perimeter.")
 
@@ -272,7 +294,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) ---
+        # Re-scale back to original size for true measurements
         orig_scale = 1.0 / scale_ratio
         if cm_per_px is None: cm_per_px = 1.0 
         
@@ -280,20 +302,35 @@ 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]))
 
-        # 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)
 
         # --- 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)
+        midline = self.get_dual_mask_midline(flesh_l, flesh_r, rind_cnt, fit_pts, cx, cy)
+        
+        # Color coding: Green=Rind, Blue=Left Flesh, Red=Right Flesh
+        output = image.copy().astype(np.float32)
+        alpha = 0.42
+        output[..., 0] = np.where(rind_mask > 0, output[..., 0] * (1 - alpha) + 0.0 * alpha, output[..., 0])
+        output[..., 1] = np.where(rind_mask > 0, output[..., 1] * (1 - alpha) + 170.0 * alpha, output[..., 1])
+        output[..., 2] = np.where(rind_mask > 0, output[..., 2] * (1 - alpha) + 0.0 * alpha, output[..., 2])
+        
+        output[..., 0] = np.where(flesh_l > 0, output[..., 0] * (1 - alpha) + 255.0 * alpha, output[..., 0])
+        output[..., 1] = np.where(flesh_l > 0, output[..., 1] * (1 - alpha) + 0.0 * alpha, output[..., 1])
+        output[..., 2] = np.where(flesh_l > 0, output[..., 2] * (1 - alpha) + 0.0 * alpha, output[..., 2])
+
+        output[..., 0] = np.where(flesh_r > 0, output[..., 0] * (1 - alpha) + 0.0 * alpha, output[..., 0])
+        output[..., 1] = np.where(flesh_r > 0, output[..., 1] * (1 - alpha) + 0.0 * alpha, output[..., 1])
+        output[..., 2] = np.where(flesh_r > 0, output[..., 2] * (1 - alpha) + 255.0 * alpha, output[..., 2])
+        
+        output = np.clip(output, 0, 255).astype(np.uint8)
         
         if checker_corners is not None:
             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:
@@ -315,6 +352,7 @@ class WatermelonProcessor:
             image_base64=img_base64, filename=source_name
         )
 
+
 app = FastAPI()
 
 app.add_middleware(
@@ -324,7 +362,7 @@ app.add_middleware(
 processor = WatermelonProcessor(MODEL_PATH)
 
 @app.get("/")
-def read_root(): return {"status": "Watermelon API is awake!"}
+def read_root(): return {"status": "Watermelon API is awake and running!"}
 
 @app.post("/process_single")
 async def process_single(file: UploadFile = File(...)):