Update main.py

main.py

@@ -12,18 +12,95 @@ from ultralytics import YOLO
 from fastapi import FastAPI, UploadFile, File
 from fastapi.middleware.cors import CORSMiddleware
 import uvicorn
+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
-import color_calibration as calib
-import test_color_eval as eval
 
 # --- CONFIGURATION ---
 MODEL_PATH = "best.pt"
-# 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
+MAX_IMAGE_SIZE = 2048  # Large enough to ensure the color checker is clearly visible
+CHECKER_WIDTH_CM = 6.3 # Physical width of the X-Rite ColorChecker Classic
+
+# ==============================================================================
+# --- COLOR CALIBRATION LOGIC (Embedded) ---
+# ==============================================================================
+def to_linear_srgb(u8_bgr):
+    rgb = cv2.cvtColor(u8_bgr, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
+    a = 0.055
+    return np.where(rgb <= 0.04045, rgb / 12.92, ((rgb + a) / (1 + a)) ** 2.4)
+
+def to_srgb_u8(lin_rgb):
+    a = 0.055
+    srgb = np.where(lin_rgb <= 0.0031308, 12.92 * lin_rgb, (1 + a) * np.power(np.maximum(lin_rgb, 0), 1/2.4) - a)
+    return cv2.cvtColor((np.clip(srgb, 0, 1) * 255.0).astype(np.uint8), cv2.COLOR_RGB2BGR)
+
+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)
+
+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]])
+    H_mat = cv2.getPerspectiveTransform(corners, dst)
+    return cv2.warpPerspective(img, H_mat, (out_w, out_h), flags=cv2.INTER_CUBIC)
+
+def sample_24_patches(warped, margin=12):
+    H, W = warped.shape[:2]
+    cell_w, cell_h = W / 6.0, H / 4.0
+    lin = to_linear_srgb(warped)
+    means =[]
+    for r in range(4):
+        for c in range(6):
+            x0, x1 = int(c * cell_w + margin), int((c+1) * cell_w - margin)
+            y0, y1 = int(r * cell_h + margin), int((r+1) * cell_h - margin)
+            means.append(np.median(lin[y0:y1, x0:x1].reshape(-1, 3), axis=0))
+    return np.stack(means, 0)
+
+def compute_deltaE_00(lin_src, lin_ref):
+    srgb_src = to_srgb_u8(lin_src).astype(np.float32) / 255.0
+    srgb_ref = to_srgb_u8(lin_ref).astype(np.float32) / 255.0
+    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):
+    # White Balance
+    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
+
+    # CCM
+    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
+
+    # Luma Curve
+    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)
+    
+    # Soft Rolloff
+    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))
+
+# ==============================================================================
+# --- CORE API & PROCESSOR ---
+# ==============================================================================
 
 @dataclass
 class ProcessResult:
@@ -42,18 +119,16 @@ class WatermelonProcessor:
     def __init__(self, model_path: str):
         self.model = YOLO(model_path)
         
+        # Load Golden Reference Patches
         self.ref24 = None
         if os.path.exists("reference.png"):
             try:
                 ref_img = cv2.imread("reference.png")
-                ref_corners = calib.detect_checker_corners(ref_img)
-                ref_warped = calib.warp_checker(ref_img, ref_corners)
-                self.ref24 = calib.sample_24_patches(ref_warped)
-                print("Reference ColorChecker loaded successfully.")
+                ref_corners = detect_checker_corners(ref_img)
+                self.ref24 = sample_24_patches(warp_checker(ref_img, ref_corners))
+                print("Reference ColorChecker patches loaded.")
             except Exception as e:
-                print(f"Failed to extract reference patches: {e}")
-        else:
-            print("WARNING: 'reference.png' not found. Color Correction will be skipped.")
+                print(f"Reference extraction failed: {e}")
 
     @staticmethod
     def watermelon_model(theta, Rx, Ry, c_a, d_top, w_top, d_bot, w_bot, phi, c_skew, c_bend):
@@ -68,26 +143,23 @@ class WatermelonProcessor:
     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
+        best_cnt, max_overlap = None, -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_combined))
-            if overlap_area > max_overlap:
-                max_overlap = overlap_area
-                best_cnt = cnt
+            temp = np.zeros_like(rind_mask)
+            cv2.drawContours(temp, [cnt], -1, 255, -1)
+            overlap = cv2.countNonZero(cv2.bitwise_and(temp, flesh_combined))
+            if overlap > max_overlap:
+                max_overlap, best_cnt = overlap, cnt
                 
         if best_cnt is None: best_cnt = max(cnts, key=cv2.contourArea)
         M = cv2.moments(best_cnt)
         if M["m00"] == 0: return None
 
