crabbly commited on
Commit
2f18545
·
verified ·
1 Parent(s): e372edd

Update main.py

Browse files
Files changed (1) hide show
  1. main.py +291 -87
main.py CHANGED
@@ -4,12 +4,14 @@ import numpy as np
4
  import base64
5
  import gc
6
  import torch
 
 
7
  from dataclasses import dataclass
8
  from typing import Optional
9
  from scipy.ndimage import median_filter
10
  from scipy.optimize import curve_fit, minimize
11
  from ultralytics import YOLO
12
- from fastapi import FastAPI, UploadFile, File
13
  from fastapi.middleware.cors import CORSMiddleware
14
  import uvicorn
15
  from skimage import color
@@ -21,6 +23,10 @@ torch.set_num_threads(1)
21
  MODEL_PATH = "best.pt"
22
  MAX_IMAGE_SIZE = 2048
23
  CHECKER_WIDTH_CM = 6.3
 
 
 
 
24
 
25
  # ==============================================================================
26
  # --- COLOR CALIBRATION LOGIC ---
@@ -84,7 +90,12 @@ def apply_color_pipeline(target_bgr, ref24, tgt24):
84
 
85
  X, Y = tgt24_wb, ref24
86
  W_init = np.linalg.inv(X.T @ X + 0.05 * np.eye(3)) @ X.T @ Y
87
- res = minimize(objective, W_init.flatten(), method='Powell')
 
 
 
 
 
88
  W_opt = res.x.reshape(3, 3)
89
 
90
  corrected_lin = (tgt_lin_wb.reshape(-1, 3) @ W_opt).reshape(tgt_lin_wb.shape)
@@ -106,6 +117,14 @@ class ProcessResult:
106
  delta_e_final: Optional[float] = None
107
  image_base64: Optional[str] = None
108
  filename: Optional[str] = None
 
 
 
 
 
 
 
 
109
 
110
  class WatermelonProcessor:
111
  def __init__(self, model_path: str):
@@ -129,6 +148,106 @@ class WatermelonProcessor:
129
  divot_bot = d_bot * np.exp(w_bot * (-np.sin(t) - 1))
130
  return (ellipse * asymmetry) - divot_top - divot_bot + c_skew * np.sin(t) + c_bend * np.cos(t) * (np.sin(t) ** 2)
131
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
132
  @staticmethod
133
  def get_stable_perimeter_data(rind_mask, flesh_combined):
134
  cnts, _ = cv2.findContours(rind_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
@@ -192,7 +311,13 @@ class WatermelonProcessor:
192
  def parabola(y_n, a, b, c): return a*(y_n**2) + b*y_n + c
193
  max_bend = w * 0.08
194
  try:
195
- 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]))
 
 
 
 
 
 
196
  except: popt_mid = [0.0, 0.0, cx]
197
 
198
  ys_extrap = np.linspace(0, h, 500)
@@ -206,8 +331,30 @@ class WatermelonProcessor:
206
  pred_cnt_cv = pred_cnt.reshape(-1, 1, 2).astype(np.int32)
207
  return np.array([pt for pt in pts_orig if cv2.pointPolygonTest(pred_cnt_cv, (float(pt[0]), float(pt[1])), False) >= 0])
208
 
209
- def process_image(self, image: np.ndarray, source_name: str, scale_ratio: float) -> ProcessResult:
210
- if image is None: return ProcessResult(success=False, message="Could not decode image.")
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
211
  h, w = image.shape[:2]
212
 
213
  # --- 1. CALIBRATION & SCALING ---
@@ -222,135 +369,167 @@ class WatermelonProcessor:
222
  tgt_warped = warp_checker(image, checker_corners)
223
  tgt24 = sample_24_patches(tgt_warped)
224
  dE_initial = float(np.mean(compute_deltaE_00(tgt24, self.ref24)))
225
-
226
  image = apply_color_pipeline(image, self.ref24, tgt24)
227
-
228
  tgt_warped_corr = warp_checker(image, checker_corners)
229
  tgt24_corr = sample_24_patches(tgt_warped_corr)
230
  dE_final = float(np.mean(compute_deltaE_00(tgt24_corr, self.ref24)))
231
  except Exception as e:
 
 
 
 
232
  print(f"Calibration skipped for {source_name}: {e}")
233
-
 
234
  # --- 2. YOLO INFERENCE (STRICTLY PARSING ALL 3 CLASSES) ---
235
  results = self.model(image, conf=0.25, retina_masks=True, verbose=False)
 
 
236
  rind_mask = np.zeros((h, w), dtype=np.uint8)
