main.py

import os
import cv2
import numpy as np
import base64
import gc
import torch
import time
import traceback
from dataclasses import dataclass
from typing import Optional
from scipy.ndimage import median_filter
from scipy.optimize import curve_fit, minimize
from ultralytics import YOLO
from fastapi import FastAPI, UploadFile, File, Query, Form
from fastapi.middleware.cors import CORSMiddleware
from fastapi.concurrency import run_in_threadpool
import uvicorn
from skimage import color
import hashlib
import asyncio

# OOM PREVENTION
torch.set_num_threads(1)

# --- CONFIGURATION ---
MODEL_PATH = "best.pt"
MAX_IMAGE_SIZE = 2048  
CHECKER_WIDTH_CM = 6.3 
MIN_RIND_FLESH_OVERLAP_RATIO = 0.10
MIN_FLESH_PIXELS_FOR_FALLBACK = 100
MAX_RIND_TO_FLESH_AREA_RATIO = 3.5
MAX_RIND_CENTER_OFFSET_RATIO = 0.60

# ==============================================================================
# --- COLOR CALIBRATION LOGIC ---
# ==============================================================================
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 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)
        
    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 
        
    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]])
    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):
    tgt_lin = to_linear_srgb(target_bgr)
    
    # 1. White Balance (Von Kries scaling)
    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. 3x3 Matrix Solver
    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))
        
    X, Y = tgt24_wb, ref24
    W_init = np.linalg.inv(X.T @ X + 0.05 * np.eye(3)) @ X.T @ Y
    res = minimize(
        objective,
        W_init.flatten(),
        method='Powell',
        options={"maxiter": 150, "xtol": 1e-4, "ftol": 1e-4},
    )
    W_opt = res.x.reshape(3, 3) 

    # 3. Apply to full image
    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))

# ==============================================================================
# --- CORE API & PROCESSOR ---
# ==============================================================================

@dataclass
class ProcessResult:
    success: bool
    message: str
    r2_rind: Optional[float] = None
    r2_flesh: Optional[float] = None
    raw_width: Optional[float] = None
    sm_width: Optional[float] = None
    raw_height: Optional[float] = None
    sm_height: Optional[float] = None
    raw_perimeter: Optional[float] = None
    sm_perimeter: Optional[float] = None
    raw_flesh_width: Optional[float] = None 
    sm_flesh_width: Optional[float] = None   
    raw_flesh_height: Optional[float] = None 
    sm_flesh_height: Optional[float] = None  
    raw_flesh_perimeter: Optional[float] = None 
    sm_flesh_perimeter: Optional[float] = None  
    raw_rind_thick: Optional[float] = None
    sm_rind_thick: Optional[float] = None
    raw_rind_ratio: Optional[float] = None
    sm_rind_ratio: Optional[float] = None
    raw_total_area: Optional[float] = None
    sm_total_area: Optional[float] = None
    raw_flesh_area: Optional[float] = None
    sm_flesh_area: Optional[float] = None
    raw_flesh_ratio: Optional[float] = None
    sm_flesh_ratio: Optional[float] = None
    raw_elongation: Optional[float] = None
    sm_elongation: Optional[float] = None
    raw_asym: Optional[float] = None
    sm_asym: Optional[float] = None
    raw_flesh_asym: Optional[float] = None
    sm_flesh_asym: Optional[float] = None
    raw_circ: Optional[float] = None
    sm_circ: Optional[float] = None
    midline_curvature: Optional[float] = None
    delta_e_initial: Optional[float] = None
    delta_e_final: Optional[float] = None
    image_raw_base64: Optional[str] = None
    image_sm_base64: Optional[str] = None
    filename: Optional[str] = None
    measurement_unit: Optional[str] = None
    area_unit: Optional[str] = None
    scale_source: Optional[str] = None
    color_checker_found: bool = False
    rind_source: Optional[str] = None
    rind_overlap_ratio: Optional[float] = None
    warnings: Optional[list] = None
    timings_ms: Optional[dict] = None
    processing_ms: Optional[int] = None

