Update main.py

main.py

@@ -17,15 +17,13 @@ from skimage import color
 # OOM PREVENTION
 torch.set_num_threads(1)
 
-from cv_helpers import blend_mask_overlays, stem_tip_tangent_deg
-
 # --- CONFIGURATION ---
 MODEL_PATH = "best.pt"
 MAX_IMAGE_SIZE = 2048  
 CHECKER_WIDTH_CM = 6.3 
 
 # ==============================================================================
-# --- DIRECT DELTA-E COLOR CALIBRATION ---
+# --- COLOR CALIBRATION LOGIC ---
 # ==============================================================================
 def to_linear_srgb(u8_bgr):
     rgb = cv2.cvtColor(u8_bgr, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
@@ -43,12 +41,11 @@ def detect_checker_corners(img_bgr):
         cc = det.getListColorChecker()[0]
         return np.array(cc.getBox() if hasattr(cc, "getBox") else cc.getCorners(), dtype=np.float32)
         
-    # Fallback: Sometimes OpenCV fails on high-res noise. Scale by 50% and try again.
     img_small = cv2.resize(img_bgr, (0,0), fx=0.5, fy=0.5)
     if det.process(img_small, cv2.mcc.MCC24):
         cc = det.getListColorChecker()[0]
         pts = np.array(cc.getBox() if hasattr(cc, "getBox") else cc.getCorners(), dtype=np.float32)
-        return pts * 2.0 # Scale corners back up
+        return pts * 2.0 
         
     raise RuntimeError("ColorChecker not found")
 
@@ -75,31 +72,22 @@ def compute_deltaE_00(lin_src, lin_ref):
     return color.deltaE_ciede2000(color.rgb2lab(srgb_src.reshape(1, -1, 3)), color.rgb2lab(srgb_ref.reshape(1, -1, 3))).flatten()
 
 def apply_color_pipeline(target_bgr, ref24, tgt24):
-    """Safely scales White Balance, then uses Powell's method to minimize Delta E."""
     tgt_lin = to_linear_srgb(target_bgr)
-    
-    # 1. White Balance (Von Kries scaling using 6 neutrals)
     gains = np.median(ref24[18:24], axis=0) / np.maximum(np.median(tgt24[18:24], axis=0), 1e-6)
     tgt_lin_wb = tgt_lin * gains.reshape(1, 1, 3)
     tgt24_wb = tgt24 * gains
 
-    # 2. Objective Function: Explicitly optimize the matrix for the lowest Delta E score
     def objective(W_flat):
         W = W_flat.reshape(3, 3)
         pred_lin = np.clip(tgt24_wb @ W, 0, 1)
         return np.mean(compute_deltaE_00(pred_lin, ref24))
         
-    # Get a fast starting point using standard Ridge Regression
     X, Y = tgt24_wb, ref24
     W_init = np.linalg.inv(X.T @ X + 0.05 * np.eye(3)) @ X.T @ Y
-    
-    # Optimize matrix to minimize CIEDE2000 natively
     res = minimize(objective, W_init.flatten(), method='Powell')
     W_opt = res.x.reshape(3, 3)
 
-    # 3. Apply to full image securely
     corrected_lin = (tgt_lin_wb.reshape(-1, 3) @ W_opt).reshape(tgt_lin_wb.shape)
-    
     return to_srgb_u8(np.clip(corrected_lin, 0, 1))
 
 # ==============================================================================
@@ -122,7 +110,6 @@ class ProcessResult:
 class WatermelonProcessor:
     def __init__(self, model_path: str):
         self.model = YOLO(model_path)
-        
         self.ref24 = None
         if os.path.exists("reference.png"):
             try:
@@ -185,8 +172,7 @@ 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]
+        l_idx, r_idx = np.where(l_rot > 0)[1], 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 
@@ -204,13 +190,14 @@ class WatermelonProcessor:
             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
+            max_bend = w * 0.08
             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]))
+                popt_mid, _ = curve_fit(parabola, (gap_points[:,0]-y_mean)/y_span, gap_points[:,1], 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)
-            pts_rot = np.vstack([xs_extrap, ys_extrap, np.ones_like(xs_extrap)])
+            pts_rot = np.vstack([xs_extrap, ys_extrap, np.ones_like(ys_extrap)])
         else:
             ys_extrap = np.linspace(0, h, 500)
             pts_rot = np.vstack([np.full_like(ys_extrap, cx), ys_extrap, np.ones_like(ys_extrap)])