237
- flesh_l_contours =[]
238
  flesh_r_contours = []
239
 
240
  if results[0].masks is None:
241
- return ProcessResult(success=False, message="No masks detected.")
 
 
 
 
 
242
 
243
  for mask_data, cls in zip(results[0].masks.xy, results[0].boxes.cls):
244
  contour = np.array(mask_data, dtype=np.int32)
245
  c_id = int(cls)
246
- if c_id == 0:
247
  cv2.drawContours(rind_mask, [contour], -1, 255, -1)
248
- elif c_id == 1:
249
  flesh_l_contours.append(contour)
250
- elif c_id == 2:
251
  flesh_r_contours.append(contour)
252
 
253
- # Failsafe: if YOLO missed one side but predicted multiple of the other
254
  if len(flesh_l_contours) >= 2 and len(flesh_r_contours) == 0:
255
- flesh_l_contours.sort(key=lambda cnt: cv2.moments(cnt)['m10'] / (cv2.moments(cnt)['m00'] + 1e-5))
256
  flesh_r_contours.append(flesh_l_contours.pop())
 
257
  elif len(flesh_r_contours) >= 2 and len(flesh_l_contours) == 0:
258
- flesh_r_contours.sort(key=lambda cnt: cv2.moments(cnt)['m10'] / (cv2.moments(cnt)['m00'] + 1e-5))
259
  flesh_l_contours.append(flesh_r_contours.pop(0))
 
260
 
261
  flesh_l_m, flesh_r_m = np.zeros((h, w), dtype=np.uint8), np.zeros((h, w), dtype=np.uint8)
262
- for cnt in flesh_l_contours: cv2.drawContours(flesh_l_m, [cnt], -1, 255, -1)
263
- for cnt in flesh_r_contours: cv2.drawContours(flesh_r_m, [cnt], -1, 255, -1)
 
 
264
 
265
  flesh_combined = cv2.bitwise_or(flesh_l_m, flesh_r_m)
 
 
 
266
 
267
  # --- 3. FIT & EXTRACTION ---
268
- perimeter_data = self.get_stable_perimeter_data(rind_mask, flesh_combined)
269
- if perimeter_data is None: return ProcessResult(success=False, message="No stable perimeter.")
 
 
 
 
 
 
 
 
270
 
271
  t_data, r_raw, (cx, cy), rind_cnt = perimeter_data
272
  scale = np.mean(r_raw)
273
- if scale <= 0: return ProcessResult(success=False, message="Invalid perimeter scale.")
274
-
 
275
  try:
276
  popt, _ = curve_fit(
277
  self.watermelon_model, t_data, r_raw / scale,
278
  p0=[1.0, 1.1, 0.0, 0.05, 3.0, 0.05, 3.0, 0.0, 0.0, 0.0],
279
- 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]),
 
280
  )
281
- except Exception as exc: return ProcessResult(success=False, message=f"Fit failed: {exc}")
 
 
282
 
283
  r2 = 1 - (np.sum((r_raw / scale - self.watermelon_model(t_data, *popt)) ** 2) / np.sum((r_raw / scale - 1) ** 2))
284
  t_fit = np.linspace(-np.pi, np.pi, 500)
285
  r_fit = self.watermelon_model(t_fit, *popt) * scale
286
  fit_pts = np.array([[r * np.cos(t) + cx, cy - r * np.sin(t)] for t, r in zip(t_fit, r_fit)])
287
-
288
- orig_scale = 1.0 / scale_ratio
289
- if cm_per_px is None: cm_per_px = 1.0
290
-
291
  width_px = float(np.max(fit_pts[:, 0]) - np.min(fit_pts[:, 0]))
292
  height_px = float(np.max(fit_pts[:, 1]) - np.min(fit_pts[:, 1]))
293
  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]))
294
 
295
  if cm_per_px is not None:
296
- # Both checker and watermelon were measured on the scaled image.
297
- # cm_per_px natively accounts for the scaling, no orig_scale needed!
298
  width_val = float(width_px * cm_per_px)
299
  height_val = float(height_px * cm_per_px)
300
  perimeter_val = float(perimeter_px * cm_per_px)
301
  else:
302
- # Fallback to true original pixels
 
303
  orig_scale = 1.0 / scale_ratio
304
  width_val = float(width_px * orig_scale)
305
  height_val = float(height_px * orig_scale)
306
  perimeter_val = float(perimeter_px * orig_scale)
 
307
 
308
  # --- 4. DRAWING ---
