crabbly's picture
Update main.py
56cb810 verified
Raw
History Blame Contribute Delete
41.3 kB
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()