@@ -223,8 +210,8 @@ 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:
             checker_corners = detect_checker_corners(image)
             top_width = np.linalg.norm(checker_corners[1] - checker_corners[2])
@@ -232,18 +219,23 @@ class WatermelonProcessor:
             cm_per_px = CHECKER_WIDTH_CM / ((top_width + bot_width) / 2.0)
 
             if self.ref24 is not None:
-                tgt24 = sample_24_patches(warp_checker(image, checker_corners))
+                tgt_warped = warp_checker(image, checker_corners)
+                tgt24 = sample_24_patches(tgt_warped)
                 dE_initial = float(np.mean(compute_deltaE_00(tgt24, self.ref24)))
                 
                 image = apply_color_pipeline(image, self.ref24, tgt24)
                 
-                tgt24_corr = sample_24_patches(warp_checker(image, checker_corners))
+                tgt_warped_corr = warp_checker(image, checker_corners)
+                tgt24_corr = sample_24_patches(tgt_warped_corr)
                 dE_final = float(np.mean(compute_deltaE_00(tgt24_corr, self.ref24)))
         except Exception as e:
             print(f"Calibration skipped for {source_name}: {e}")
         
+        # --- 2. YOLO INFERENCE (STRICTLY PARSING ALL 3 CLASSES) ---
         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.")
@@ -251,19 +243,28 @@ 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)
+            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 missed one side but predicted multiple of the other
+        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_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)
+        for cnt in flesh_l_contours: cv2.drawContours(flesh_l_m, [cnt], -1, 255, -1)
+        for cnt in flesh_r_contours: cv2.drawContours(flesh_r_m, [cnt], -1, 255, -1)
 
         flesh_combined = cv2.bitwise_or(flesh_l_m, flesh_r_m)
+
+        # --- 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.")
 
@@ -295,24 +296,43 @@ class WatermelonProcessor:
         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)
         
+        # 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_m > 0, output[..., 0] * (1 - alpha) + 255.0 * alpha, output[..., 0])
+        output[..., 1] = np.where(flesh_l_m > 0, output[..., 1] * (1 - alpha) + 0.0 * alpha, output[..., 1])
+        output[..., 2] = np.where(flesh_l_m > 0, output[..., 2] * (1 - alpha) + 0.0 * alpha, output[..., 2])
+
+        output[..., 0] = np.where(flesh_r_m > 0, output[..., 0] * (1 - alpha) + 0.0 * alpha, output[..., 0])
+        output[..., 1] = np.where(flesh_r_m > 0, output[..., 1] * (1 - alpha) + 0.0 * alpha, output[..., 1])
+        output[..., 2] = np.where(flesh_r_m > 0, output[..., 2] * (1 - alpha) + 255.0 * alpha, output[..., 2])
+        
+        output = np.clip(output, 0, 255).astype(np.uint8)
+        
+        # Draw Color Checker Box
         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)
-
-        stem = stem_tip_tangent_deg(rind_cnt, (cx, cy))
-        if stem is not None:
-            tx, ty, tdeg = stem
-            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)
-            cv2.line(output, p1, p2, (255, 0, 255), 2)
+        # Draw Midline and Midline Intersections
+        if len(midline) > 1: 
+            cv2.polylines(output, [midline.astype(np.int32)], False, (0, 255, 255), 3)
+            # Add intersection dots (White inner, Black outer)
+            pt_top = (int(midline[0][0]), int(midline[0][1]))
+            pt_bot = (int(midline[-1][0]), int(midline[-1][1]))
+            cv2.circle(output, pt_top, 10, (0, 0, 0), 2)
+            cv2.circle(output, pt_top, 8, (255, 255, 255), -1)
+            cv2.circle(output, pt_bot, 10, (0, 0, 0), 2)
+            cv2.circle(output, pt_bot, 8, (255, 255, 255), -1)
+
+        # Draw Predicted Perimeter Boundary
+        cv2.polylines(output, [fit_pts.astype(np.int32)], True, (0, 255, 0), 3)
 
         _, buffer = cv2.imencode('.jpg', output,[cv2.IMWRITE_JPEG_QUALITY, 85])
         img_base64 = base64.b64encode(buffer).decode('utf-8')