309
- midline = self.get_dual_mask_midline(flesh_l_m, flesh_r_m, rind_cnt, fit_pts, cx, cy)
310
-
311
- # Color coding: Green=Rind, Blue=Left Flesh, Red=Right Flesh
312
- output = image.copy().astype(np.float32)
313
- alpha = 0.42
314
- output[..., 0] = np.where(rind_mask > 0, output[..., 0] * (1 - alpha) + 0.0 * alpha, output[..., 0])
315
- output[..., 1] = np.where(rind_mask > 0, output[..., 1] * (1 - alpha) + 170.0 * alpha, output[..., 1])
316
- output[..., 2] = np.where(rind_mask > 0, output[..., 2] * (1 - alpha) + 0.0 * alpha, output[..., 2])
317
-
318
- output[..., 0] = np.where(flesh_l_m > 0, output[..., 0] * (1 - alpha) + 255.0 * alpha, output[..., 0])
319
- output[..., 1] = np.where(flesh_l_m > 0, output[..., 1] * (1 - alpha) + 0.0 * alpha, output[..., 1])
320
- output[..., 2] = np.where(flesh_l_m > 0, output[..., 2] * (1 - alpha) + 0.0 * alpha, output[..., 2])
321
-
322
- output[..., 0] = np.where(flesh_r_m > 0, output[..., 0] * (1 - alpha) + 0.0 * alpha, output[..., 0])
323
- output[..., 1] = np.where(flesh_r_m > 0, output[..., 1] * (1 - alpha) + 0.0 * alpha, output[..., 1])
324
- output[..., 2] = np.where(flesh_r_m > 0, output[..., 2] * (1 - alpha) + 255.0 * alpha, output[..., 2])
325
-
326
- output = np.clip(output, 0, 255).astype(np.uint8)
327
-
328
- # Draw Color Checker Box
329
- if checker_corners is not None:
330
- cv2.polylines(output, [np.int32(checker_corners)], True, (0, 165, 255), 4)
331
-
332
- # Draw Midline and Midline Intersections
333
- if len(midline) > 1:
334
- cv2.polylines(output, [midline.astype(np.int32)], False, (0, 255, 255), 3)
335
- # Add intersection dots (White inner, Black outer)
336
- pt_top = (int(midline[0][0]), int(midline[0][1]))
337
- pt_bot = (int(midline[-1][0]), int(midline[-1][1]))
338
- cv2.circle(output, pt_top, 10, (0, 0, 0), 2)
339
- cv2.circle(output, pt_top, 8, (255, 255, 255), -1)
340
- cv2.circle(output, pt_bot, 10, (0, 0, 0), 2)
341
- cv2.circle(output, pt_bot, 8, (255, 255, 255), -1)
342
-
343
- # Draw Predicted Perimeter Boundary
344
- cv2.polylines(output, [fit_pts.astype(np.int32)], True, (0, 255, 0), 3)
345
-
346
- _, buffer = cv2.imencode('.jpg', output,[cv2.IMWRITE_JPEG_QUALITY, 85])
347
- img_base64 = base64.b64encode(buffer).decode('utf-8')
348
 
349
  return ProcessResult(
350
  success=True, message="Success", r2_score=float(r2),
351
  width_val=width_val, height_val=height_val, perimeter_val=perimeter_val,
352
  delta_e_initial=dE_initial, delta_e_final=dE_final,
353
- image_base64=img_base64, filename=source_name
 
 
 
 
354
  )
355
 
356
 
@@ -366,19 +545,44 @@ processor = WatermelonProcessor(MODEL_PATH)
366
  def read_root(): return {"status": "Watermelon API is awake and running!"}
367
 
368
  @app.post("/process_single")
369
- async def process_single(file: UploadFile = File(...)):
370
- contents = await file.read()
371
- img = cv2.imdecode(np.frombuffer(contents, np.uint8), cv2.IMREAD_COLOR)
372
-
373
- scale_ratio = 1.0
374
- h, w = img.shape[:2]
375
- if max(h, w) > MAX_IMAGE_SIZE:
376
- scale_ratio = MAX_IMAGE_SIZE / float(max(h, w))
377
- img = cv2.resize(img, (int(w * scale_ratio), int(h * scale_ratio)), interpolation=cv2.INTER_AREA)
378
-
379
- res = processor.process_image(img, file.filename, scale_ratio)
380
-
381
- del img, contents
382
- gc.collect()
383
-
384
- return res.__dict__
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
4
  import base64
5
  import gc
6
  import torch
7
+ import time
8
+ import traceback
9
  from dataclasses import dataclass
10
  from typing import Optional
11
  from scipy.ndimage import median_filter
12
  from scipy.optimize import curve_fit, minimize
13
  from ultralytics import YOLO
