Update main.py

main.py

@@ -7,7 +7,7 @@ import torch
 from dataclasses import dataclass
 from typing import Optional
 from scipy.ndimage import median_filter
-from scipy.optimize import curve_fit
+from scipy.optimize import curve_fit, minimize
 from ultralytics import YOLO
 from fastapi import FastAPI, UploadFile, File
 from fastapi.middleware.cors import CORSMiddleware
@@ -25,7 +25,7 @@ MAX_IMAGE_SIZE = 2048
 CHECKER_WIDTH_CM = 6.3 
 
 # ==============================================================================
-# --- COLOR CALIBRATION LOGIC (Polynomial Color Correction) ---
+# --- DIRECT DELTA-E COLOR CALIBRATION ---
 # ==============================================================================
 def to_linear_srgb(u8_bgr):
     rgb = cv2.cvtColor(u8_bgr, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
@@ -39,10 +39,18 @@ def to_srgb_u8(lin_rgb):
 
 def detect_checker_corners(img_bgr):
     det = cv2.mcc.CCheckerDetector_create()
-    if not det.process(img_bgr, cv2.mcc.MCC24):
-        raise RuntimeError("ColorChecker not found")
-    cc = det.getListColorChecker()[0]
-    return np.array(cc.getBox() if hasattr(cc, "getBox") else cc.getCorners(), dtype=np.float32)
+    if det.process(img_bgr, cv2.mcc.MCC24):
+        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
+        
+    raise RuntimeError("ColorChecker not found")
 
 def warp_checker(img, corners, out_w=600, out_h=400):
     dst = np.float32([[0, out_h-1],[0, 0],[out_w-1, 0],[out_w-1, out_h-1]])
@@ -67,28 +75,31 @@ 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):
-    """
-    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)
+    """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))
 
 # ==============================================================================
@@ -199,7 +210,7 @@ class WatermelonProcessor:
             
             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)])
+            pts_rot = np.vstack([xs_extrap, ys_extrap, np.ones_like(xs_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)])
@@ -232,8 +243,7 @@ class WatermelonProcessor:
             print(f"Calibration skipped for {source_name}: {e}")
         
         results = self.model(image, conf=0.25, retina_masks=True, verbose=False)
-        rind_mask = np.zeros((h, w), dtype=np.uint8)
-        flesh_contours = []
+        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.")
@@ -241,18 +251,14 @@ 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)
-            elif c_id == 2: 
-                flesh_contours.append(contour)
+            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 len(flesh_contours) >= 2:
-            cv2.drawContours(flesh_l_m, [flesh_contours[0]], -1, 255, -1)
+            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)
@@ -290,22 +296,7 @@ class WatermelonProcessor:
         perimeter_val = float(perimeter_px * cm_per_px * orig_scale)
 
         midline = self.get_dual_mask_midline(flesh_l_m, flesh_r_m, rind_cnt, fit_pts, cx, cy)
-        
-        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)
+        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)
@@ -323,13 +314,14 @@ class WatermelonProcessor:
             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(
             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=base64.b64encode(buffer).decode('utf-8'), filename=source_name
+            image_base64=img_base64, filename=source_name
         )