crabbly commited on
Commit
ced90d4
·
1 Parent(s): ba914d4

Deploying FastAPI Backend Code

Browse files
Files changed (4) hide show
  1. Dockerfile +17 -0
  2. cv_helpers.py +78 -0
  3. main.py +240 -0
  4. requirements.txt +10 -0
Dockerfile ADDED
@@ -0,0 +1,17 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Use official Python image
2
+ FROM python:3.11-slim
3
+
4
+ # Set the working directory
5
+ WORKDIR /app
6
+
7
+ # Copy your requirements
8
+ COPY requirements.txt .
9
+
10
+ # Install dependencies (CPU-only PyTorch to save space)
11
+ RUN pip install --no-cache-dir -r requirements.txt
12
+
13
+ # Copy all your scripts and weights into the container
14
+ COPY . .
15
+
16
+ # Hugging Face requires Docker spaces to run on port 7860!
17
+ CMD["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "7860"]
cv_helpers.py ADDED
@@ -0,0 +1,78 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Shared CV helpers: mask visualization and stem-tip heuristic."""
2
+
3
+ import numpy as np
4
+
5
+
6
+ def blend_mask_overlays(bgr, rind_mask, flesh_mask, alpha=0.42):
7
+ """
8
+ Semi-transparent rind (green tint) and flesh (orange tint) on top of the BGR image.
9
+ Flesh is drawn after rind so overlap reads clearly.
10
+ """
11
+ out = bgr.astype(np.float32)
12
+ rind_m = (rind_mask > 0).astype(np.float32)
13
+ flesh_m = (flesh_mask > 0).astype(np.float32)
14
+ rind_color = np.array([0.0, 170.0, 0.0], dtype=np.float32)
15
+ flesh_color = np.array([60.0, 120.0, 255.0], dtype=np.float32)
16
+ for c in range(3):
17
+ ch = out[..., c]
18
+ ch[:] = ch * (1.0 - alpha * rind_m) + rind_color[c] * (alpha * rind_m)
19
+ for c in range(3):
20
+ ch = out[..., c]
21
+ ch[:] = ch * (1.0 - alpha * flesh_m) + flesh_color[c] * (alpha * flesh_m)
22
+ return np.clip(out, 0, 255).astype(np.uint8)
23
+
24
+
25
+ def stem_tip_tangent_deg(contour, centroid_xy):
26
+ """
27
+ Heuristic "stem / neck" pole on the rind contour: take PCA major-axis extremes,
28
+ then pick the end with sharper local turning (inward-curving neck). Tie-break:
29
+ smaller image y (overhead shots often have stem toward top of frame).
30
+
31
+ Returns (tip_x, tip_y, tangent_deg) where tangent_deg is atan2(dy, dx) in degrees,
32
+ or None if not enough contour points.
33
+ """
34
+ cnt = contour.reshape(-1, 2).astype(np.float64)
35
+ n = len(cnt)
36
+ if n < 9:
37
+ return None
38
+
39
+ cx, cy = float(centroid_xy[0]), float(centroid_xy[1])
40
+ X = cnt - np.array([cx, cy])
41
+ cov = np.cov(X.T)
42
+ eigvals, eigvecs = np.linalg.eigh(cov)
43
+ u = eigvecs[:, int(np.argmax(eigvals))]
44
+ un = np.linalg.norm(u)
45
+ if un < 1e-9:
46
+ return None
47
+ u /= un
48
+
49
+ s = X @ u
50
+ idx_a = int(np.argmax(s))
51
+ idx_b = int(np.argmin(s))
52
+ span = max(3, min(25, n // 30))
53
+
54
+ def curvature_score(i):
55
+ p = cnt[i % n]
56
+ prev = cnt[(i - span) % n]
57
+ nxt = cnt[(i + span) % n]
58
+ v1 = p - prev
59
+ v2 = nxt - p
60
+ nv1 = np.linalg.norm(v1)
61
+ nv2 = np.linalg.norm(v2)
62
+ if nv1 < 1e-6 or nv2 < 1e-6:
63
+ return 0.0
64
+ v1u = v1 / nv1
65
+ v2u = v2 / nv2
66
+ return abs(v1u[0] * v2u[1] - v1u[1] * v2u[0])
67
+
68
+ ka, kb = curvature_score(idx_a), curvature_score(idx_b)
69
+ if abs(ka - kb) < 0.05:
70
+ stem_idx = idx_a if cnt[idx_a, 1] < cnt[idx_b, 1] else idx_b
71
+ else:
72
+ stem_idx = idx_a if ka > kb else idx_b
73
+
74
+ span_t = max(2, span // 2)
75
+ d = cnt[(stem_idx + span_t) % n] - cnt[(stem_idx - span_t) % n]
76
+ tang_deg = float(np.degrees(np.arctan2(d[1], d[0])))
77
+ tip = cnt[stem_idx]
78
+ return float(tip[0]), float(tip[1]), tang_deg
main.py ADDED
@@ -0,0 +1,240 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import cv2
3
+ 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
11
+ from ultralytics import YOLO
12
+ from fastapi import FastAPI, UploadFile, File
13
+ from fastapi.middleware.cors import CORSMiddleware
14
+ import uvicorn
15
+
16
+ # OOM PREVENTION 1: Force PyTorch to use minimal memory overhead
17
+ torch.set_num_threads(1)
18
+
19
+ # Import your helpers (assuming cv_helpers.py is in the same folder)
20
+ from cv_helpers import blend_mask_overlays, stem_tip_tangent_deg
21
+
22
+ # --- CONFIGURATION ---
23
+ MODEL_PATH = "best.pt"
24
+ PIXELS_TO_CM = 1.0
25
+ MAX_IMAGE_SIZE = 1024 # OOM PREVENTION 2: Max pixels on the longest side
26
+
27
+ @dataclass
28
+ class ProcessResult:
29
+ success: bool
30
+ message: str
31
+ r2_score: Optional[float] = None
32
+ width_val: Optional[float] = None
33
+ height_val: Optional[float] = None
34
+ perimeter_val: Optional[float] = None
35
+ image_base64: Optional[str] = None
36
+ filename: Optional[str] = None
37
+
38
+ class WatermelonProcessor:
39
+ def __init__(self, model_path: str):
40
+ self.model = YOLO(model_path)
41
+
42
+ @staticmethod
43
+ def watermelon_model(theta, Rx, Ry, c_a, d_top, w_top, d_bot, w_bot, phi, c_skew, c_bend):
44
+ t = theta - phi
45
+ ellipse = (Rx * Ry) / np.sqrt((Ry * np.cos(t)) ** 2 + (Rx * np.sin(t)) ** 2)
46
+ asymmetry = 1 + c_a * np.cos(t) ** 3
47
+ divot_top = d_top * np.exp(w_top * (np.sin(t) - 1))
48
+ divot_bot = d_bot * np.exp(w_bot * (-np.sin(t) - 1))
49
+ return (ellipse * asymmetry) - divot_top - divot_bot + c_skew * np.sin(t) + c_bend * np.cos(t) * (np.sin(t) ** 2)
50
+
51
+ @staticmethod
52
+ def get_stable_perimeter_data(rind_mask, flesh_mask):
53
+ cnts, _ = cv2.findContours(rind_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
54
+ if not cnts: return None
55
+ best_cnt = None
56
+ max_overlap = -1
57
+ for cnt in cnts:
58
+ temp_mask = np.zeros_like(rind_mask)
59
+ cv2.drawContours(temp_mask, [cnt], -1, 255, -1)
60
+ overlap_area = cv2.countNonZero(cv2.bitwise_and(temp_mask, flesh_mask))
61
+ if overlap_area > max_overlap:
62
+ max_overlap = overlap_area
63
+ best_cnt = cnt
64
+ if best_cnt is None: best_cnt = max(cnts, key=cv2.contourArea)
65
+ moments = cv2.moments(best_cnt)
66
+ if moments["m00"] == 0: return None
67
+
68
+ cx, cy = moments["m10"] / moments["m00"], moments["m01"] / moments["m00"]
69
+ pts = best_cnt.reshape(-1, 2)
70
+ dx, dy = pts[:, 0] - cx, cy - pts[:, 1]
71
+ r_vals, t_vals = np.sqrt(dx**2 + dy**2), np.arctan2(dy, dx)
72
+ num_bins = 360
73
+ bins = np.linspace(-np.pi, np.pi, num_bins + 1)
74
+ raw_r = np.full(num_bins, np.nan)
75
+ for i in range(num_bins):
76
+ mask = (t_vals >= bins[i]) & (t_vals < bins[i + 1])
77
+ if np.any(mask): raw_r[i] = np.max(r_vals[mask])
78
+
79
+ valid_idx = np.where(~np.isnan(raw_r))[0]
80
+ if len(valid_idx) == 0: return None
81
+ raw_r[np.isnan(raw_r)] = np.interp(np.where(np.isnan(raw_r))[0], valid_idx, raw_r[valid_idx], period=360)
82
+ final_r = median_filter(raw_r, size=7, mode="wrap")
83
+ final_theta = (bins[:-1] + bins[1:]) / 2.0
84
+ return final_theta, final_r, (cx, cy), best_cnt
85
+
86
+ @staticmethod
87
+ def get_ray_scan_midline(flesh_mask, rind_cnt, predicted_cnt, cx, cy):
88
+ h, w = flesh_mask.shape
89
+ if len(rind_cnt) > 5:
90
+ _, (ma, Ma), angle = cv2.fitEllipse(rind_cnt)
91
+ rot_angle = angle if ma < Ma else angle + 90
92
+ else: rot_angle = 0
93
+
94
+ m_rot = cv2.getRotationMatrix2D((cx, cy), rot_angle, 1.0)
95
+ m_inv = cv2.getRotationMatrix2D((cx, cy), -rot_angle, 1.0)
96
+ f_rot = cv2.warpAffine(flesh_mask, m_rot, (w, h))
97
+
98
+ gap_points =[]
99
+ y_indices, _ = np.where(f_rot > 0)
100
+ if len(y_indices) > 0:
101
+ for y in range(np.min(y_indices), np.max(y_indices)):
102
+ row = f_rot[y, :]
103
+ white_px = np.where(row > 0)[0]
104
+ if len(white_px) >= 2:
105
+ blanks = np.where(row[white_px[0]:white_px[-1]] == 0)[0] + white_px[0]
106
+ if len(blanks) > 0: gap_points.append([y, np.median(blanks)])
107
+
108
+ gap_points = np.array(gap_points)
109
+ if len(gap_points) > 10:
110
+ y_min, y_max = np.min(gap_points[:, 0]), np.max(gap_points[:, 0])
111
+ y_span, y_mean = max(y_max - y_min, 1), (y_max + y_min) / 2.0
112
+ y_norm = (gap_points[:, 0] - y_mean) / y_span
113
+ x_data = gap_points[:, 1]
114
+
115
+ def parabola(y_n, a, b, c): return a * (y_n**2) + b * y_n + c
116
+ try:
117
+ popt_mid, _ = curve_fit(parabola, y_norm, x_data, bounds=([-w*0.08, -np.inf, -np.inf],[w*0.08, np.inf, np.inf]))
118
+ except: popt_mid =[0.0, 0.0, cx]
119
+
120
+ ys_extrap = np.linspace(0, h, 500)
121
+ xs_extrap = parabola((ys_extrap - y_mean) / y_span, *popt_mid)
122
+ pts_rot = np.vstack([xs_extrap, ys_extrap, np.ones_like(xs_extrap)])
123
+ else:
124
+ ys_extrap = np.linspace(0, h, 500)
125
+ pts_rot = np.vstack([np.full_like(ys_extrap, cx), ys_extrap, np.ones_like(ys_extrap)])
126
+
127
+ pts_orig = (m_inv @ pts_rot).T
128
+ pred_cnt_cv = predicted_cnt.reshape(-1, 1, 2).astype(np.int32)
129
+ return np.array([pt for pt in pts_orig if cv2.pointPolygonTest(pred_cnt_cv, (float(pt[0]), float(pt[1])), False) >= 0])
130
+
131
+ def process_image(self, image: np.ndarray, source_name: str, scale_ratio: float) -> ProcessResult:
132
+ if image is None: return ProcessResult(success=False, message="Could not decode image.")
133
+ h, w = image.shape[:2]
134
+
135
+ results = self.model(image, conf=0.25, verbose=False)
136
+ rind_mask, flesh_mask = np.zeros((h, w), dtype=np.uint8), np.zeros((h, w), dtype=np.uint8)
137
+
138
+ if results[0].masks is None:
139
+ return ProcessResult(success=False, message="No masks detected.")
140
+
141
+ for mask_data, cls in zip(results[0].masks.xy, results[0].boxes.cls):
142
+ contour = np.array(mask_data, dtype=np.int32)
143
+ if int(cls) == 0: cv2.drawContours(rind_mask, [contour], -1, 255, -1)
144
+ elif int(cls) == 1: cv2.drawContours(flesh_mask, [contour], -1, 255, -1)
145
+
146
+ perimeter_data = self.get_stable_perimeter_data(rind_mask, flesh_mask)
147
+ if perimeter_data is None: return ProcessResult(success=False, message="No stable perimeter.")
148
+
149
+ t_data, r_raw, (cx, cy), rind_cnt = perimeter_data
150
+ scale = np.mean(r_raw)
151
+
152
+ try:
153
+ popt, _ = curve_fit(
154
+ self.watermelon_model, t_data, r_raw / scale,
155
+ p0=[1.0, 1.1, 0.0, 0.05, 3.0, 0.05, 3.0, 0.0, 0.0, 0.0],
156
+ 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]),
157
+ )
158
+ except Exception as exc: return ProcessResult(success=False, message=f"Fit failed: {exc}")
159
+
160
+ r2 = 1 - (np.sum((r_raw / scale - self.watermelon_model(t_data, *popt)) ** 2) / np.sum((r_raw / scale - 1) ** 2))
161
+
162
+ t_fit = np.linspace(-np.pi, np.pi, 500)
163
+ r_fit = self.watermelon_model(t_fit, *popt) * scale
164
+ fit_pts = np.array([[r * np.cos(t) + cx, cy - r * np.sin(t)] for t, r in zip(t_fit, r_fit)])
165
+
166
+ # --- FEATURE EXTRACTION (WITH TRUE-SIZE CORRECTION) ---
167
+ width_px = float(np.max(fit_pts[:, 0]) - np.min(fit_pts[:, 0]))
168
+ height_px = float(np.max(fit_pts[:, 1]) - np.min(fit_pts[:, 1]))
169
+ diffs = np.diff(fit_pts, axis=0)
170
+ perimeter_px = float(np.sum(np.linalg.norm(diffs, axis=1)) + np.linalg.norm(fit_pts[-1] - fit_pts[0]))
171
+
172
+ # We divide by scale_ratio to perfectly undo the downscaling for measurements!
173
+ orig_scale = 1.0 / scale_ratio
174
+ width_val = width_px * PIXELS_TO_CM * orig_scale
175
+ height_val = height_px * PIXELS_TO_CM * orig_scale
176
+ perimeter_val = perimeter_px * PIXELS_TO_CM * orig_scale
177
+
178
+ # --- DRAWING ---
179
+ midline = self.get_ray_scan_midline(flesh_mask, rind_cnt, fit_pts, cx, cy)
180
+ output = blend_mask_overlays(image, rind_mask, flesh_mask)
181
+ if len(midline) > 1: cv2.polylines(output,[midline.astype(np.int32)], False, (0, 255, 255), 3)
182
+ cv2.polylines(output,[fit_pts.astype(np.int32)], True, (0, 255, 0), 3)
183
+
184
+ stem = stem_tip_tangent_deg(rind_cnt, (cx, cy))
185
+ if stem is not None:
186
+ tx, ty, tdeg = stem
187
+ L = min(w, h) * 0.08
188
+ rad = np.deg2rad(tdeg)
189
+ p1 = (int(round(tx)), int(round(ty)))
190
+ p2 = (int(round(tx + L * np.cos(rad))), int(round(ty + L * np.sin(rad))))
191
+ cv2.circle(output, p1, 6, (255, 0, 255), -1)
192
+ cv2.line(output, p1, p2, (255, 0, 255), 2)
193
+
194
+ _, buffer = cv2.imencode('.jpg', output, [cv2.IMWRITE_JPEG_QUALITY, 85])
195
+ img_base64 = base64.b64encode(buffer).decode('utf-8')
196
+
197
+ return ProcessResult(
198
+ success=True, message="Success", r2_score=float(r2),
199
+ width_val=width_val, height_val=height_val, perimeter_val=perimeter_val,
200
+ image_base64=img_base64, filename=source_name
201
+ )
202
+
203
+
204
+ # --- FASTAPI APP ---
205
+ app = FastAPI()
206
+
207
+ app.add_middleware(
208
+ CORSMiddleware,
209
+ allow_origins=["*"],
210
+ allow_credentials=True,
211
+ allow_methods=["*"],
212
+ allow_headers=["*"],
213
+ )
214
+
215
+ processor = WatermelonProcessor(MODEL_PATH)
216
+
217
+ @app.get("/")
218
+ def read_root():
219
+ return {"status": "Watermelon API is awake and running!"}
220
+
221
+ @app.post("/process_single")
222
+ async def process_single(file: UploadFile = File(...)):
223
+ contents = await file.read()
224
+ nparr = np.frombuffer(contents, np.uint8)
225
+ img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
226
+
227
+ # OOM PREVENTION 3: Resize image if it's massive
228
+ h, w = img.shape[:2]
229
+ scale_ratio = 1.0
230
+ if max(h, w) > MAX_IMAGE_SIZE:
231
+ scale_ratio = MAX_IMAGE_SIZE / float(max(h, w))
232
+ img = cv2.resize(img, (int(w * scale_ratio), int(h * scale_ratio)), interpolation=cv2.INTER_AREA)
233
+
234
+ res = processor.process_image(img, file.filename, scale_ratio)
235
+
236
+ # OOM PREVENTION 4: Force garbage collection immediately after processing
237
+ del img, nparr, contents
238
+ gc.collect()
239
+
240
+ return res.__dict__
requirements.txt ADDED
@@ -0,0 +1,10 @@
 
 
 
 
 
 
 
 
 
 
 
1
+ fastapi
2
+ uvicorn[standard]
3
+ python-multipart
4
+ numpy
5
+ scipy
6
+ ultralytics
7
+ opencv-python-headless
8
+ --extra-index-url https://download.pytorch.org/whl/cpu
9
+ torch
10
+ torchvision