14
+ from fastapi import FastAPI, UploadFile, File, Query
15
  from fastapi.middleware.cors import CORSMiddleware
16
  import uvicorn
17
  from skimage import color
 
23
  MODEL_PATH = "best.pt"
24
  MAX_IMAGE_SIZE = 2048
25
  CHECKER_WIDTH_CM = 6.3
26
+ MIN_RIND_FLESH_OVERLAP_RATIO = 0.10
27
+ MIN_FLESH_PIXELS_FOR_FALLBACK = 100
28
+ MAX_RIND_TO_FLESH_AREA_RATIO = 3.5
29
+ MAX_RIND_CENTER_OFFSET_RATIO = 0.60
30
 
31
  # ==============================================================================
32
  # --- COLOR CALIBRATION LOGIC ---
 
90
 
91
  X, Y = tgt24_wb, ref24
92
  W_init = np.linalg.inv(X.T @ X + 0.05 * np.eye(3)) @ X.T @ Y
93
+ res = minimize(
94
+ objective,
95
+ W_init.flatten(),
96
+ method='Powell',
97
+ options={"maxiter": 150, "xtol": 1e-4, "ftol": 1e-4},
98
+ )
99
  W_opt = res.x.reshape(3, 3)
100
 
101
  corrected_lin = (tgt_lin_wb.reshape(-1, 3) @ W_opt).reshape(tgt_lin_wb.shape)
 
117
  delta_e_final: Optional[float] = None
118
  image_base64: Optional[str] = None
119
  filename: Optional[str] = None
120
+ measurement_unit: Optional[str] = None
121
+ scale_source: Optional[str] = None
122
+ color_checker_found: bool = False
123
+ rind_source: Optional[str] = None
124
+ rind_overlap_ratio: Optional[float] = None
125
+ warnings: Optional[list] = None
126
+ timings_ms: Optional[dict] = None
127
+ processing_ms: Optional[int] = None
128
 
129
  class WatermelonProcessor:
130
  def __init__(self, model_path: str):
 
148
  divot_bot = d_bot * np.exp(w_bot * (-np.sin(t) - 1))
149
  return (ellipse * asymmetry) - divot_top - divot_bot + c_skew * np.sin(t) + c_bend * np.cos(t) * (np.sin(t) ** 2)
150
 