class WatermelonProcessor:
    def __init__(self, model_path: str):
        self.model = YOLO(model_path)
        self.ref24 = None
        if os.path.exists("reference.png"):
            try:
                ref_img = cv2.imread("reference.png")
                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"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):
        t = theta - phi
        ellipse = (Rx * Ry) / np.sqrt((Ry * np.cos(t)) ** 2 + (Rx * np.sin(t)) ** 2)
        asymmetry = 1 + c_a * np.cos(t) ** 3
        divot_top = d_top * np.exp(w_top * (np.sin(t) - 1))
        divot_bot = d_bot * np.exp(w_bot * (-np.sin(t) - 1))
        return (ellipse * asymmetry) - divot_top - divot_bot + c_skew * np.sin(t) + c_bend * np.cos(t) * (np.sin(t) ** 2)

    @staticmethod
    def calculate_axis_metrics(cx, cy, phi, rind_mask, flesh_mask):
        """Instantly finds axes and rind thickness using fast OpenCV bitwise operations."""
        h, w = rind_mask.shape

        def get_intersections(theta, mask):
            temp = np.zeros((h, w), dtype=np.uint8)
            L = max(h, w)
            # Draw a line spanning across the entire image
            p1 = (int(cx + L * np.cos(theta)), int(cy - L * np.sin(theta)))
            p2 = (int(cx - L * np.cos(theta)), int(cy + L * np.sin(theta)))
            cv2.line(temp, p1, p2, 255, 1)
            
            # Find where the line overlaps the mask
            overlap = cv2.bitwise_and(mask, temp)
            y_pts, x_pts = np.where(overlap > 0)
            
            if len(x_pts) == 0:
                return (int(cx), int(cy)), (int(cx), int(cy)), 0.0
                
            # Project points to find the two extreme ends of the line segment
            dx, dy = x_pts - cx, y_pts - cy
            proj = dx * np.cos(theta) - dy * np.sin(theta)
            
            idx_max, idx_min = np.argmax(proj), np.argmin(proj)
            pt1 = (int(x_pts[idx_max]), int(y_pts[idx_max]))
            pt2 = (int(x_pts[idx_min]), int(y_pts[idx_min]))
            dist = float(np.hypot(pt1[0] - pt2[0], pt1[1] - pt2[1]))
            
            return pt1, pt2, dist

        # Height line (perpendicular to phi)
        pt_top, pt_bot, height_px = get_intersections(phi + np.pi/2, rind_mask)
        
        # Width line (parallel to phi)
        pt_right, pt_left, width_px = get_intersections(phi, rind_mask)
        
        # Flesh width along the exact same width line
        _, _, flesh_width_px = get_intersections(phi, flesh_mask)
        
        rind_thick_px = None
        if width_px > 0 and flesh_width_px > 0:
            rind_thick_px = float(max(0.0, (width_px - flesh_width_px) / 2.0))

        return height_px, width_px, rind_thick_px, (pt_top, pt_bot), (pt_left, pt_right)
    
    @staticmethod
    def contour_centroid(contour):
        M = cv2.moments(contour)
        if M["m00"] == 0:
            return None
        return np.array([M["m10"] / M["m00"], M["m01"] / M["m00"]], dtype=np.float32)

    @staticmethod
    def mask_centroid(mask):
        M = cv2.moments(mask)
        if M["m00"] == 0:
            return None
        return np.array([M["m10"] / M["m00"], M["m01"] / M["m00"]], dtype=np.float32)

    @staticmethod
    def draw_single_contour(shape, contour):
        mask = np.zeros(shape, dtype=np.uint8)
        cv2.drawContours(mask, [contour.astype(np.int32)], -1, 255, -1)
        return mask

    @staticmethod
    def flesh_envelope_mask(flesh_combined):
        ys, xs = np.where(flesh_combined > 0)
        if len(xs) < MIN_FLESH_PIXELS_FOR_FALLBACK:
            return None

        pts = np.column_stack([xs, ys]).astype(np.int32)
        hull = cv2.convexHull(pts.reshape(-1, 1, 2))
        envelope = np.zeros_like(flesh_combined)
        cv2.drawContours(envelope, [hull], -1, 255, -1)

        _, _, bw, bh = cv2.boundingRect(hull)
        pad = int(max(12, min(80, round(max(bw, bh) * 0.035))))
        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (pad * 2 + 1, pad * 2 + 1))
        envelope = cv2.dilate(envelope, kernel, iterations=1)

        close_size = max(5, (pad // 2) * 2 + 1)
        close_kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (close_size, close_size))
        return cv2.morphologyEx(envelope, cv2.MORPH_CLOSE, close_kernel)

    @staticmethod
    def choose_target_rind_mask(rind_mask, flesh_combined):
        warnings = []
        flesh_area = cv2.countNonZero(flesh_combined)
        cnts, _ = cv2.findContours(rind_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)

        if not cnts:
            envelope = WatermelonProcessor.flesh_envelope_mask(flesh_combined)
            if envelope is None:
                return rind_mask, "missing", None, ["No whole-watermelon mask and not enough flesh mask for fallback."]
            return envelope, "flesh_envelope", 1.0, ["No whole-watermelon mask; estimated perimeter from flesh masks."]

        scored = []
        for cnt in cnts:
            temp = WatermelonProcessor.draw_single_contour(rind_mask.shape, cnt)
            overlap = cv2.countNonZero(cv2.bitwise_and(temp, flesh_combined))
            ratio = overlap / max(flesh_area, 1)
            scored.append((ratio, cv2.contourArea(cnt), cnt, temp))

        scored.sort(key=lambda item: (item[0], item[1]), reverse=True)
        best_ratio, best_area, best_cnt, best_mask = scored[0]

        if flesh_area == 0:
            warnings.append("No flesh masks detected; using largest whole-watermelon mask.")
            largest = max(cnts, key=cv2.contourArea)
            return WatermelonProcessor.draw_single_contour(rind_mask.shape, largest), "whole_mask_no_flesh", None, warnings

        flesh_center = WatermelonProcessor.mask_centroid(flesh_combined)
        rind_center = WatermelonProcessor.contour_centroid(best_cnt)
        ys, xs = np.where(flesh_combined > 0)
        flesh_extent = max(float(np.ptp(xs)) if len(xs) else 1.0, float(np.ptp(ys)) if len(ys) else 1.0, 1.0)
        center_offset_ratio = 0.0
        if flesh_center is not None and rind_center is not None:
            center_offset_ratio = float(np.linalg.norm(flesh_center - rind_center) / flesh_extent)
        area_ratio = float(best_area / max(flesh_area, 1))

        if best_ratio >= MIN_RIND_FLESH_OVERLAP_RATIO:
            if area_ratio <= MAX_RIND_TO_FLESH_AREA_RATIO and center_offset_ratio <= MAX_RIND_CENTER_OFFSET_RATIO:
                return best_mask, "whole_mask_overlap", float(best_ratio), warnings
            warnings.append("Whole-watermelon mask overlapped flesh but looked too large or off-center; using fallback.")

        if flesh_center is not None and rind_center is not None:
            shifted_cnt = best_cnt.astype(np.float32) + (flesh_center - rind_center).reshape(1, 1, 2)
            shifted_mask = WatermelonProcessor.draw_single_contour(rind_mask.shape, shifted_cnt)
            shifted_ratio = cv2.countNonZero(cv2.bitwise_and(shifted_mask, flesh_combined)) / max(flesh_area, 1)
            if shifted_ratio >= MIN_RIND_FLESH_OVERLAP_RATIO and area_ratio <= MAX_RIND_TO_FLESH_AREA_RATIO:
                if best_ratio >= MIN_RIND_FLESH_OVERLAP_RATIO:
                    warnings.append("Whole-watermelon mask was suspicious; translated it to the flesh-mask centroid.")
                else:
                    warnings.append("Whole-watermelon mask did not overlap flesh; translated it to the flesh-mask centroid.")
                return shifted_mask, "translated_whole_mask", float(shifted_ratio), warnings

        envelope = WatermelonProcessor.flesh_envelope_mask(flesh_combined)
        if envelope is not None:
            warnings.append("Whole-watermelon mask did not overlap flesh; estimated perimeter from flesh masks.")
            return envelope, "flesh_envelope", float(best_ratio), warnings

        warnings.append("Whole-watermelon mask did not overlap flesh and fallback was unavailable.")
        return best_mask, "low_overlap_whole_mask", float(best_ratio), warnings

    @staticmethod
    def contour_area_px(points):
        cnt = points.reshape(-1, 1, 2).astype(np.float32)
        return float(abs(cv2.contourArea(cnt)))

    @staticmethod
    def elongation_from_points(points):
        cnt = points.reshape(-1, 1, 2).astype(np.float32)
        if len(cnt) >= 5:
            _, (axis_a, axis_b), _ = cv2.fitEllipse(cnt)
            minor = max(min(axis_a, axis_b), 1e-6)
            return float(max(axis_a, axis_b) / minor)

        _, _, bw, bh = cv2.boundingRect(cnt.astype(np.int32))
        return float(max(bw, bh) / max(min(bw, bh), 1))

    @staticmethod
    def split_asymmetry(region_mask, midline, thickness=5):
        if midline is None or len(midline) < 2 or cv2.countNonZero(region_mask) == 0: return None
        split_mask = region_mask.copy()
        
        # FIX: Extrapolate the line 100 pixels in both directions to guarantee it completely bisects smoothed/expanded masks
        m = midline.astype(np.float32)
        p0, p1, pn, pn_1 = m[0], m[1], m[-1], m[-2]
        n0, n1 = np.linalg.norm(p0 - p1), np.linalg.norm(pn - pn_1)
        ext_start = p0 + (p0 - p1) / n0 * 100 if n0 > 1e-5 else p0
        ext_end = pn + (pn - pn_1) / n1 * 100 if n1 > 1e-5 else pn
        ext_midline = np.vstack([ext_start, m, ext_end]).astype(np.int32)
        
        cv2.polylines(split_mask, [ext_midline], False, 0, thickness)
        n_labels, _, stats, _ = cv2.connectedComponentsWithStats((split_mask > 0).astype(np.uint8), connectivity=8)
        if n_labels <= 2: return None

        areas = sorted([int(stats[i, cv2.CC_STAT_AREA]) for i in range(1, n_labels)], reverse=True)
        if len(areas) < 2 or areas[0] + areas[1] == 0: return None
        return float(abs(areas[0] - areas[1]) / (areas[0] + areas[1]))

    @staticmethod
    def midline_curvature_score(midline):
        if midline is None or len(midline) < 3:
            return None

        diffs = np.diff(midline.astype(np.float32), axis=0)
        path_len = float(np.sum(np.linalg.norm(diffs, axis=1)))
        chord_len = float(np.linalg.norm(midline[-1] - midline[0]))
        if chord_len <= 1e-6:
            return None

        return float(max(0.0, (path_len / chord_len) - 1.0))

    @staticmethod
    def get_polar_data(mask):
        """Universal polar extractor for either Rind or Flesh masks."""
        cnts, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
        if not cnts: return None
        cnt = max(cnts, key=cv2.contourArea)
        M = cv2.moments(cnt)
        if M["m00"] == 0: return None

        cx, cy = M["m10"]/M["m00"], M["m01"]/M["m00"]
        pts = 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)
        for i in range(num_bins):
            b_mask = (t_vals >= bins[i]) & (t_vals < bins[i + 1])
            if np.any(b_mask): raw_r[i] = np.max(r_vals[b_mask])

        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)
        return (bins[:-1] + bins[1:])/2.0, median_filter(raw_r, size=7, mode="wrap"), (cx, cy), cnt

    @staticmethod
    def calculate_axis_metrics(cx, cy, phi, rind_mask, flesh_mask):
        """Instantly finds axes and rind thickness using fast OpenCV bitwise operations."""
        h, w = rind_mask.shape
        def get_intersections(theta, mask):
            temp = np.zeros((h, w), dtype=np.uint8)
            L = max(h, w)
            p1 = (int(cx + L * np.cos(theta)), int(cy - L * np.sin(theta)))
            p2 = (int(cx - L * np.cos(theta)), int(cy + L * np.sin(theta)))
            cv2.line(temp, p1, p2, 255, 1)
            
            overlap = cv2.bitwise_and(mask, temp)
            y_pts, x_pts = np.where(overlap > 0)
            if len(x_pts) == 0: return (int(cx), int(cy)), (int(cx), int(cy)), 0.0
                
            dx, dy = x_pts - cx, y_pts - cy
            proj = dx * np.cos(theta) - dy * np.sin(theta)
            idx_max, idx_min = np.argmax(proj), np.argmin(proj)
            pt1 = (int(x_pts[idx_max]), int(y_pts[idx_max]))
            pt2 = (int(x_pts[idx_min]), int(y_pts[idx_min]))
            dist = float(np.hypot(pt1[0] - pt2[0], pt1[1] - pt2[1]))
            return pt1, pt2, dist

        pt_top, pt_bot, height_px = get_intersections(phi + np.pi/2, rind_mask)
        pt_right, pt_left, width_px = get_intersections(phi, rind_mask)
        _, _, f_height_px = get_intersections(phi + np.pi/2, flesh_mask) 
        _, _, f_width_px = get_intersections(phi, flesh_mask)
        
        rind_thick_px = None
        if width_px > 0 and f_width_px > 0:
            rind_thick_px = float(max(0.0, (width_px - f_width_px) / 2.0))

        return height_px, width_px, rind_thick_px, (pt_top, pt_bot), (pt_left, pt_right), f_height_px, f_width_px

    @staticmethod
    def get_dual_mask_midline(f_left, f_right, rind_cnt, pred_cnt, cx, cy):
        h, w = f_left.shape
        _, (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, r_rot = cv2.warpAffine(f_left, m_rot, (w, h)), cv2.warpAffine(f_right, m_rot, (w, h))

        l_idx, r_idx = np.where(l_rot > 0)[1], 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_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_g, y_max_g = np.min(gap_points[:, 0]), np.max(gap_points[:, 0])
            y_span, y_mean = max(y_max_g - y_min_g, 1), (y_max_g + y_min_g) / 2.0
            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,
                    (gap_points[:,0]-y_mean)/y_span,
                    gap_points[:,1],
                    bounds=([-max_bend, -np.inf, -np.inf],[max_bend, np.inf, np.inf]),
                    max_nfev=1500,
                )
            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(ys_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)])

        pts_orig = (m_inv @ pts_rot).T
        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, include_image: bool = True) -> ProcessResult:
        timings = {}
        stage_t = time.perf_counter()

        def mark(stage_name):
            nonlocal stage_t
            now = time.perf_counter()
            timings[stage_name] = int(round((now - stage_t) * 1000))
            stage_t = now

        def fail(message, **extra):
            return ProcessResult(success=False, message=message, filename=source_name, warnings=warnings or None, timings_ms=timings, **extra)

        warnings =[]
        if image is None: return ProcessResult(success=False, message="Could not decode image.", filename=source_name)

        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)
            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)

            if self.ref24 is not None:
                tgt_warped = warp_checker(image, checker_corners)
                tgt24 = sample_24_patches(tgt_warped)
                dE_initial = float(np.mean(compute_deltaE_00(tgt24, self.ref24)))
                image = apply_color_pipeline(image, self.ref24, tgt24)
                tgt_warped_corr = warp_checker(image, checker_corners)
                tgt24_corr = sample_24_patches(tgt_warped_corr)
                dE_final = float(np.mean(compute_deltaE_00(tgt24_corr, self.ref24)))
        except Exception as e:
            if cm_per_px is None: warnings.append("ColorChecker not found; dimensions in original-image pixels.")
            else: warnings.append("Color correction skipped after ColorChecker detection.")
        mark("calibration")

        # 2. YOLO INFERENCE
        results = self.model(image, conf=0.25, retina_masks=True, verbose=False)
        mark("yolo_inference")

        rind_mask = np.zeros((h, w), dtype=np.uint8)
        f_l_cnts, f_r_cnts = [], []
        if results[0].masks is None:
            return fail("No masks detected.", measurement_unit="cm" if cm_per_px is not None else "px", scale_source="color_checker" if cm_per_px is not None else "original_pixels", color_checker_found=checker_corners is not None)

        for mask_data, cls in zip(results[0].masks.xy, results[0].boxes.cls):
            c = np.array(mask_data, dtype=np.int32)
            if int(cls) == 0: cv2.drawContours(rind_mask, [c], -1, 255, -1)
            elif int(cls) == 1: f_l_cnts.append(c)
            elif int(cls) == 2: f_r_cnts.append(c)

        if len(f_l_cnts) >= 2 and len(f_r_cnts) == 0:
            f_l_cnts.sort(key=lambda cnt: cv2.moments(cnt)["m10"] / (cv2.moments(cnt)["m00"] + 1e-5))
            f_r_cnts.append(f_l_cnts.pop())
            warnings.append("Only flesh_left was detected; split by x-position.")
        elif len(f_r_cnts) >= 2 and len(f_l_cnts) == 0:
            f_r_cnts.sort(key=lambda cnt: cv2.moments(cnt)["m10"] / (cv2.moments(cnt)["m00"] + 1e-5))
            f_l_cnts.append(f_r_cnts.pop(0))
            warnings.append("Only flesh_right was detected; split by x-position.")

        flesh_l_m, flesh_r_m = np.zeros((h, w), dtype=np.uint8), np.zeros((h, w), dtype=np.uint8)
        for c in f_l_cnts: cv2.drawContours(flesh_l_m, [c], -1, 255, -1)
        for c in f_r_cnts: cv2.drawContours(flesh_r_m, [c], -1, 255, -1)
        flesh_combined = cv2.bitwise_or(flesh_l_m, flesh_r_m)

        target_rind_mask, rind_source, rind_overlap_ratio, r_warn = self.choose_target_rind_mask(rind_mask, flesh_combined)
        warnings.extend(r_warn)
        
        mark("mask_parse")

        # 3. EXTRACTION (Rind First)
        rind_data = self.get_polar_data(target_rind_mask)
        if rind_data is None:
            return fail("No stable perimeter.", measurement_unit="cm" if cm_per_px else "px", scale_source="color_checker" if cm_per_px else "original_pixels", color_checker_found=checker_corners is not None, rind_source=rind_source, rind_overlap_ratio=rind_overlap_ratio)

        t_r, raw_r_r, (cx, cy), raw_rind_cnt = rind_data
        pts_r_raw = raw_rind_cnt.reshape(-1, 2).astype(np.float32)

        # --- SCAN-LINE FLESH GAP FILLING (Preserves V-shape) ---
        _, (ma, Ma), angle = cv2.fitEllipse(raw_rind_cnt) if len(raw_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(flesh_l_m, M_rot, (w, h))
        r_rot = cv2.warpAffine(flesh_r_m, 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 and np.mean(l_idx) > np.mean(r_idx):
            l_rot, r_rot = r_rot, l_rot
            
        flesh_closed_rot = cv2.bitwise_or(l_rot, r_rot)
        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 = np.where(l_rot[y, :] > 0)[0]
                row_r = np.where(r_rot[y, :] > 0)[0]
                if len(row_l) > 0 and len(row_r) > 0:
                    x_start = row_l[-1]
                    x_end = row_r[0]
                    if x_start < x_end:
                        flesh_closed_rot[y, x_start:x_end] = 255

        flesh_closed = cv2.warpAffine(flesh_closed_rot, M_inv, (w, h))
        _, flesh_closed = cv2.threshold(flesh_closed, 127, 255, cv2.THRESH_BINARY)

        # NOW we extract the flesh_data!
        flesh_data = self.get_polar_data(flesh_closed)

        raw_flesh_perim = None
        if flesh_data:
            _, _, _, raw_flesh_cnt = flesh_data
            raw_flesh_perim = float(cv2.arcLength(raw_flesh_cnt, True))

        # 4. RAW FEATURES
        _, _, r_angle = cv2.fitEllipse(raw_rind_cnt)
        raw_phi = np.deg2rad(180 - r_angle) if r_angle > 90 else np.deg2rad(-r_angle)
        
        # Capture raw_f_h and raw_f_w
        raw_h, raw_w, raw_rt, raw_h_line, raw_w_line, raw_f_h, raw_f_w = self.calculate_axis_metrics(cx, cy, raw_phi, target_rind_mask, flesh_closed)
        
        # Midline is found using flesh_combined (with gap) and clipped to the raw rind contour
        midline = self.get_dual_mask_midline(flesh_l_m, flesh_r_m, raw_rind_cnt, pts_r_raw, cx, cy)
        
        raw_perim = float(cv2.arcLength(raw_rind_cnt, True))
        raw_tot_a = self.contour_area_px(pts_r_raw)
        raw_f_a = float(cv2.countNonZero(flesh_combined))
        raw_f_rat = float(raw_f_a / raw_tot_a) if raw_tot_a > 0 else None
        raw_elong = self.elongation_from_points(pts_r_raw)
        raw_circ = float((4.0 * np.pi * raw_tot_a) / (raw_perim ** 2)) if raw_perim > 0 and raw_tot_a > 0 else None
        raw_asym = self.split_asymmetry(target_rind_mask, midline)
        raw_f_asym = self.split_asymmetry(flesh_combined, midline, thickness=3)
        midline_curve = self.midline_curvature_score(midline)

        # 5. SMOOTHED FEATURES
        sm_w, sm_h, sm_perim, sm_rt, sm_tot_a, sm_f_a, sm_f_rat = None, None, None, None, None, None, None
        sm_elong, sm_circ, sm_asym, sm_f_asym, r2_rind, r2_flesh = None, None, None, None, None, None
        sm_rind_cnt, sm_flesh_cnt = None, None
        sm_h_line, sm_w_line = None, None
        
        try:
            # Fit Rind
            scale_r = np.mean(raw_r_r)
            popt_r, _ = curve_fit(self.watermelon_model, t_r, raw_r_r/scale_r, 
                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]), max_nfev=3000)
            d_r = np.sum((raw_r_r/scale_r - 1)**2)
            r2_rind = 1 - (np.sum((raw_r_r/scale_r - self.watermelon_model(t_r, *popt_r))**2) / d_r) if d_r != 0 else None

            # R² Warning Flag
            if r2_rind is not None and r2_rind < 0.85:
                warnings.append(f"R² below 0.85 ({r2_rind:.2f}). Fruit may be damaged or irregular.")
                
            t_fit = np.linspace(-np.pi, np.pi, 500)
            fit_r = self.watermelon_model(t_fit, *popt_r) * scale_r
            sm_rind_pts = np.array([[r*np.cos(t)+cx, cy-r*np.sin(t)] for t, r in zip(t_fit, fit_r)], dtype=np.float32)
            sm_rind_cnt = sm_rind_pts.reshape(-1, 1, 2).astype(np.int32)
            
            sm_rind_mask = np.zeros_like(target_rind_mask)
            cv2.fillPoly(sm_rind_mask, [sm_rind_cnt], 255)
            
            sm_perim = float(np.sum(np.linalg.norm(np.diff(sm_rind_pts, axis=0), axis=1)) + np.linalg.norm(sm_rind_pts[-1]-sm_rind_pts[0]))
            sm_tot_a = self.contour_area_px(sm_rind_pts)
            sm_elong = self.elongation_from_points(sm_rind_pts)
            sm_circ = float((4.0 * np.pi * sm_tot_a) / (sm_perim ** 2)) if sm_perim > 0 and sm_tot_a > 0 else None
            sm_asym = self.split_asymmetry(sm_rind_mask, midline)
            sm_phi = popt_r[7]

            # Fit Flesh
            if flesh_data:
                t_f, raw_r_f, (fcx, fcy), _ = flesh_data
                scale_f = np.mean(raw_r_f)
                popt_f, _ = curve_fit(self.watermelon_model, t_f, raw_r_f/scale_f, 
                    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]), max_nfev=3000)
                d_f = np.sum((raw_r_f/scale_f - 1)**2)
                r2_flesh = 1 - (np.sum((raw_r_f/scale_f - self.watermelon_model(t_f, *popt_f))**2) / d_f) if d_f != 0 else None
                
                fit_f = self.watermelon_model(t_fit, *popt_f) * scale_f
                sm_flesh_pts = np.array([[r*np.cos(t)+fcx, fcy-r*np.sin(t)] for t, r in zip(t_fit, fit_f)], dtype=np.float32)
                sm_flesh_cnt = sm_flesh_pts.reshape(-1, 1, 2).astype(np.int32)
                
                sm_flesh_mask = np.zeros_like(target_rind_mask)
                cv2.fillPoly(sm_flesh_mask, [sm_flesh_cnt], 255)
                
                sm_f_a = self.contour_area_px(sm_flesh_pts)
                sm_flesh_asym = self.split_asymmetry(sm_flesh_mask, midline, thickness=3)
                sm_flesh_perim = float(np.sum(np.linalg.norm(np.diff(sm_flesh_pts, axis=0), axis=1)) + np.linalg.norm(sm_flesh_pts[-1]-sm_flesh_pts[0]))
            else:
                sm_flesh_mask = np.zeros_like(target_rind_mask)
                sm_flesh_perim = None

            # Smooth axes (using the filled smooth masks)
            sm_h, sm_w, sm_rt, sm_h_line, sm_w_line, sm_f_h, sm_f_w = self.calculate_axis_metrics(cx, cy, sm_phi, sm_rind_mask, sm_flesh_mask)
            if sm_tot_a > 0 and sm_f_a is not None:
                sm_f_rat = float(sm_f_a / sm_tot_a)

        except Exception as exc:
            warnings.append(f"Smoothing fit failed: {exc}")
        mark("fit")

        # 6. APPLY SCALES
        sc_src = "color_checker" if cm_per_px else "original_pixels"
        m_unit = "cm" if cm_per_px else "px"
        a_unit = "cm2" if cm_per_px else "px2"
        scaler = cm_per_px if cm_per_px else (1.0 / scale_ratio)
        a_scaler = scaler ** 2

        def s(v): return float(v * scaler) if v is not None else None
        def a(v): return float(v * a_scaler) if v is not None else None
        def rt_rat(thick, w): return float((thick * 2.0) / w) if thick is not None and w and w > 0 else None

        res = ProcessResult(
            success=True, message="Success", filename=source_name, measurement_unit=m_unit, area_unit=a_unit,
            scale_source=sc_src, color_checker_found=bool(cm_per_px),
            rind_source=rind_source, rind_overlap_ratio=rind_overlap_ratio, warnings=warnings or None,
            r2_rind=r2_rind, r2_flesh=r2_flesh, midline_curvature=midline_curve,
            delta_e_initial=dE_initial, delta_e_final=dE_final, timings_ms=timings,
            
            raw_width=s(raw_w), sm_width=s(sm_w), raw_height=s(raw_h), sm_height=s(sm_h),
            raw_perimeter=s(raw_perim), sm_perimeter=s(sm_perim),
            raw_flesh_width=s(raw_f_w), sm_flesh_width=s(sm_f_w),
            raw_flesh_height=s(raw_f_h), sm_flesh_height=s(sm_f_h),
            raw_flesh_perimeter=s(raw_flesh_perim), sm_flesh_perimeter=s(sm_flesh_perim),
            raw_rind_thick=s(raw_rt), sm_rind_thick=s(sm_rt),
            raw_rind_ratio=rt_rat(raw_rt, raw_w), sm_rind_ratio=rt_rat(sm_rt, sm_w),
            raw_total_area=a(raw_tot_a), sm_total_area=a(sm_tot_a),
            raw_flesh_area=a(raw_f_a), sm_flesh_area=a(sm_f_a),
            raw_flesh_ratio=raw_f_rat, sm_flesh_ratio=sm_f_rat,
            raw_elongation=raw_elong, sm_elongation=sm_elong,
            raw_asym=raw_asym, sm_asym=sm_asym, raw_flesh_asym=raw_f_asym, sm_flesh_asym=sm_flesh_asym,
            raw_circ=raw_circ, sm_circ=sm_circ
        )

        # 7. DRAW PREVIEWS
        if include_image:
            def encode_img(canvas):
                _, b = cv2.imencode(".jpg", canvas, [cv2.IMWRITE_JPEG_QUALITY, 85])
                return base64.b64encode(b).decode("utf-8")
                
            def draw_base(r_m, f_m):
                out = image.copy().astype(np.float32)
                alpha = 0.18 # Transparency
                
                # --- THE FIX: Only paint the rind where there is NO flesh ---
                rind_only = (r_m > 0) & (f_m == 0)
                
                # Rind fill: Green
                out[..., 0] = np.where(rind_only, out[..., 0]*(1-alpha) + 0, out[..., 0])
                out[..., 1] = np.where(rind_only, out[..., 1]*(1-alpha) + 170, out[..., 1])
                out[..., 2] = np.where(rind_only, out[..., 2]*(1-alpha) + 0, out[..., 2])
                
                # Flesh fill: Red
                out[..., 0] = np.where(f_m > 0, out[..., 0]*(1-alpha) + 0, out[..., 0])
                out[..., 1] = np.where(f_m > 0, out[..., 1]*(1-alpha) + 0, out[..., 1])
                out[..., 2] = np.where(f_m > 0, out[..., 2]*(1-alpha) + 255, out[..., 2])
                
                out = np.clip(out, 0, 255).astype(np.uint8)
                if checker_corners is not None: 
                    cv2.polylines(out, [np.int32(checker_corners)], True, (0, 165, 255), 4)
                if len(midline) > 1: 
                    cv2.polylines(out, [midline.astype(np.int32)], False, (0, 255, 255), 3)
                    pt1, pt2 = tuple(midline[0].astype(int)), tuple(midline[-1].astype(int))
                    for pt in (pt1, pt2):
                        cv2.circle(out, pt, 8, (0,0,0), 2)
                        cv2.circle(out, pt, 6, (255,255,255), -1)
                return out

            # RAW
            out_raw = draw_base(target_rind_mask, flesh_closed) 
            cv2.line(out_raw, raw_h_line[0], raw_h_line[1], (255, 100, 255), 2)
            cv2.line(out_raw, raw_w_line[0], raw_w_line[1], (255, 255, 100), 2)
            
            f_cnts_raw, _ = cv2.findContours(flesh_closed, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
            if f_cnts_raw: 
                cv2.polylines(out_raw, [max(f_cnts_raw, key=cv2.contourArea)], True, (0, 0, 255), 2) # Red, 2px
            cv2.polylines(out_raw, [raw_rind_cnt], True, (0, 200, 0), 2) # Dark Green, 2px
            res.image_raw_base64 = encode_img(out_raw)

            # SMOOTH
            if sm_rind_cnt is not None:
                # Use raw fills for the smoothed preview (decoupled visual)
                out_sm = draw_base(target_rind_mask, flesh_closed) 
                cv2.line(out_sm, sm_h_line[0], sm_h_line[1], (255, 100, 255), 2)
                cv2.line(out_sm, sm_w_line[0], sm_w_line[1], (255, 255, 100), 2)
                
                if sm_flesh_cnt is not None: 
                    cv2.polylines(out_sm, [sm_flesh_cnt], True, (0, 0, 255), 2) # Red, 2px
                cv2.polylines(out_sm, [sm_rind_cnt], True, (0, 200, 0), 2) # Dark Green, 2px
                res.image_sm_base64 = encode_img(out_sm)
            
        mark("render")
        return res
app = FastAPI()

app.add_middleware(
    CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"],
)

processor = WatermelonProcessor(MODEL_PATH)

# --- CONCURRENCY & QUEUE MANAGEMENT ---
dev_lock = asyncio.Lock()
gen_lock = asyncio.Lock()
dev_queue_count = 0
gen_queue_count = 0

@app.get("/")
def read_root(): return {"status": "Watermelon API is awake and running!"}

@app.get("/queue_status")
def get_queue_status():
    return {"dev_queue": dev_queue_count, "gen_queue": gen_queue_count}

@app.post("/process_single")
async def process_single(
    file: UploadFile = File(...), 
    include_image: bool = Query(True), 
    password: str = Form(""),
    username: str = Form("")
):
    global dev_queue_count, gen_queue_count
    request_t = time.perf_counter()
    
    expected_hash = "9139eb3676d5dfafced7613f044d86d9e7c84f40a04c83ddce062878621315d0"
    if hashlib.sha256(password.encode('utf-8')).hexdigest() != expected_hash:
        return ProcessResult(success=False, message="Unauthorized.", filename=file.filename).__dict__

    contents, img = None, None

    try:
        contents = await file.read()
        if not contents: return ProcessResult(success=False, message="Empty.", filename=file.filename).__dict__

        img = cv2.imdecode(np.frombuffer(contents, np.uint8), cv2.IMREAD_COLOR)
        if img is None: return ProcessResult(success=False, message="Decode error.", filename=file.filename).__dict__

        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)

        # --- CPU CORE ROUTING & QUEUE LOGIC ---
        is_dev = (username.strip().lower() == 'devtest')
        
        # run_in_threadpool prevents OpenCV/YOLO from freezing the API so /queue_status can still answer
        if is_dev:
            dev_queue_count += 1
            try:
                async with dev_lock:
                    res = await run_in_threadpool(processor.process_image, img, file.filename, scale_ratio, include_image)
            finally:
                dev_queue_count -= 1
        else:
            gen_queue_count += 1
            try:
                async with gen_lock:
                    res = await run_in_threadpool(processor.process_image, img, file.filename, scale_ratio, include_image)
            finally:
                gen_queue_count -= 1

        res.processing_ms = int(round((time.perf_counter() - request_t) * 1000))
        return res.__dict__

    except Exception as exc:
        traceback.print_exc()
        return ProcessResult(success=False, message=str(exc), filename=file.filename).__dict__

    finally:
        del img, contents
        gc.collect()