Update main.py

main.py

@@ -22,11 +22,11 @@ from cv_helpers import blend_mask_overlays, stem_tip_tangent_deg
 
 # --- CONFIGURATION ---
 MODEL_PATH = "best.pt"
-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
+MAX_IMAGE_SIZE = 2048  
+CHECKER_WIDTH_CM = 6.3 
 
 # ==============================================================================
-# --- COLOR CALIBRATION LOGIC (Embedded) ---
+# --- COLOR CALIBRATION LOGIC ---
 # ==============================================================================
 def to_linear_srgb(u8_bgr):
     rgb = cv2.cvtColor(u8_bgr, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
@@ -46,7 +46,7 @@ def detect_checker_corners(img_bgr):
     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]])
+    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)
 
@@ -68,27 +68,23 @@ 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):
-    # 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)
+    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)
@@ -119,7 +115,6 @@ 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:
@@ -196,13 +191,16 @@ class WatermelonProcessor:
             y_span, y_mean = max(y_max - y_min, 1), (y_max + y_min) / 2.0
             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]))
+                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_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)])
+            
+            # FIXED TYPO HERE (ys_extrap)
+            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)])
         else:
-            ys_ex = np.linspace(0, h, 500)
-            pts_rot = np.vstack([np.full_like(ys_ex, cx), ys_ex, np.ones_like(ys_ex)])
+            ys_extrap = np.linspace(0, h, 500)
+            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 = pred_cnt.reshape(-1, 1, 2).astype(np.int32)
@@ -212,12 +210,10 @@ 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)
-            # 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])
             cm_per_px = CHECKER_WIDTH_CM / ((top_width + bot_width) / 2.0)
@@ -226,16 +222,13 @@ class WatermelonProcessor:
                 tgt24 = sample_24_patches(warp_checker(image, checker_corners))
                 dE_initial = float(np.mean(compute_deltaE_00(tgt24, self.ref24)))
                 
-                # Apply pipeline to full image BEFORE YOLO inference
                 image = apply_color_pipeline(image, self.ref24, tgt24)
                 
-                # 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 (3 CLASSES) ---
         results = self.model(image, conf=0.25, retina_masks=True, verbose=False)
         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)
 
@@ -255,11 +248,14 @@ class WatermelonProcessor:
 
         t_data, r_raw, (cx, cy), rind_cnt = perimeter_data
         scale = np.mean(r_raw)
+        if scale <= 0: return ProcessResult(success=False, message="Invalid perimeter scale.")
         
         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))
@@ -267,15 +263,17 @@ 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)])
         
-        # --- 3. DIMENSIONAL MATH (IN CM) ---
         orig_scale = 1.0 / scale_ratio
-        if cm_per_px is None: cm_per_px = 1.0
+        if cm_per_px is None: cm_per_px = 1.0 
         
-        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
+        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
 
-        # --- 4. DRAWING ---
         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)
         
@@ -288,22 +286,23 @@ 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 + min(w, h)*0.08 * np.cos(rad))), int(round(ty + min(w, h)*0.08 * np.sin(rad))))
+            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)
 
-        _, 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
         )
 
-
 app = FastAPI()
 
 app.add_middleware(
@@ -318,17 +317,18 @@ def read_root(): return {"status": "Watermelon API is awake!"}
 @app.post("/process_single")
 async def process_single(file: UploadFile = File(...)):
     contents = await file.read()
-    img = cv2.imdecode(np.frombuffer(contents, np.uint8), cv2.IMREAD_COLOR)
+    nparr = np.frombuffer(contents, np.uint8)
+    img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
     
-    scale_ratio = 1.0
     h, w = img.shape[:2]
+    scale_ratio = 1.0
     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, contents
+    del img, nparr, contents
     gc.collect()
 
     return res.__dict__