151
+ @staticmethod
152
+ def contour_centroid(contour):
153
+ M = cv2.moments(contour)
154
+ if M["m00"] == 0:
155
+ return None
156
+ return np.array([M["m10"] / M["m00"], M["m01"] / M["m00"]], dtype=np.float32)
157
+
158
+ @staticmethod
159
+ def mask_centroid(mask):
160
+ M = cv2.moments(mask)
161
+ if M["m00"] == 0:
162
+ return None
163
+ return np.array([M["m10"] / M["m00"], M["m01"] / M["m00"]], dtype=np.float32)
164
+
165
+ @staticmethod
166
+ def draw_single_contour(shape, contour):
167
+ mask = np.zeros(shape, dtype=np.uint8)
168
+ cv2.drawContours(mask, [contour.astype(np.int32)], -1, 255, -1)
169
+ return mask
170
+
171
+ @staticmethod
172
+ def flesh_envelope_mask(flesh_combined):
173
+ ys, xs = np.where(flesh_combined > 0)
174
+ if len(xs) < MIN_FLESH_PIXELS_FOR_FALLBACK:
175
+ return None
176
+
177
+ pts = np.column_stack([xs, ys]).astype(np.int32)
178
+ hull = cv2.convexHull(pts.reshape(-1, 1, 2))
179
+ envelope = np.zeros_like(flesh_combined)
180
+ cv2.drawContours(envelope, [hull], -1, 255, -1)
181
+
182
+ _, _, bw, bh = cv2.boundingRect(hull)
183
+ pad = int(max(12, min(80, round(max(bw, bh) * 0.035))))
184
+ kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (pad * 2 + 1, pad * 2 + 1))
185
+ envelope = cv2.dilate(envelope, kernel, iterations=1)
186
+
187
+ close_size = max(5, (pad // 2) * 2 + 1)
188
+ close_kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (close_size, close_size))
189
+ return cv2.morphologyEx(envelope, cv2.MORPH_CLOSE, close_kernel)
190
+
191
+ @staticmethod
192
+ def choose_target_rind_mask(rind_mask, flesh_combined):
193
+ warnings = []
194
+ flesh_area = cv2.countNonZero(flesh_combined)
195
+ cnts, _ = cv2.findContours(rind_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
196
+
197
+ if not cnts:
198
+ envelope = WatermelonProcessor.flesh_envelope_mask(flesh_combined)
199
+ if envelope is None:
200
+ return rind_mask, "missing", None, ["No whole-watermelon mask and not enough flesh mask for fallback."]
201
+ return envelope, "flesh_envelope", 1.0, ["No whole-watermelon mask; estimated perimeter from flesh masks."]
202
+
203
+ scored = []
204
+ for cnt in cnts:
205
+ temp = WatermelonProcessor.draw_single_contour(rind_mask.shape, cnt)
206
+ overlap = cv2.countNonZero(cv2.bitwise_and(temp, flesh_combined))
207
+ ratio = overlap / max(flesh_area, 1)
208
+ scored.append((ratio, cv2.contourArea(cnt), cnt, temp))
209
+
210
+ scored.sort(key=lambda item: (item[0], item[1]), reverse=True)
211
+ best_ratio, best_area, best_cnt, best_mask = scored[0]
212
+
213
+ if flesh_area == 0:
214
+ warnings.append("No flesh masks detected; using largest whole-watermelon mask.")
215
+ largest = max(cnts, key=cv2.contourArea)
216
+ return WatermelonProcessor.draw_single_contour(rind_mask.shape, largest), "whole_mask_no_flesh", None, warnings
217
+
218
+ flesh_center = WatermelonProcessor.mask_centroid(flesh_combined)
219
+ rind_center = WatermelonProcessor.contour_centroid(best_cnt)
220
+ ys, xs = np.where(flesh_combined > 0)
221
+ 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)
222
+ center_offset_ratio = 0.0
223
+ if flesh_center is not None and rind_center is not None:
224
+ center_offset_ratio = float(np.linalg.norm(flesh_center - rind_center) / flesh_extent)
225
+ area_ratio = float(best_area / max(flesh_area, 1))
226
+
227
+ if best_ratio >= MIN_RIND_FLESH_OVERLAP_RATIO:
228
+ if area_ratio <= MAX_RIND_TO_FLESH_AREA_RATIO and center_offset_ratio <= MAX_RIND_CENTER_OFFSET_RATIO:
229
+ return best_mask, "whole_mask_overlap", float(best_ratio), warnings
230
+ warnings.append("Whole-watermelon mask overlapped flesh but looked too large or off-center; using fallback.")
231
+
232
+ if flesh_center is not None and rind_center is not None:
233
+ shifted_cnt = best_cnt.astype(np.float32) + (flesh_center - rind_center).reshape(1, 1, 2)
234
+ shifted_mask = WatermelonProcessor.draw_single_contour(rind_mask.shape, shifted_cnt)
235
+ shifted_ratio = cv2.countNonZero(cv2.bitwise_and(shifted_mask, flesh_combined)) / max(flesh_area, 1)
236
+ if shifted_ratio >= MIN_RIND_FLESH_OVERLAP_RATIO and area_ratio <= MAX_RIND_TO_FLESH_AREA_RATIO:
237
+ if best_ratio >= MIN_RIND_FLESH_OVERLAP_RATIO:
238
+ warnings.append("Whole-watermelon mask was suspicious; translated it to the flesh-mask centroid.")
239
+ else:
240
+ warnings.append("Whole-watermelon mask did not overlap flesh; translated it to the flesh-mask centroid.")
241
+ return shifted_mask, "translated_whole_mask", float(shifted_ratio), warnings
242
+
243
+ envelope = WatermelonProcessor.flesh_envelope_mask(flesh_combined)
244
+ if envelope is not None:
245
+ warnings.append("Whole-watermelon mask did not overlap flesh; estimated perimeter from flesh masks.")
246
+ return envelope, "flesh_envelope", float(best_ratio), warnings
247
+
248
+ warnings.append("Whole-watermelon mask did not overlap flesh and fallback was unavailable.")
249
+ return best_mask, "low_overlap_whole_mask", float(best_ratio), warnings
250
+
251
  @staticmethod
252
  def get_stable_perimeter_data(rind_mask, flesh_combined):
253
  cnts, _ = cv2.findContours(rind_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
 
311
  def parabola(y_n, a, b, c): return a*(y_n**2) + b*y_n + c
312
  max_bend = w * 0.08
313
  try:
314
+ popt_mid, _ = curve_fit(
315
+ parabola,
316
+ (gap_points[:,0]-y_mean)/y_span,
317
+ gap_points[:,1],
318
+ bounds=([-max_bend, -np.inf, -np.inf],[max_bend, np.inf, np.inf]),
319
+ max_nfev=1500,
320
+ )
321
  except: popt_mid = [0.0, 0.0, cx]
322
 
323
  ys_extrap = np.linspace(0, h, 500)
 
331
  pred_cnt_cv = pred_cnt.reshape(-1, 1, 2).astype(np.int32)
332
  return np.array([pt for pt in pts_orig if cv2.pointPolygonTest(pred_cnt_cv, (float(pt[0]), float(pt[1])), False) >= 0])
333
 
334
+ def process_image(self, image: np.ndarray, source_name: str, scale_ratio: float, include_image: bool = True) -> ProcessResult:
335
+ timings = {}
336
+ stage_t = time.perf_counter()
337
+
338
+ def mark(stage_name):
339
+ nonlocal stage_t
340
+ now = time.perf_counter()
341
+ timings[stage_name] = int(round((now - stage_t) * 1000))
342
+ stage_t = now
343
+
344
+ def fail(message, **extra):
345
+ return ProcessResult(
346
+ success=False,
347
+ message=message,
348
+ filename=source_name,
349
+ warnings=warnings or None,
350
+ timings_ms=timings,
351
+ **extra,
352
+ )
353
+
354
+ warnings = []
355
+ if image is None:
356
+ return ProcessResult(success=False, message="Could not decode image.", filename=source_name)
357
+
358
  h, w = image.shape[:2]
359
 
360
  # --- 1. CALIBRATION & SCALING ---
 
369
  tgt_warped = warp_checker(image, checker_corners)
370
  tgt24 = sample_24_patches(tgt_warped)
371
  dE_initial = float(np.mean(compute_deltaE_00(tgt24, self.ref24)))
372
+
373
  image = apply_color_pipeline(image, self.ref24, tgt24)
374
+
375
  tgt_warped_corr = warp_checker(image, checker_corners)
376
  tgt24_corr = sample_24_patches(tgt_warped_corr)
377
  dE_final = float(np.mean(compute_deltaE_00(tgt24_corr, self.ref24)))
378
  except Exception as e:
379
+ if cm_per_px is None:
380
+ warnings.append("ColorChecker not found; dimensions are returned in original-image pixels.")
381
+ else:
382
+ warnings.append("Color correction skipped after ColorChecker detection; dimensions are still in centimeters.")
383
  print(f"Calibration skipped for {source_name}: {e}")
384
+ mark("calibration")
385
+
386
  # --- 2. YOLO INFERENCE (STRICTLY PARSING ALL 3 CLASSES) ---
387
  results = self.model(image, conf=0.25, retina_masks=True, verbose=False)
388
+ mark("yolo_inference")
389
+
390
  rind_mask = np.zeros((h, w), dtype=np.uint8)
391
+ flesh_l_contours = []
392
  flesh_r_contours = []
393
 
394
  if results[0].masks is None:
395
+ return fail(
396
+ "No masks detected.",
397
+ measurement_unit="cm" if cm_per_px is not None else "px",
398
+ scale_source="color_checker" if cm_per_px is not None else "original_pixels",
399
+ color_checker_found=checker_corners is not None,
400
+ )
401
 
402
  for mask_data, cls in zip(results[0].masks.xy, results[0].boxes.cls):
403
  contour = np.array(mask_data, dtype=np.int32)
404
  c_id = int(cls)
405
+ if c_id == 0:
406
  cv2.drawContours(rind_mask, [contour], -1, 255, -1)
407
+ elif c_id == 1:
408
  flesh_l_contours.append(contour)
409
+ elif c_id == 2:
410
  flesh_r_contours.append(contour)
411
 
412
+ # Failsafe: if YOLO missed one side but predicted multiple of the other.
413
  if len(flesh_l_contours) >= 2 and len(flesh_r_contours) == 0:
414
+ flesh_l_contours.sort(key=lambda cnt: cv2.moments(cnt)["m10"] / (cv2.moments(cnt)["m00"] + 1e-5))
415
  flesh_r_contours.append(flesh_l_contours.pop())
416
+ warnings.append("Only flesh_left was detected; split the two left detections into left/right by x-position.")
417
  elif len(flesh_r_contours) >= 2 and len(flesh_l_contours) == 0:
418
+ flesh_r_contours.sort(key=lambda cnt: cv2.moments(cnt)["m10"] / (cv2.moments(cnt)["m00"] + 1e-5))
419
  flesh_l_contours.append(flesh_r_contours.pop(0))
420
+ warnings.append("Only flesh_right was detected; split the two right detections into left/right by x-position.")
421
 
422
  flesh_l_m, flesh_r_m = np.zeros((h, w), dtype=np.uint8), np.zeros((h, w), dtype=np.uint8)
423
+ for cnt in flesh_l_contours:
424
+ cv2.drawContours(flesh_l_m, [cnt], -1, 255, -1)
425
+ for cnt in flesh_r_contours:
426
+ cv2.drawContours(flesh_r_m, [cnt], -1, 255, -1)
427
 
428
  flesh_combined = cv2.bitwise_or(flesh_l_m, flesh_r_m)
429
+ target_rind_mask, rind_source, rind_overlap_ratio, rind_warnings = self.choose_target_rind_mask(rind_mask, flesh_combined)
430
+ warnings.extend(rind_warnings)
431
+ mark("mask_parse")
432
 
433
  # --- 3. FIT & EXTRACTION ---
434
+ perimeter_data = self.get_stable_perimeter_data(target_rind_mask, flesh_combined)
435
+ if perimeter_data is None:
436
+ return fail(
437
+ "No stable perimeter.",
438
+ measurement_unit="cm" if cm_per_px is not None else "px",
439
+ scale_source="color_checker" if cm_per_px is not None else "original_pixels",
440
+ color_checker_found=checker_corners is not None,
441
+ rind_source=rind_source,
442
+ rind_overlap_ratio=rind_overlap_ratio,
443
+ )
444
 
445
  t_data, r_raw, (cx, cy), rind_cnt = perimeter_data
446
  scale = np.mean(r_raw)
447
+ if scale <= 0:
448
+ return fail("Invalid perimeter scale.", rind_source=rind_source, rind_overlap_ratio=rind_overlap_ratio)
449
+
450
  try:
451
  popt, _ = curve_fit(
452
  self.watermelon_model, t_data, r_raw / scale,
453
  p0=[1.0, 1.1, 0.0, 0.05, 3.0, 0.05, 3.0, 0.0, 0.0, 0.0],
454
+ 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]),
455
+ max_nfev=3000,
456
  )