-        cx, cy = M["m10"] / M["m00"], M["m01"] / M["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)
+        
         num_bins = 360
         bins = np.linspace(-np.pi, np.pi, num_bins + 1)
         raw_r = np.full(num_bins, np.nan)
@@ -98,28 +170,17 @@ class WatermelonProcessor:
         valid_idx = np.where(~np.isnan(raw_r))[0]
         if len(valid_idx) == 0: return None
         raw_r[np.isnan(raw_r)] = np.interp(np.where(np.isnan(raw_r))[0], valid_idx, raw_r[valid_idx], period=360)
-        final_r = median_filter(raw_r, size=7, mode="wrap")
-        final_theta = (bins[:-1] + bins[1:]) / 2.0
-        return final_theta, final_r, (cx, cy), best_cnt
+        return (bins[:-1] + bins[1:])/2.0, median_filter(raw_r, size=7, mode="wrap"), (cx, cy), best_cnt
 
     @staticmethod
-    def get_dual_mask_midline(f_left, f_right, rind_cnt, predicted_cnt, cx, cy):
+    def get_dual_mask_midline(f_left, f_right, rind_cnt, pred_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
-        else: rot_angle = 0
+        _, (ma, Ma), angle = cv2.fitEllipse(rind_cnt) if len(rind_cnt) > 5 else (None, (0,0), 0)
+        rot_angle = angle if ma < Ma else angle + 90
 
         m_rot = cv2.getRotationMatrix2D((cx, cy), rot_angle, 1.0)
         m_inv = cv2.getRotationMatrix2D((cx, cy), -rot_angle, 1.0)
-        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 
+        l_rot, r_rot = cv2.warpAffine(f_left, m_rot, (w, h)), cv2.warpAffine(f_right, m_rot, (w, h))
 
         gap_points =[]
         y_l, y_r = np.where(l_rot > 0)[0], np.where(r_rot > 0)[0]
@@ -133,22 +194,18 @@ class WatermelonProcessor:
         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, 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]))
+            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]))
             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)])
+            ys_ex = np.linspace(0, h, 500)
+            pts_rot = np.vstack([parabola((ys_ex - y_mean)/y_span, *popt_mid), ys_extrap, np.ones_like(ys_ex)])
         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)])
+            ys_ex = np.linspace(0, h, 500)
+            pts_rot = np.vstack([np.full_like(ys_ex, cx), ys_ex, np.ones_like(ys_ex)])
 
         pts_orig = (m_inv @ pts_rot).T
-        pred_cnt_cv = predicted_cnt.reshape(-1, 1, 2).astype(np.int32)
+        pred_cnt_cv = pred_cnt.reshape(-1, 1, 2).astype(np.int32)
         return np.array([pt for pt in pts_orig if cv2.pointPolygonTest(pred_cnt_cv, (float(pt[0]), float(pt[1])), False) >= 0])
 
     def process_image(self, image: np.ndarray, source_name: str, scale_ratio: float) -> ProcessResult:
@@ -159,53 +216,40 @@ class WatermelonProcessor:
         dE_initial, dE_final, cm_per_px, checker_corners = None, None, None, None
         
         try:
-            checker_corners = calib.detect_checker_corners(image)
+            checker_corners = detect_checker_corners(image)
+            # Physical width logic
             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
+            cm_per_px = CHECKER_WIDTH_CM / ((top_width + bot_width) / 2.0)
 
             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 = sample_24_patches(warp_checker(image, checker_corners))
+                dE_initial = float(np.mean(compute_deltaE_00(tgt24, self.ref24)))
                 
-                # Apply Color Correction Pipeline
-                image = calib.apply_pipeline(image, self.ref24, tgt24)
+                # Apply pipeline to full image BEFORE YOLO inference
+                image = apply_color_pipeline(image, self.ref24, tgt24)
                 
-                # 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)))
+                # Re-check the fully corrected image
+                tgt24_corr = sample_24_patches(warp_checker(image, checker_corners))
+                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 ---
+        # --- 2. YOLO INFERENCE (3 CLASSES) ---
         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 ---
