Update main.py

main.py

@@ -13,18 +13,17 @@ from fastapi import FastAPI, UploadFile, File
 from fastapi.middleware.cors import CORSMiddleware
 import uvicorn
 
-# OOM PREVENTION
 torch.set_num_threads(1)
 
-# Import your helpers
 from cv_helpers import blend_mask_overlays, stem_tip_tangent_deg
 import color_calibration as calib
 import test_color_eval as eval
 
 # --- CONFIGURATION ---
 MODEL_PATH = "best.pt"
-MAX_IMAGE_SIZE = 1024  
-CHECKER_WIDTH_MM = 63.0  # Physical width of the ColorChecker
+# HF has 16GB RAM, so we can use a higher resolution to ensure the checkerboard is detected!
+MAX_IMAGE_SIZE = 2048  
+CHECKER_WIDTH_CM = 6.3  # 63 mm = 6.3 cm
 
 @dataclass
 class ProcessResult:
@@ -43,7 +42,6 @@ class WatermelonProcessor:
     def __init__(self, model_path: str):
         self.model = YOLO(model_path)
         
-        # Load the Golden Reference once on startup
         self.ref24 = None
         if os.path.exists("reference.png"):
             try:
@@ -55,7 +53,7 @@ class WatermelonProcessor:
             except Exception as e:
                 print(f"Failed to extract reference patches: {e}")
         else:
-            print("WARNING: 'reference.png' not found. Calibration will be skipped.")
+            print("WARNING: 'reference.png' not found. Color Correction will be skipped.")
 
     @staticmethod
     def watermelon_model(theta, Rx, Ry, c_a, d_top, w_top, d_bot, w_bot, phi, c_skew, c_bend):
@@ -67,23 +65,26 @@ class WatermelonProcessor:
         return (ellipse * asymmetry) - divot_top - divot_bot + c_skew * np.sin(t) + c_bend * np.cos(t) * (np.sin(t) ** 2)
 
     @staticmethod
