Update main.py

main.py

@@ -17,7 +17,6 @@ from skimage import color
 # OOM PREVENTION
 torch.set_num_threads(1)
 
-# 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 +25,7 @@ MAX_IMAGE_SIZE = 2048
 CHECKER_WIDTH_CM = 6.3 
 
 # ==============================================================================
-# --- COLOR CALIBRATION LOGIC (Embedded) ---
+# --- COLOR CALIBRATION LOGIC (Polynomial Color Correction) ---
 # ==============================================================================
 def to_linear_srgb(u8_bgr):
     rgb = cv2.cvtColor(u8_bgr, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
@@ -68,31 +67,29 @@ 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):
-    gains = np.median(ref24[18:24], axis=0) / np.maximum(np.median(tgt24[18:24], axis=0), 1e-6)
-    lin = to_linear_srgb(target_bgr)
-    lin_wb = lin * gains.reshape(1, 1, 3)
-    tgt24_wb = tgt24 * gains
-
-    M = np.linalg.lstsq(tgt24_wb[:18], ref24[:18], rcond=None)[0].T
-    corrected = (lin_wb.reshape(-1, 3) @ M.T).reshape(lin_wb.shape)
-    tgt24_ccm = tgt24_wb @ M.T
-
-    w = np.array([0.2126, 0.7152, 0.0722], np.float32)
-    Ls, Lt = (tgt24_ccm[19:23] @ w), (ref24[19:23] @ w)
-    sort_idx = np.argsort(Ls)
-    Ls, Lt = np.concatenate([[0.01], Ls[sort_idx], [0.98]]), np.concatenate([[0.01], Lt[sort_idx],[0.98]])
-    
-    L = np.clip(np.tensordot(corrected, w, axes=([2],[0])), 0, 1)
-    Lt_mapped = np.interp(L, Ls, Lt)
-    
-    knee, strength = 0.90, 0.6
-    below = Lt_mapped < knee
-    Lt_final = np.empty_like(Lt_mapped)
-    Lt_final[below] = Lt_mapped[below]
-    Lt_final[~below] = knee + (1.0 - knee) * (1.0 - np.exp(-strength * ((Lt_mapped[~below] - knee) / (1.0 - knee))))
-    
-    scale = (Lt_final + 1e-6) / (L + 1e-6)
-    return to_srgb_u8(np.clip(corrected * scale[..., None], 0, 1))
+    """
+    Polynomial Color Correction Matrix (PCCM).
+    Fits a 10-term polynomial to robustly map Target colors -> Reference colors.
+    """
+    def extract_features(rgb_array):
+        R, G, B = rgb_array[..., 0], rgb_array[..., 1], rgb_array[..., 2]
+        return np.stack([
+            R, G, B, 
+            R*G, R*B, G*B, 
+            R**2, G**2, B**2, 
+            np.ones_like(R)
+        ], axis=-1)
+
+    # 1. Fit the polynomial weights to the 24 patches
+    X_tgt = extract_features(tgt24)  # Shape: (24, 10)
+    W, _, _, _ = np.linalg.lstsq(X_tgt, ref24, rcond=None)  # Shape: (10, 3)
+
+    # 2. Apply the weights to the entire target image
+    lin_img = to_linear_srgb(target_bgr)
+    img_features = extract_features(lin_img) # Shape: (H, W, 10)
+    corrected_lin = img_features @ W         # Shape: (H, W, 3)
+
+    return to_srgb_u8(np.clip(corrected_lin, 0, 1))
 
 # ==============================================================================
 # --- CORE API & PROCESSOR ---
@@ -215,7 +212,6 @@ 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:
@@ -235,12 +231,9 @@ class WatermelonProcessor:
         except Exception as e:
             print(f"Calibration skipped for {source_name}: {e}")
         
-        # --- 2. YOLO INFERENCE (3 CLASSES FIXED) ---
         results = self.model(image, conf=0.25, retina_masks=True, verbose=False)
         rind_mask = np.zeros((h, w), dtype=np.uint8)
-        
-        flesh_l_contours =[]
-        flesh_r_contours = []
+        flesh_contours = []
 
         if results[0].masks is None:
             return ProcessResult(success=False, message="No masks detected.")
@@ -251,29 +244,20 @@ class WatermelonProcessor:
             if c_id == 0: 
                 cv2.drawContours(rind_mask, [contour], -1, 255, -1)
             elif c_id == 1: 
-                flesh_l_contours.append(contour)
+                flesh_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)
-        
-        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_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)
 
-        # --- 3. FIT & EXTRACTION ---
+        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.")
 
@@ -294,7 +278,6 @@ 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)])
         
-        # 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 
         
@@ -306,28 +289,26 @@ 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, flesh_r, rind_cnt, fit_pts, cx, cy)
+        midline = self.get_dual_mask_midline(flesh_l_m, flesh_r_m, 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_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 > 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[..., 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)
         
         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)
@@ -343,13 +324,12 @@ class WatermelonProcessor:
             cv2.line(output, p1, p2, (255, 0, 255), 2)
 
         _, buffer = cv2.imencode('.jpg', output, [cv2.IMWRITE_JPEG_QUALITY, 85])
-        img_base64 = base64.b64encode(buffer).decode('utf-8')
 
         return ProcessResult(
             success=True, message="Success", r2_score=float(r2),
             width_val=width_val, height_val=height_val, perimeter_val=perimeter_val,
             delta_e_initial=dE_initial, delta_e_final=dE_final,
-            image_base64=img_base64, filename=source_name
+            image_base64=base64.b64encode(buffer).decode('utf-8'), filename=source_name
         )