+        rind_mask, flesh_l, flesh_r = np.zeros((h, w), dtype=np.uint8), 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)
+            c_id = int(cls)
+            if c_id == 0: cv2.drawContours(rind_mask, [contour], -1, 255, -1)
+            elif c_id == 1: cv2.drawContours(flesh_l, [contour], -1, 255, -1)
+            elif c_id == 2: cv2.drawContours(flesh_r, [contour], -1, 255, -1)
+
+        flesh_combined = cv2.bitwise_or(flesh_l, flesh_r)
         perimeter_data = self.get_stable_perimeter_data(rind_mask, flesh_combined)
         if perimeter_data is None: return ProcessResult(success=False, message="No stable perimeter.")
 
@@ -213,11 +257,9 @@ class WatermelonProcessor:
         scale = np.mean(r_raw)
         
         try:
-            popt, _ = curve_fit(
-                self.watermelon_model, t_data, r_raw / scale,
+            popt, _ = curve_fit(self.watermelon_model, t_data, r_raw / scale,
                 p0=[1.0, 1.1, 0.0, 0.05, 3.0, 0.05, 3.0, 0.0, 0.0, 0.0],
-                bounds=([0.5, 0.5, -0.4, 0.0, 0.1, 0.0, 0.1, -1.5, -0.2, -0.2],[2.0, 2.0, 0.4, 0.5, 50.0, 0.5, 50.0, 1.5, 0.2, 0.2]),
-            )
+                bounds=([0.5, 0.5, -0.4, 0.0, 0.1, 0.0, 0.1, -1.5, -0.2, -0.2],[2.0, 2.0, 0.4, 0.5, 50.0, 0.5, 50.0, 1.5, 0.2, 0.2]))
         except Exception as exc: return ProcessResult(success=False, message=f"Fit failed: {exc}")
 
         r2 = 1 - (np.sum((r_raw / scale - self.watermelon_model(t_data, *popt)) ** 2) / np.sum((r_raw / scale - 1) ** 2))
@@ -225,23 +267,18 @@ 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
+        # --- 3. DIMENSIONAL MATH (IN CM) ---
         orig_scale = 1.0 / scale_ratio
-        if cm_per_px is None: cm_per_px = 1.0 # Fallback to pixel units if checker failed
+        if cm_per_px is None: cm_per_px = 1.0
         
-        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]))
-        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]))
-
-        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
+        width_val = float(np.max(fit_pts[:, 0]) - np.min(fit_pts[:, 0])) * cm_per_px * orig_scale
+        height_val = float(np.max(fit_pts[:, 1]) - np.min(fit_pts[:, 1])) * cm_per_px * orig_scale
+        perimeter_val = float(np.sum(np.linalg.norm(np.diff(fit_pts, axis=0), axis=1)) + np.linalg.norm(fit_pts[-1] - fit_pts[0])) * 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)
+        midline = self.get_dual_mask_midline(flesh_l, flesh_r, 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)
 
@@ -251,20 +288,19 @@ class WatermelonProcessor:
         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(np.deg2rad(tdeg)))), int(round(ty + L * np.sin(np.deg2rad(tdeg)))))
+            p2 = (int(round(tx + min(w, h)*0.08 * np.cos(rad))), int(round(ty + min(w, h)*0.08 * np.sin(rad))))
             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])
-        img_base64 = base64.b64encode(buffer).decode('utf-8')
+        _, buffer = cv2.imencode('.jpg', output,[cv2.IMWRITE_JPEG_QUALITY, 85])
 
         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
         )
 
 
@@ -282,18 +318,17 @@ def read_root(): return {"status": "Watermelon API is awake!"}
 @app.post("/process_single")
 async def process_single(file: UploadFile = File(...)):
     contents = await file.read()
-    nparr = np.frombuffer(contents, np.uint8)
-    img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
+    img = cv2.imdecode(np.frombuffer(contents, np.uint8), cv2.IMREAD_COLOR)
     
-    h, w = img.shape[:2]
     scale_ratio = 1.0
+    h, w = img.shape[:2]
     if max(h, w) > MAX_IMAGE_SIZE:
         scale_ratio = MAX_IMAGE_SIZE / float(max(h, w))
         img = cv2.resize(img, (int(w * scale_ratio), int(h * scale_ratio)), interpolation=cv2.INTER_AREA)
 
     res = processor.process_image(img, file.filename, scale_ratio)
     
-    del img, nparr, contents
+    del img, contents
     gc.collect()
 
     return res.__dict__