-    def get_stable_perimeter_data(rind_mask, flesh_mask):
+    def get_stable_perimeter_data(rind_mask, flesh_combined):
         cnts, _ = cv2.findContours(rind_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
         if not cnts: return None
+        
+        # Pick the rind that specifically surrounds the flesh
         best_cnt = None
         max_overlap = -1
         for cnt in cnts:
             temp_mask = np.zeros_like(rind_mask)
             cv2.drawContours(temp_mask, [cnt], -1, 255, -1)
-            overlap_area = cv2.countNonZero(cv2.bitwise_and(temp_mask, flesh_mask))
+            overlap_area = cv2.countNonZero(cv2.bitwise_and(temp_mask, flesh_combined))
             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
+        M = cv2.moments(best_cnt)
+        if M["m00"] == 0: return None
 
-        cx, cy = moments["m10"] / moments["m00"], moments["m01"] / moments["m00"]
+        cx, cy = M["m10"] / M["m00"], M["m01"] / M["m00"]
         pts = best_cnt.reshape(-1, 2)
         dx, dy = pts[:, 0] - cx, cy - pts[:, 1]
         r_vals, t_vals = np.sqrt(dx**2 + dy**2), np.arctan2(dy, dx)
@@ -102,8 +103,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
@@ -111,42 +112,40 @@ 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))
+
+        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:
-            y_min, y_max = np.min(y_indices), np.max(y_indices)
-            for y in range(y_min, y_max):
-                row = f_rot[y, :]
-                white_px = np.where(row > 0)[0]
-                if len(white_px) >= 2:
-                    first, last = white_px[0], white_px[-1]
-                    blanks_in_between = np.where(row[first:last] == 0)[0] + first
-                    if len(blanks_in_between) > 0:
-                        gap_points.append([y, np.median(blanks_in_between)])
+        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:
+            for y in range(max(np.min(y_l), np.min(y_r)), min(np.max(y_l), np.max(y_r))):
+                row_l, row_r = np.where(l_rot[y, :] > 0)[0], np.where(r_rot[y, :] > 0)[0]
+                if len(row_l) > 0 and len(row_r) > 0:
+                    gap_points.append([y, (row_l[-1] + row_r[0]) / 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_norm = (gap_points[:, 0] - y_mean) / y_span
-            x_data = gap_points[:, 1]
+            y_norm, x_data = (gap_points[:, 0] - y_mean) / y_span, 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=([-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)
-            ys_extrap_norm = (ys_extrap - y_mean) / y_span
-            xs_extrap = parabola(ys_extrap_norm, *popt_mid)
+            xs_extrap = parabola((ys_extrap - y_mean) / y_span, *popt_mid)
             pts_rot = np.vstack([xs_extrap, ys_extrap, np.ones_like(xs_extrap)])
         else:
             ys_extrap = np.linspace(0, h, 500)
-            xs_extrap = np.full_like(ys_extrap, cx)
-            pts_rot = np.vstack([xs_extrap, ys_extrap, np.ones_like(xs_extrap)])
+            pts_rot = np.vstack([np.full_like(ys_extrap, cx), ys_extrap, np.ones_like(ys_extrap)])
 
         pts_orig = (m_inv @ pts_rot).T
         pred_cnt_cv = predicted_cnt.reshape(-1, 1, 2).astype(np.int32)
@@ -156,50 +155,58 @@ class WatermelonProcessor:
         if image is None: return ProcessResult(success=False, message="Could not decode image.")
         h, w = image.shape[:2]
 
-        # --- CALIBRATION & SCALING ---
-        dE_initial, dE_final, mm_per_px, checker_corners = None, None, None, None
+        # --- 1. CALIBRATION & SCALING ---
+        dE_initial, dE_final, cm_per_px, checker_corners = None, None, None, None
         
-        if self.ref24 is not None:
-            try:
-                # 1. Detect Checker
-                checker_corners = calib.detect_checker_corners(image)
-                
-                # 2. Calculate mm/px Scale
-                # Corners:[Bottom-Left, Top-Left, Top-Right, Bottom-Right]
-                top_width = np.linalg.norm(checker_corners[1] - checker_corners[2])
-                bot_width = np.linalg.norm(checker_corners[0] - checker_corners[3])
-                px_width = (top_width + bot_width) / 2.0
-                mm_per_px = CHECKER_WIDTH_MM / px_width
-                
-                # 3. Apply Calibration Pipeline
+        try:
+            checker_corners = calib.detect_checker_corners(image)
+            top_width = np.linalg.norm(checker_corners[1] - checker_corners[2])
+            bot_width = np.linalg.norm(checker_corners[0] - checker_corners[3])
+            px_width = (top_width + bot_width) / 2.0
+            cm_per_px = CHECKER_WIDTH_CM / px_width
+
+            if self.ref24 is not None:
                 tgt_warped = calib.warp_checker(image, checker_corners)
                 tgt24 = calib.sample_24_patches(tgt_warped)
                 
                 dE_initial = float(np.mean(eval.compute_deltaE_00(tgt24, self.ref24)))
-                tgt24_corr = eval.apply_pipeline_to_patches(tgt24, self.ref24)
-                dE_final = float(np.mean(eval.compute_deltaE_00(tgt24_corr, self.ref24)))
                 
+                # Apply Color Correction Pipeline
                 image = calib.apply_pipeline(image, self.ref24, tgt24)
-            except Exception as e:
-                print(f"Calibration skipped for {source_name}: {e}")
+                
+                # Re-check the corrected patches for Final dE
+                tgt_warped_corr = calib.warp_checker(image, checker_corners)
+                tgt24_corr = calib.sample_24_patches(tgt_warped_corr)
+                dE_final = float(np.mean(eval.compute_deltaE_00(tgt24_corr, self.ref24)))
+        except Exception as e:
+            print(f"Calibration skipped for {source_name}: {e}")
         
-        # Fallback to pure pixels if checker not found
-        if mm_per_px is None:
-            mm_per_px = 1.0 / scale_ratio 
-
-        # --- YOLO INFERENCE ---
-        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)
-
-        if results[0].masks is None:
-            return ProcessResult(success=False, message="No masks detected.")
-
-        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)
-
-        perimeter_data = self.get_stable_perimeter_data(rind_mask, flesh_mask)
+        # --- 2. YOLO INFERENCE ---
+        results = self.model(image, conf=0.25, retina_masks=True, verbose=False)
+        rind_mask = np.zeros((h, w), dtype=np.uint8)
+        flesh_contours = []
+
+        if results[0].masks is not None:
+            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: 
+                    flesh_contours.append(contour)
+
+        # Split Class 1 into Left and Right physically!
+        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)
+
+        # --- 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.")
 
         t_data, r_raw, (cx, cy), rind_cnt = perimeter_data
@@ -218,37 +225,39 @@ 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)])
         
-        # --- FEATURE EXTRACTION (IN MM) ---
+        # 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 # Fallback to pixel units if checker failed
+        
         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]))
-
-        width_val = width_px * mm_per_px
-        height_val = height_px * mm_per_px
-        perimeter_val = perimeter_px * mm_per_px
+        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]))
 
-        # --- 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)
+        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
 
-        # Draw Color Checker Box
+        # --- 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 Checker Box
         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)
 
         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)
             p1 = (int(round(tx)), int(round(ty)))
-            p2 = (int(round(tx + L * np.cos(rad))), int(round(ty + L * np.sin(rad))))
+            p2 = (int(round(tx + L * np.cos(np.deg2rad(tdeg)))), int(round(ty + L * np.sin(np.deg2rad(tdeg)))))
             cv2.circle(output, p1, 6, (255, 0, 255), -1)
             cv2.line(output, p1, p2, (255, 0, 255), 2)
 
-        _, buffer = cv2.imencode('.jpg', output,[cv2.IMWRITE_JPEG_QUALITY, 85])
+        _, buffer = cv2.imencode('.jpg', output, [cv2.IMWRITE_JPEG_QUALITY, 85])
         img_base64 = base64.b64encode(buffer).decode('utf-8')
 
         return ProcessResult(
@@ -259,22 +268,16 @@ class WatermelonProcessor:
         )
 
 
-# --- FASTAPI APP ---
 app = FastAPI()
 
 app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"], 
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
+    CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"],
 )
 
 processor = WatermelonProcessor(MODEL_PATH)
 
 @app.get("/")
-def read_root():
-    return {"status": "Watermelon API is awake and running!"}
+def read_root(): return {"status": "Watermelon API is awake!"}
 
 @app.post("/process_single")
 async def process_single(file: UploadFile = File(...)):