457
+ except Exception as exc:
458
+ mark("fit")
459
+ return fail(f"Fit failed: {exc}", rind_source=rind_source, rind_overlap_ratio=rind_overlap_ratio)
460
 
461
  r2 = 1 - (np.sum((r_raw / scale - self.watermelon_model(t_data, *popt)) ** 2) / np.sum((r_raw / scale - 1) ** 2))
462
  t_fit = np.linspace(-np.pi, np.pi, 500)
463
  r_fit = self.watermelon_model(t_fit, *popt) * scale
464
  fit_pts = np.array([[r * np.cos(t) + cx, cy - r * np.sin(t)] for t, r in zip(t_fit, r_fit)])
465
+
 
 
 
466
  width_px = float(np.max(fit_pts[:, 0]) - np.min(fit_pts[:, 0]))
467
  height_px = float(np.max(fit_pts[:, 1]) - np.min(fit_pts[:, 1]))
468
  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]))
469
 
470
  if cm_per_px is not None:
471
+ measurement_unit = "cm"
472
+ scale_source = "color_checker"
473
  width_val = float(width_px * cm_per_px)
474
  height_val = float(height_px * cm_per_px)
475
  perimeter_val = float(perimeter_px * cm_per_px)
476
  else:
477
+ measurement_unit = "px"
478
+ scale_source = "original_pixels"
479
  orig_scale = 1.0 / scale_ratio
480
  width_val = float(width_px * orig_scale)
481
  height_val = float(height_px * orig_scale)
482
  perimeter_val = float(perimeter_px * orig_scale)
483
+ mark("fit")
484
 
485
  # --- 4. DRAWING ---
486
+ img_base64 = None
487
+ if include_image:
488
+ midline = self.get_dual_mask_midline(flesh_l_m, flesh_r_m, rind_cnt, fit_pts, cx, cy)
489
+
490
+ # Color coding: Green=chosen rind, Blue=Left Flesh, Red=Right Flesh.
491
+ output = image.copy().astype(np.float32)
492
+ alpha = 0.42
493
+ output[..., 0] = np.where(target_rind_mask > 0, output[..., 0] * (1 - alpha) + 0.0 * alpha, output[..., 0])
494
+ output[..., 1] = np.where(target_rind_mask > 0, output[..., 1] * (1 - alpha) + 170.0 * alpha, output[..., 1])
495
+ output[..., 2] = np.where(target_rind_mask > 0, output[..., 2] * (1 - alpha) + 0.0 * alpha, output[..., 2])
496
+
497
+ output[..., 0] = np.where(flesh_l_m > 0, output[..., 0] * (1 - alpha) + 255.0 * alpha, output[..., 0])
498
+ output[..., 1] = np.where(flesh_l_m > 0, output[..., 1] * (1 - alpha) + 0.0 * alpha, output[..., 1])
499
+ output[..., 2] = np.where(flesh_l_m > 0, output[..., 2] * (1 - alpha) + 0.0 * alpha, output[..., 2])
500
+
501
+ output[..., 0] = np.where(flesh_r_m > 0, output[..., 0] * (1 - alpha) + 0.0 * alpha, output[..., 0])
502
+ output[..., 1] = np.where(flesh_r_m > 0, output[..., 1] * (1 - alpha) + 0.0 * alpha, output[..., 1])
503
+ output[..., 2] = np.where(flesh_r_m > 0, output[..., 2] * (1 - alpha) + 255.0 * alpha, output[..., 2])
504
+
505
+ output = np.clip(output, 0, 255).astype(np.uint8)
506
+
507
+ if checker_corners is not None:
508
+ cv2.polylines(output, [np.int32(checker_corners)], True, (0, 165, 255), 4)
509
+
510
+ if len(midline) > 1:
511
+ cv2.polylines(output, [midline.astype(np.int32)], False, (0, 255, 255), 3)
512
+ pt_top = (int(midline[0][0]), int(midline[0][1]))
513
+ pt_bot = (int(midline[-1][0]), int(midline[-1][1]))
514
+ cv2.circle(output, pt_top, 10, (0, 0, 0), 2)
515
+ cv2.circle(output, pt_top, 8, (255, 255, 255), -1)
516
+ cv2.circle(output, pt_bot, 10, (0, 0, 0), 2)
517
+ cv2.circle(output, pt_bot, 8, (255, 255, 255), -1)
518
+
519
+ cv2.polylines(output, [fit_pts.astype(np.int32)], True, (0, 255, 0), 3)
520
+ _, buffer = cv2.imencode(".jpg", output, [cv2.IMWRITE_JPEG_QUALITY, 85])
521
+ img_base64 = base64.b64encode(buffer).decode("utf-8")
522
+ mark("render")
 
 
523
 
524
  return ProcessResult(
525
  success=True, message="Success", r2_score=float(r2),
526
  width_val=width_val, height_val=height_val, perimeter_val=perimeter_val,
527
  delta_e_initial=dE_initial, delta_e_final=dE_final,
528
+ image_base64=img_base64, filename=source_name,
529
+ measurement_unit=measurement_unit, scale_source=scale_source,
530
+ color_checker_found=checker_corners is not None,
531
+ rind_source=rind_source, rind_overlap_ratio=rind_overlap_ratio,
532
+ warnings=warnings or None, timings_ms=timings
533
  )
534
 
535
 
 
545
  def read_root(): return {"status": "Watermelon API is awake and running!"}
546
 
547
  @app.post("/process_single")
548
+ async def process_single(file: UploadFile = File(...), include_image: bool = Query(True)):
549
+ request_t = time.perf_counter()
550
+ contents = None
551
+ img = None
552
+
553
+ try:
554
+ contents = await file.read()
555
+ if not contents:
556
+ res = ProcessResult(success=False, message="Empty upload.", filename=file.filename)
557
+ res.processing_ms = int(round((time.perf_counter() - request_t) * 1000))
558
+ return res.__dict__
559
+
560
+ img = cv2.imdecode(np.frombuffer(contents, np.uint8), cv2.IMREAD_COLOR)
561
+ if img is None:
562
+ res = ProcessResult(success=False, message="Could not decode image.", filename=file.filename)
563
+ res.processing_ms = int(round((time.perf_counter() - request_t) * 1000))
564
+ return res.__dict__
565
+
566
+ scale_ratio = 1.0
567
+ h, w = img.shape[:2]
568
+ if max(h, w) > MAX_IMAGE_SIZE:
569
+ scale_ratio = MAX_IMAGE_SIZE / float(max(h, w))
570
+ img = cv2.resize(img, (int(w * scale_ratio), int(h * scale_ratio)), interpolation=cv2.INTER_AREA)
571
+
572
+ res = processor.process_image(img, file.filename, scale_ratio, include_image=include_image)
573
+ res.processing_ms = int(round((time.perf_counter() - request_t) * 1000))
574
+ return res.__dict__
575
+
576
+ except Exception as exc:
577
+ traceback.print_exc()
578
+ res = ProcessResult(
579
+ success=False,
580
+ message=f"Server error: {type(exc).__name__}: {exc}",
581
+ filename=file.filename,
582
+ processing_ms=int(round((time.perf_counter() - request_t) * 1000)),
583
+ )
584
+ return res.__dict__
585
+
586
+ finally:
587
+ del img, contents
588
+ gc.collect()