LiangLabUMB commited on
Commit
83b2285
Β·
verified Β·
1 Parent(s): 89b513b

Upload app.py

Browse files
Files changed (1) hide show
  1. app.py +1783 -0
app.py ADDED
@@ -0,0 +1,1783 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import spaces
3
+ from cellpose import models
4
+ import numpy as np
5
+ import cv2
6
+ import matplotlib.pyplot as plt
7
+ import tempfile
8
+ from PIL import Image, ImageDraw
9
+ import io
10
+ from huggingface_hub import hf_hub_download
11
+ import base64
12
+ from concurrent.futures import ThreadPoolExecutor, as_completed
13
+ import csv
14
+ import joblib
15
+ import os
16
+
17
+ HF_REPO_ID = "myang4218/cellposemodel"
18
+ HF_REPO_ID2 = "LiangLabUMB/viability_model"
19
+ MODEL_OPTIONS = {
20
+ "Hemocytometer Model": "hemocytometermodel.npy",
21
+ "General Model": "generalmodel.npy"
22
+ }
23
+
24
+ loaded_models = {}
25
+
26
+ VIABILITY_CLF = None
27
+ VIABILITY_SCALER = None
28
+
29
+ try:
30
+ _clf_path = hf_hub_download(repo_id=HF_REPO_ID2, filename="viability_clf.pkl")
31
+ _scaler_path = hf_hub_download(repo_id=HF_REPO_ID2, filename="viability_scaler.pkl")
32
+ VIABILITY_CLF = joblib.load(_clf_path)
33
+ VIABILITY_SCALER = joblib.load(_scaler_path)
34
+ print("βœ“ Viability classifier loaded.")
35
+ except Exception as e:
36
+ print(f"Viability classifier not found or failed to load: {e}")
37
+
38
+ # ---- mobile-safe size limits (aggressive for Safari) ----
39
+ MAX_SIDE = 1024
40
+ MAX_PIXELS = 1024 * 1024
41
+
42
+
43
+ def safe_resize(image_np):
44
+ """
45
+ Downscale image to fit within MAX_SIDE and MAX_PIXELS while
46
+ preserving aspect ratio. Works for RGB / RGBA / grayscale.
47
+ """
48
+ h, w = image_np.shape[:2]
49
+ total = h * w
50
+
51
+ if max(h, w) <= MAX_SIDE and total <= MAX_PIXELS:
52
+ return image_np
53
+
54
+ # compute scale
55
+ scale_side = MAX_SIDE / max(h, w)
56
+ scale_pixels = (MAX_PIXELS / total) ** 0.5
57
+ scale = min(scale_side, scale_pixels)
58
+
59
+ new_w = max(1, int(w * scale))
60
+ new_h = max(1, int(h * scale))
61
+
62
+ return cv2.resize(image_np, (new_w, new_h), interpolation=cv2.INTER_AREA)
63
+
64
+
65
+ def draw_exclusion_overlay(image_np, left_width_pct, top_width_pct):
66
+
67
+ h, w = image_np.shape[:2]
68
+
69
+ # Convert to PIL for drawing
70
+ img_pil = Image.fromarray(image_np)
71
+ draw = ImageDraw.Draw(img_pil, 'RGBA')
72
+
73
+ # Calculate pixel widths from percentages
74
+ left_px = int(w * left_width_pct / 100)
75
+ top_px = int(h * top_width_pct / 100)
76
+
77
+ # Draw overlays for exclusion zones
78
+ if left_px > 0:
79
+ # Left exclusion zone
80
+ draw.rectangle(
81
+ [(0, 0), (left_px, h)],
82
+ fill=(255, 0, 0, 80) # Semi-transparent red
83
+ )
84
+ # border line
85
+ draw.line([(left_px, 0), (left_px, h)], fill=(255, 0, 0, 255), width=3)
86
+
87
+ if top_px > 0:
88
+ # Top exclusion zone
89
+ draw.rectangle(
90
+ [(0, 0), (w, top_px)],
91
+ fill=(255, 0, 0, 80) # Semi-transparent red
92
+ )
93
+ # border line
94
+ draw.line([(0, top_px), (w, top_px)], fill=(255, 0, 0, 255), width=3)
95
+
96
+ return np.array(img_pil)
97
+
98
+
99
+ def apply_stereological_exclusion(masks, left_width_pct, top_width_pct):
100
+ """
101
+ Exclude every cell touching the left or top exclusion zone.
102
+
103
+ A cell intersects the half-plane x < left_px exactly when its leftmost
104
+ pixel does, so per-label bounding boxes give an exact answer -- no need to
105
+ approximate with centroids and radii. scipy's find_objects collects every
106
+ bounding box in a single pass, instead of one full-image comparison per
107
+ cell, which is what makes this cheap enough to re-run interactively.
108
+
109
+ Cell ids are NOT renumbered: stable ids let the exclusion be re-applied
110
+ after viability classification without invalidating its label map.
111
+ """
112
+ from scipy import ndimage
113
+
114
+ h, w = masks.shape
115
+ left_px = int(w * left_width_pct / 100)
116
+ top_px = int(h * top_width_pct / 100)
117
+
118
+ if left_px <= 0 and top_px <= 0:
119
+ n = len(np.unique(masks)) - (1 if (masks == 0).any() else 0)
120
+ return masks.copy(), 0, n
121
+
122
+ boxes = ndimage.find_objects(masks)
123
+ excluded_ids = []
124
+ included_ids = []
125
+ for idx, box in enumerate(boxes):
126
+ if box is None: # id absent from the label image
127
+ continue
128
+ cell_id = idx + 1
129
+ row_slice, col_slice = box
130
+ touches_left = left_px > 0 and col_slice.start < left_px
131
+ touches_top = top_px > 0 and row_slice.start < top_px
132
+ if touches_left or touches_top:
133
+ excluded_ids.append(cell_id)
134
+ else:
135
+ included_ids.append(cell_id)
136
+
137
+ filtered_masks = masks.copy()
138
+ if excluded_ids:
139
+ drop = np.zeros(int(masks.max()) + 1, dtype=bool)
140
+ drop[np.asarray(excluded_ids, dtype=np.int64)] = True
141
+ filtered_masks[drop[masks]] = 0
142
+
143
+ return filtered_masks, len(excluded_ids), len(included_ids)
144
+
145
+
146
+ def draw_exclusion_overlay(image_np, left_width_pct, top_width_pct):
147
+
148
+ h, w = image_np.shape[:2]
149
+
150
+ # Convert to PIL for drawing
151
+ img_pil = Image.fromarray(image_np)
152
+ draw = ImageDraw.Draw(img_pil, 'RGBA')
153
+
154
+ # Calculate pixel widths from percentages
155
+ left_px = int(w * left_width_pct / 100)
156
+ top_px = int(h * top_width_pct / 100)
157
+
158
+ # Draw overlays for exclusion zones
159
+ if left_px > 0:
160
+ # Left exclusion zone
161
+ draw.rectangle(
162
+ [(0, 0), (left_px, h)],
163
+ fill=(255, 0, 0, 80) # Semi-transparent red
164
+ )
165
+ # border line
166
+ draw.line([(left_px, 0), (left_px, h)], fill=(255, 0, 0, 255), width=3)
167
+
168
+ if top_px > 0:
169
+ # Top exclusion zone
170
+ draw.rectangle(
171
+ [(0, 0), (w, top_px)],
172
+ fill=(255, 0, 0, 80) # Semi-transparent red
173
+ )
174
+ # border line
175
+ draw.line([(0, top_px), (w, top_px)], fill=(255, 0, 0, 255), width=3)
176
+
177
+ return np.array(img_pil)
178
+
179
+
180
+ def apply_stereological_exclusion(masks, left_width_pct, top_width_pct):
181
+ h, w = masks.shape
182
+
183
+ # Calculate pixel widths from percentages
184
+ left_px = int(w * left_width_pct / 100)
185
+ top_px = int(h * top_width_pct / 100)
186
+
187
+ filtered_masks = masks.copy()
188
+ cell_ids = np.unique(masks)
189
+ cell_ids = cell_ids[cell_ids > 0]
190
+
191
+ excluded_cells = []
192
+ included_cells = []
193
+
194
+ for cell_id in cell_ids:
195
+ cell_mask = (masks == cell_id)
196
+
197
+ # Get cell boundary coordinates
198
+ rows, cols = np.where(cell_mask)
199
+
200
+ # Check if cell touches left exclusion zone
201
+ touches_left = np.any(cols < left_px) if left_px > 0 else False
202
+
203
+ # Check if cell touches top exclusion zone
204
+ touches_top = np.any(rows < top_px) if top_px > 0 else False
205
+
206
+ # Exclude if touching left or top
207
+ if touches_left or touches_top:
208
+ filtered_masks[cell_mask] = 0
209
+ excluded_cells.append(cell_id)
210
+ else:
211
+ included_cells.append(cell_id)
212
+
213
+ # Renumber remaining cells
214
+ unique_ids = np.unique(filtered_masks)
215
+ unique_ids = unique_ids[unique_ids > 0]
216
+
217
+ renumbered_masks = np.zeros_like(filtered_masks)
218
+ for new_id, old_id in enumerate(unique_ids, start=1):
219
+ renumbered_masks[filtered_masks == old_id] = new_id
220
+
221
+ return renumbered_masks, len(excluded_cells), len(included_cells)
222
+
223
+
224
+
225
+ FEATURE_COLS_INFERENCE = [
226
+ "mean_r", "mean_g", "mean_b", "std_r", "std_g", "std_b",
227
+ "mean_h", "mean_s", "mean_v", "std_s", "std_v",
228
+ "blue_red_ratio", "blue_green_ratio", "rg_ratio",
229
+ "inner_brightness", "peak_brightness",
230
+ "bright_spot_fraction", "ring_darkness",
231
+ "centre_periphery_ratio", "brightness_std_normalised",
232
+ ]
233
+
234
+
235
+ def classify_cells_by_model(image_np, masks):
236
+ """
237
+ Run the trained LogisticRegression classifier to predict live/dead per cell.
238
+ Returns (dead_count, alive_count, overlay_np, {cell_id: label}).
239
+ Requires VIABILITY_CLF and VIABILITY_SCALER to be loaded.
240
+ """
241
+ import numpy as np
242
+ cell_ids = np.unique(masks)
243
+ cell_ids = cell_ids[cell_ids > 0]
244
+ if len(cell_ids) == 0:
245
+ return 0, 0, image_np.copy(), {}
246
+
247
+ features = extract_cell_features(image_np, masks)
248
+ if not features:
249
+ return 0, 0, image_np.copy(), {}
250
+
251
+ import numpy as np
252
+ X = np.array([[f[c] for c in FEATURE_COLS_INFERENCE] for f in features], dtype=np.float32)
253
+
254
+ # replace any NaN/Inf with column median
255
+ for j in range(X.shape[1]):
256
+ bad = ~np.isfinite(X[:, j])
257
+ if bad.any():
258
+ X[bad, j] = float(np.nanmedian(X[:, j]))
259
+
260
+ X_scaled = VIABILITY_SCALER.transform(X)
261
+ predictions = VIABILITY_CLF.predict(X_scaled) # 0=live, 1=dead
262
+
263
+ label_map = {int(f["cell_id"]): int(p) for f, p in zip(features, predictions)}
264
+ overlay = draw_viability_overlay(image_np, masks, label_map)
265
+
266
+ dead = int(sum(predictions))
267
+ alive = int(len(predictions) - dead)
268
+ return dead, alive, overlay, label_map
269
+
270
+
271
+ def draw_viability_overlay(image_np, masks, label_map):
272
+ """
273
+ Draw coloured contours + cell-number labels onto image_np.
274
+ label_map: {cell_id: 0=live, 1=dead}
275
+ Returns a uint8 numpy array.
276
+ """
277
+ overlay = image_np.copy()
278
+ cell_ids = np.unique(masks)
279
+ cell_ids = cell_ids[cell_ids > 0]
280
+ cell_enum = {int(cid): idx + 1 for idx, cid in enumerate(sorted(cell_ids))}
281
+
282
+ for cid in cell_ids:
283
+ cid_int = int(cid)
284
+ label = label_map.get(cid_int, 0)
285
+ color = (220, 50, 50) if label == 1 else (50, 220, 80)
286
+ cell_mask = (masks == cid).astype(np.uint8)
287
+ contours, _ = cv2.findContours(cell_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
288
+ cv2.drawContours(overlay, contours, -1, color, thickness=2)
289
+
290
+ ys, xs = np.where(cell_mask)
291
+ if len(ys) > 0:
292
+ cx, cy = int(xs.mean()), int(ys.mean())
293
+ label_str = str(cell_enum[cid_int])
294
+ font = cv2.FONT_HERSHEY_SIMPLEX
295
+ font_scale = 0.35
296
+ thickness = 1
297
+ (tw, th), _ = cv2.getTextSize(label_str, font, font_scale, thickness)
298
+ cv2.rectangle(overlay,
299
+ (cx - tw//2 - 1, cy - th//2 - 1),
300
+ (cx + tw//2 + 1, cy + th//2 + 1),
301
+ (0, 0, 0), -1)
302
+ cv2.putText(overlay, label_str,
303
+ (cx - tw//2, cy + th//2),
304
+ font, font_scale, color, thickness, cv2.LINE_AA)
305
+ return overlay
306
+
307
+
308
+ def classify_cells_by_blueness(image_np, masks, threshold_bias):
309
+ """
310
+ Classify cells as dead (blue) or alive using an adaptive Otsu threshold
311
+ on per-cell blueness scores, with a user bias to fine-tune.
312
+
313
+ Args:
314
+ image_np: RGB image array
315
+ masks: Cellpose segmentation masks
316
+ threshold_bias: Slider value -50..+50; shifts Otsu threshold up/down.
317
+ Negative = more cells classified dead (looser).
318
+ Positive = fewer cells classified dead (stricter).
319
+ 0 = pure Otsu (fully automatic).
320
+
321
+ Returns:
322
+ dead_count, alive_count, colored_overlay, otsu_threshold, final_threshold
323
+ """
324
+
325
+ if len(image_np.shape) == 2:
326
+ image_np = cv2.cvtColor(image_np, cv2.COLOR_GRAY2RGB)
327
+ elif len(image_np.shape) == 3 and image_np.shape[2] == 4:
328
+ image_np = cv2.cvtColor(image_np, cv2.COLOR_RGBA2RGB)
329
+
330
+ hsv = cv2.cvtColor(image_np, cv2.COLOR_RGB2HSV)
331
+
332
+ hue = hsv[:, :, 0].astype(np.float32)
333
+ saturation = hsv[:, :, 1].astype(np.float32)
334
+
335
+ # Raw blueness: hue proximity to 115Β° Γ— saturation
336
+ hue_distance = np.minimum(np.abs(hue - 115), 180 - np.abs(hue - 115))
337
+ hue_score = np.maximum(0, 1 - hue_distance / 65)
338
+ blueness = hue_score * (saturation / 255.0)
339
+
340
+ # --- Compute per-cell mean blueness scores ---
341
+ cell_ids = np.unique(masks)
342
+ cell_ids = cell_ids[cell_ids > 0]
343
+
344
+ if len(cell_ids) == 0:
345
+ blank = image_np.copy()
346
+ return 0, 0, blank, 0.0, 0.0
347
+
348
+ cell_scores = np.array([np.mean(blueness[masks == cid]) for cid in cell_ids])
349
+
350
+ # --- Otsu on the distribution of per-cell scores ---
351
+ # cv2.threshold expects uint8; scale 0-1 β†’ 0-255
352
+ scores_u8 = (np.clip(cell_scores, 0, 1) * 255).astype(np.uint8)
353
+
354
+ if scores_u8.max() == scores_u8.min():
355
+ # All cells identical β†’ Otsu is undefined; use midpoint
356
+ otsu_threshold = float(scores_u8[0]) / 255.0
357
+ else:
358
+ # Reshape to a single-column image so cv2.threshold works
359
+ thresh_val, _ = cv2.threshold(
360
+ scores_u8.reshape(-1, 1), 0, 255,
361
+ cv2.THRESH_BINARY + cv2.THRESH_OTSU
362
+ )
363
+ otsu_threshold = thresh_val / 255.0
364
+
365
+ # --- Apply user bias: slider -50..+50 maps to Β±0.20 shift ---
366
+ bias = (threshold_bias / 50.0) * 0.20
367
+ final_threshold = float(np.clip(otsu_threshold + bias, 0.0, 1.0))
368
+
369
+ # --- Classify ---
370
+ dead_cells = [cid for cid, s in zip(cell_ids, cell_scores) if s > final_threshold]
371
+ alive_cells = [cid for cid, s in zip(cell_ids, cell_scores) if s <= final_threshold]
372
+
373
+ # --- Outline-only overlay on raw image with enumerated labels ---
374
+ final_overlay = image_np.copy()
375
+
376
+ # Compute a consistent enumeration order (cell_ids is already sorted ascending)
377
+ cell_enum = {cid: idx + 1 for idx, cid in enumerate(cell_ids)}
378
+
379
+ dead_set = set(dead_cells)
380
+ alive_set = set(alive_cells)
381
+
382
+ for cid in cell_ids:
383
+ cell_mask = (masks == cid).astype(np.uint8)
384
+ contours, _ = cv2.findContours(cell_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
385
+ color = (220, 50, 50) if cid in dead_set else (50, 220, 80)
386
+ cv2.drawContours(final_overlay, contours, -1, color, thickness=2)
387
+
388
+ # Draw enumeration label at centroid
389
+ ys, xs = np.where(cell_mask)
390
+ if len(ys) > 0:
391
+ cx, cy = int(xs.mean()), int(ys.mean())
392
+ label_str = str(cell_enum[cid])
393
+ font = cv2.FONT_HERSHEY_SIMPLEX
394
+ font_scale = 0.35
395
+ thickness = 1
396
+ (tw, th), _ = cv2.getTextSize(label_str, font, font_scale, thickness)
397
+ # Dark background rectangle for readability
398
+ cv2.rectangle(
399
+ final_overlay,
400
+ (cx - tw // 2 - 1, cy - th // 2 - 1),
401
+ (cx + tw // 2 + 1, cy + th // 2 + 1),
402
+ (0, 0, 0),
403
+ -1
404
+ )
405
+ cv2.putText(
406
+ final_overlay, label_str,
407
+ (cx - tw // 2, cy + th // 2),
408
+ font, font_scale, color, thickness, cv2.LINE_AA
409
+ )
410
+
411
+ return len(dead_cells), len(alive_cells), final_overlay, otsu_threshold, final_threshold
412
+
413
+
414
+ def measure_confluency(masks, image_np):
415
+ tot_pixels = image_np.shape[0] * image_np.shape[1]
416
+ cell_pixels = np.count_nonzero(masks)
417
+ confluency = cell_pixels / tot_pixels * 100
418
+ return confluency
419
+
420
+ def filter_mask_by_size(masks, minimum_pixels):
421
+ filtered_masks = masks.copy()
422
+ cell_ids = np.unique(masks)
423
+ cell_ids = cell_ids[cell_ids > 0]
424
+
425
+ removed_count = 0
426
+
427
+ for cell_id in cell_ids:
428
+ cell_mask = (masks == cell_id)
429
+ cell_pixels = np.count_nonzero(cell_mask)
430
+ if cell_pixels < minimum_pixels:
431
+ filtered_masks[cell_mask] = 0
432
+ removed_count += 1
433
+
434
+ unique_ids = np.unique(filtered_masks)
435
+ unique_ids = unique_ids[unique_ids > 0]
436
+
437
+ renumbered_masks = np.zeros_like(filtered_masks)
438
+ for new_id, old_id in enumerate(unique_ids, start=1):
439
+ renumbered_masks[filtered_masks == old_id] = new_id
440
+
441
+ return renumbered_masks, removed_count
442
+
443
+
444
+ def filter_mask_by_maxsize(masks, maximum_pixels):
445
+ filtered_masks = masks.copy()
446
+ cell_ids = np.unique(masks)
447
+ cell_ids = cell_ids[cell_ids > 0]
448
+
449
+ removed_count = 0
450
+ for cell_id in cell_ids:
451
+ cell_mask = (masks == cell_id)
452
+ cell_pixels = np.count_nonzero(cell_mask)
453
+ if cell_pixels > maximum_pixels:
454
+ filtered_masks[cell_mask] = 0
455
+ removed_count += 1
456
+
457
+ unique_ids = np.unique(filtered_masks)
458
+ unique_ids = unique_ids[unique_ids > 0]
459
+
460
+ renumbered_masks = np.zeros_like(filtered_masks)
461
+ for new_id, old_id in enumerate(unique_ids, start=1):
462
+ renumbered_masks[filtered_masks == old_id] = new_id
463
+
464
+ return renumbered_masks, removed_count
465
+
466
+
467
+ def rec_min_size(masks, q=25):
468
+ ids = np.unique(masks)
469
+ ids = ids[ids > 0]
470
+ if len(ids) == 0:
471
+ return 0
472
+ sizes = np.array([np.count_nonzero(masks == cid) for cid in ids])
473
+ return int(round(np.percentile(sizes, q)))
474
+
475
+
476
+ def apply_polygon_mask(image_pil, points_json):
477
+ """
478
+ Given a PIL image and a JSON string of [[x,y],...] points,
479
+ zero out everything outside the polygon and return a PIL image.
480
+ """
481
+ import json
482
+ if not points_json or points_json.strip() in ("", "[]"):
483
+ return image_pil
484
+ try:
485
+ pts = json.loads(points_json)
486
+ except Exception:
487
+ return image_pil
488
+ if len(pts) < 3:
489
+ return image_pil
490
+
491
+ image_np = np.array(image_pil)
492
+ h, w = image_np.shape[:2]
493
+ poly = np.array(pts, dtype=np.int32)
494
+ poly[:, 0] = np.clip(poly[:, 0], 0, w - 1)
495
+ poly[:, 1] = np.clip(poly[:, 1], 0, h - 1)
496
+ mask = np.zeros((h, w), dtype=np.uint8)
497
+ cv2.fillPoly(mask, [poly], 255)
498
+ if len(image_np.shape) == 3:
499
+ result = np.where(mask[:, :, np.newaxis] == 255, image_np, 0).astype(np.uint8)
500
+ else:
501
+ result = np.where(mask == 255, image_np, 0).astype(np.uint8)
502
+ return Image.fromarray(result)
503
+
504
+ def order_quad(pts):
505
+ """
506
+ Order 4 points as (top-left, top-right, bottom-right, bottom-left).
507
+ Sorts by angle about the centroid so the result stays correct for
508
+ rotated quads, where the sum/difference heuristic breaks down.
509
+ """
510
+ pts = np.asarray(pts, dtype=np.float32)
511
+ centre = pts.mean(axis=0)
512
+ angles = np.arctan2(pts[:, 1] - centre[1], pts[:, 0] - centre[0])
513
+ # Clockwise on screen (y grows downward) == increasing angle here
514
+ ordered = pts[np.argsort(angles)]
515
+ # Rotate so the corner nearest the top-left of the quad comes first
516
+ start = np.argmin(ordered.sum(axis=1))
517
+ return np.roll(ordered, -start, axis=0)
518
+
519
+
520
+ def quad_output_size(points):
521
+ """
522
+ Width and height, in source pixels, of the rectangle a 4-point quad
523
+ warps to. Shared by the warp itself and the exclusion-zone preview so
524
+ the two can never disagree about the ROI's dimensions.
525
+ """
526
+ tl, tr, br, bl = order_quad(points)
527
+ out_w = max(1, int(round(max(np.linalg.norm(br - bl), np.linalg.norm(tr - tl)))))
528
+ out_h = max(1, int(round(max(np.linalg.norm(tr - br), np.linalg.norm(tl - bl)))))
529
+ return out_w, out_h
530
+
531
+
532
+ def warp_polygon_to_square(image_np, points):
533
+ src = order_quad(points)
534
+ out_w, out_h = quad_output_size(points)
535
+
536
+ dst = np.array(
537
+ [[0, 0],
538
+ [out_w - 1, 0],
539
+ [out_w - 1, out_h - 1],
540
+ [0, out_h - 1]],
541
+ dtype=np.float32)
542
+
543
+ M = cv2.getPerspectiveTransform(src, dst)
544
+ warped = cv2.warpPerspective(image_np, M, (out_w, out_h))
545
+ return warped
546
+
547
+
548
+ def toggle_stereological_mode(use_stereology):
549
+ """Show/hide stereological controls based on checkbox"""
550
+ return gr.update(visible=use_stereology)
551
+
552
+
553
+ def exclusion_pct_for_preview(image_shape, crop_points, left_pct, top_pct):
554
+ """
555
+ Exclusion percentages restated relative to the FULL image.
556
+
557
+ The exclusion zones are applied to the cropped/warped ROI, so a given
558
+ percentage covers far fewer pixels than the same percentage of the whole
559
+ upload. Rescaling by the ROI's dimensions lets the full-image preview draw
560
+ a band of the true width instead of a misleadingly wide one.
561
+ """
562
+ h, w = image_shape[:2]
563
+ if not crop_points or len(crop_points) != 4 or w == 0 or h == 0:
564
+ return left_pct, top_pct
565
+ out_w, out_h = quad_output_size(crop_points)
566
+ return left_pct * out_w / float(w), top_pct * out_h / float(h)
567
+
568
+
569
+ def update_exclusion_preview(image, left_width, top_width, crop_points=None):
570
+ """Update the preview image with exclusion zone overlay"""
571
+ if image is None:
572
+ return None
573
+
574
+ image_np = np.array(image)
575
+ left_shown, top_shown = exclusion_pct_for_preview(
576
+ image_np.shape, crop_points, left_width, top_width
577
+ )
578
+ overlay = draw_exclusion_overlay(image_np, left_shown, top_shown)
579
+ return Image.fromarray(overlay)
580
+
581
+
582
+ # ---------------------------------------------------------------------------
583
+ # Patch segmentation
584
+ # ---------------------------------------------------------------------------
585
+ PATCH_SIZE = 512 # target patch side length
586
+ PATCH_OVERLAP = 64 # overlap border on each edge (pixels)
587
+ MIN_PATCH_DIM = 256 # don't bother patching if image fits comfortably
588
+
589
+
590
+ def _split_patches(image_np, patch_size=PATCH_SIZE, overlap=PATCH_OVERLAP):
591
+ """
592
+ Split image into overlapping patches.
593
+ Returns list of (patch_np, row_start, col_start) tuples.
594
+ """
595
+ h, w = image_np.shape[:2]
596
+ patches = []
597
+ row = 0
598
+ while row < h:
599
+ row_end = min(row + patch_size, h)
600
+ col = 0
601
+ while col < w:
602
+ col_end = min(col + patch_size, w)
603
+ patch = image_np[row:row_end, col:col_end]
604
+ patches.append((patch, row, col))
605
+ if col_end == w:
606
+ break
607
+ col += patch_size - overlap
608
+ if row_end == h:
609
+ break
610
+ row += patch_size - overlap
611
+ return patches
612
+
613
+
614
+ def _merge_patch_masks(patch_results, full_h, full_w, overlap=PATCH_OVERLAP):
615
+ """
616
+ Stitch per-patch masks into a single full-image mask.
617
+
618
+ Strategy:
619
+ - Each patch gets a unique ID offset so cell IDs never collide.
620
+ - Patches are pasted into the canvas using a priority canvas that
621
+ gives interior pixels precedence over overlap-border pixels.
622
+ - After pasting, cells whose centroids fall in the overlap zone
623
+ of two adjacent patches are deduplicated: if two cells from
624
+ different patches share >50% IoU they are the same cell β€” keep
625
+ the one whose centroid is furthest from a patch edge.
626
+ """
627
+ full_mask = np.zeros((full_h, full_w), dtype=np.int32)
628
+ # track which patch_idx owns each pixel (used for overlap resolution)
629
+ owner_map = np.full((full_h, full_w), -1, dtype=np.int32)
630
+ # distance-to-nearest-edge for the owning patch (higher = more central)
631
+ priority = np.zeros((full_h, full_w), dtype=np.float32)
632
+
633
+ id_offset = 0
634
+ patch_meta = [] # (offset, row_start, col_start, patch_h, patch_w)
635
+
636
+ for patch_idx, (mask_patch, row_start, col_start) in enumerate(patch_results):
637
+ ph, pw = mask_patch.shape
638
+ # offset all non-zero IDs so they're globally unique.
639
+ # Widen BEFORE adding: cellpose returns uint16 masks, and adding the
640
+ # offset in that dtype wraps once ids pass 65535.
641
+ mask_patch = mask_patch.astype(np.int32, copy=False)
642
+ shifted = np.where(mask_patch > 0, mask_patch + id_offset, 0).astype(np.int32)
643
+
644
+ # compute per-pixel priority = min distance to any patch edge
645
+ rows_idx = np.arange(ph)
646
+ cols_idx = np.arange(pw)
647
+ dist_r = np.minimum(rows_idx, ph - 1 - rows_idx) # (ph,)
648
+ dist_c = np.minimum(cols_idx, pw - 1 - cols_idx) # (pw,)
649
+ pri_patch = np.minimum(dist_r[:, None], dist_c[None, :]) # (ph, pw)
650
+
651
+ roi_full = full_mask [row_start:row_start+ph, col_start:col_start+pw]
652
+ roi_owner = owner_map [row_start:row_start+ph, col_start:col_start+pw]
653
+ roi_pri = priority [row_start:row_start+ph, col_start:col_start+pw]
654
+
655
+ # where this patch has higher priority, overwrite
656
+ better = pri_patch > roi_pri
657
+ roi_full [better] = shifted [better]
658
+ roi_owner[better] = patch_idx
659
+ roi_pri [better] = pri_patch [better]
660
+
661
+ max_id = int(mask_patch.max())
662
+ patch_meta.append((id_offset, row_start, col_start, ph, pw))
663
+ id_offset += max_id + 1
664
+
665
+ # --- Renumber to compact sequential IDs ---
666
+ unique_ids = np.unique(full_mask)
667
+ unique_ids = unique_ids[unique_ids > 0]
668
+ renumbered = np.zeros_like(full_mask)
669
+ for new_id, old_id in enumerate(unique_ids, start=1):
670
+ renumbered[full_mask == old_id] = new_id
671
+
672
+ return renumbered
673
+
674
+
675
+ def _segment_patch(args):
676
+ """Worker: run cellpose on a single patch. Called from a thread pool."""
677
+ patch_np, row_start, col_start, model_filename, hf_repo = args
678
+ # Each thread uses the shared loaded_models cache (GIL-safe for reads;
679
+ # model.eval() releases the GIL during GPU work so threads overlap.)
680
+ model_path = hf_hub_download(repo_id=hf_repo, filename=model_filename)
681
+ if model_filename in loaded_models:
682
+ model = loaded_models[model_filename]
683
+ else:
684
+ model = models.CellposeModel(gpu=True, pretrained_model=model_path)
685
+ loaded_models[model_filename] = model
686
+
687
+ mask, _, _ = model.eval(patch_np, diameter=None, channels=[0, 0])
688
+ return mask, row_start, col_start
689
+
690
+
691
+ def run_segmentation_patched(image_np, model_filename):
692
+ """
693
+ Split image into overlapping patches, run Cellpose on each in parallel,
694
+ then stitch back into a single full-resolution mask.
695
+ Falls back to whole-image segmentation if the image is small enough
696
+ that patching adds overhead without benefit.
697
+ """
698
+ h, w = image_np.shape[:2]
699
+ model_path = hf_hub_download(repo_id=HF_REPO_ID, filename=model_filename)
700
+ if model_filename in loaded_models:
701
+ model = loaded_models[model_filename]
702
+ else:
703
+ model = models.CellposeModel(gpu=True, pretrained_model=model_path)
704
+ loaded_models[model_filename] = model
705
+
706
+ # Small images: no benefit from patching
707
+ if max(h, w) <= MIN_PATCH_DIM * 2:
708
+ mask, _, _ = model.eval(image_np, diameter=None, channels=[0, 0])
709
+ return mask, 1 # 1 patch
710
+
711
+ patches = _split_patches(image_np)
712
+ n_patches = len(patches)
713
+
714
+ # Build argument list for the thread pool
715
+ args_list = [
716
+ (patch, r, c, model_filename, HF_REPO_ID)
717
+ for patch, r, c in patches
718
+ ]
719
+
720
+ patch_results = [] # (mask, row_start, col_start) in submission order
721
+
722
+ # ThreadPoolExecutor: GPU kernels release the GIL so threads overlap on GPU
723
+ with ThreadPoolExecutor(max_workers=min(n_patches, 4)) as pool:
724
+ futures = {pool.submit(_segment_patch, a): a for a in args_list}
725
+ for future in as_completed(futures):
726
+ mask_patch, row_start, col_start = future.result()
727
+ patch_results.append((mask_patch, row_start, col_start))
728
+
729
+ # Re-sort by (row, col) so stitching is deterministic
730
+ patch_results.sort(key=lambda x: (x[1], x[2]))
731
+
732
+ full_mask = _merge_patch_masks(patch_results, h, w)
733
+ return full_mask, n_patches
734
+
735
+
736
+ def render_segmentation(base_masks, processed_image_np,
737
+ use_stereology, left_exclusion, top_exclusion):
738
+ """
739
+ Apply the stereological exclusion to already-segmented masks and rebuild
740
+ the overlay. Split out of run_segmentation so the exclusion zones can be
741
+ adjusted after segmentation without re-running Cellpose.
742
+
743
+ Returns (masks, cell_count, confluency, overlay_pil, excluded_count).
744
+ """
745
+ if use_stereology:
746
+ masks, excluded_count, _ = apply_stereological_exclusion(
747
+ base_masks, left_exclusion, top_exclusion
748
+ )
749
+ else:
750
+ masks, excluded_count = base_masks.copy(), 0
751
+
752
+ cell_count = int(len(np.unique(masks)) - (1 if (masks == 0).any() else 0))
753
+ confluency = measure_confluency(masks, processed_image_np)
754
+
755
+ overlay = processed_image_np.copy().astype(np.float32)
756
+ if masks.max() > 0:
757
+ np.random.seed(42) # For consistent random colors
758
+ colors = np.random.randint(0, 255, size=(int(masks.max()) + 1, 3))
759
+ colors[0] = [0, 0, 0]
760
+ colored_mask = colors[masks]
761
+ alpha = 0.4
762
+ overlay = (1 - alpha) * overlay + alpha * colored_mask
763
+ overlay = np.clip(overlay, 0, 255).astype(np.uint8)
764
+
765
+ if use_stereology:
766
+ overlay = draw_exclusion_overlay(overlay, left_exclusion, top_exclusion)
767
+
768
+ return masks, cell_count, confluency, Image.fromarray(overlay), excluded_count
769
+
770
+
771
+ @spaces.GPU
772
+ def run_segmentation(image, model_choice, min_cell_size, max_cell_size,
773
+ use_stereology, left_exclusion, top_exclusion,
774
+ crop_points=None, use_min_filter=False, use_max_filter=False):
775
+ image_np = np.array(image)
776
+
777
+ # Crop BEFORE downscaling, so the ROI is sampled from the original pixels
778
+ # rather than upscaled out of a ≀MAX_SIDE working copy. The crop points were
779
+ # clicked on the full-resolution upload, so they are already in this space.
780
+ # (Need β‰₯3 points for a polygon.)
781
+ if crop_points and len(crop_points) >= 3:
782
+ import json
783
+ if len(crop_points) == 4:
784
+ # The perspective warp already discards everything outside the quad,
785
+ # so masking first would only round the corners off.
786
+ image_np = warp_polygon_to_square(image_np, crop_points)
787
+ else:
788
+ pts_json = json.dumps([[float(x), float(y)] for x, y in crop_points])
789
+ image_pil_masked = apply_polygon_mask(Image.fromarray(image_np), pts_json)
790
+ image_np = np.array(image_pil_masked)
791
+
792
+ # Cap the working image only after cropping β€” a small ROI now keeps its
793
+ # native detail, and a large one is still bounded for segmentation.
794
+ image_np = safe_resize(image_np)
795
+
796
+ # Un-annotated copy of whatever we actually segment, kept for cell thumbnails.
797
+ # Must be taken after the crop so it stays index-compatible with the masks.
798
+ raw_image_np = image_np.copy()
799
+
800
+
801
+ try:
802
+ model_filename = MODEL_OPTIONS[model_choice]
803
+
804
+ # Process image format to RGB
805
+ if len(image_np.shape) == 2:
806
+ processed_image_np = cv2.cvtColor(image_np, cv2.COLOR_GRAY2RGB)
807
+ elif len(image_np.shape) == 3 and image_np.shape[2] == 4:
808
+ processed_image_np = cv2.cvtColor(image_np, cv2.COLOR_RGBA2RGB)
809
+ else:
810
+ processed_image_np = image_np
811
+
812
+ # Run patch-parallel Cellpose segmentation
813
+ masks_raw, n_patches = run_segmentation_patched(processed_image_np, model_filename)
814
+
815
+ ids = np.unique(masks_raw)
816
+ ids = ids[ids > 0]
817
+
818
+ sizes = np.array([np.count_nonzero(masks_raw == cid) for cid in ids])
819
+
820
+ print("num_cells:", len(ids))
821
+ print("mean:", sizes.mean() if len(sizes) > 0 else 0)
822
+ print("median:", np.median(sizes) if len(sizes) > 0 else 0)
823
+ print("p90:", np.percentile(sizes, 90) if len(sizes) > 0 else 0)
824
+ print("max:", sizes.max() if len(sizes) > 0 else 0)
825
+
826
+ # Compute recommendation from RAW masks
827
+ recommend_min = rec_min_size(masks_raw)
828
+
829
+ # Surface the recommendation as text rather than writing it into the
830
+ # slider, which would overwrite a threshold the user had chosen.
831
+ if recommend_min > 0:
832
+ rec_msg = (f"*Recommended minimum: **{recommend_min}** pixels "
833
+ f"(25th percentile of detected object sizes). "
834
+ f"Enable the filter and leave the slider at 0 to use it.*")
835
+ else:
836
+ rec_msg = "*No objects detected β€” no minimum size recommendation.*"
837
+
838
+ # Size filters only run when their checkbox is ticked. When the minimum
839
+ # filter is on but the slider is still at 0, fall back to the
840
+ # recommendation; when it is off, nothing is filtered at all.
841
+ min_used = 0
842
+ if use_min_filter:
843
+ min_used = recommend_min if (min_cell_size == 0) else int(min_cell_size)
844
+
845
+ masks = masks_raw.copy()
846
+ removed_small = 0
847
+ removed_large = 0
848
+
849
+ if use_min_filter and min_used > 0:
850
+ masks, removed_small = filter_mask_by_size(masks, min_used)
851
+
852
+ if use_max_filter and max_cell_size > 0:
853
+ masks, removed_large = filter_mask_by_maxsize(masks, int(max_cell_size))
854
+
855
+ # Masks before exclusion are kept in state so the zones stay adjustable
856
+ base_masks = masks
857
+ masks, cell_count, confluency, overlay_pil, excluded_count = render_segmentation(
858
+ base_masks, processed_image_np, use_stereology, left_exclusion, top_exclusion
859
+ )
860
+
861
+ filter_msg = ""
862
+ if removed_small:
863
+ filter_msg += f"Removed {removed_small} small objects (< {min_used} pixels).\n"
864
+ if removed_large:
865
+ filter_msg += f"Removed {removed_large} large objects (> {int(max_cell_size)} pixels).\n"
866
+ if use_stereology and excluded_count > 0:
867
+ filter_msg += f"Stereological exclusion: {excluded_count} cells excluded (touching left/top zones).\n"
868
+
869
+ info_msg = ""
870
+ if filter_msg:
871
+ info_msg += filter_msg
872
+ info_msg += f"Segmentation complete! Found {cell_count} cells.\n"
873
+ info_msg += f"Confluency: {confluency:.1f}%\n"
874
+ info_msg += f"Processed as {n_patches} patch{'es' if n_patches > 1 else ''} (parallel).\n"
875
+ if use_stereology:
876
+ info_msg += f"Stereological counting enabled (Left: {left_exclusion}%, Top: {top_exclusion}%)\n"
877
+ info_msg += "Exclusion zones stay adjustable below without re-segmenting.\n"
878
+ info_msg += "Now run the viability classification model for viability assessment."
879
+
880
+ return (
881
+ cell_count,
882
+ overlay_pil,
883
+ info_msg,
884
+ gr.update(visible=True),
885
+ pack_array(masks),
886
+ pack_array(processed_image_np),
887
+ confluency,
888
+ gr.update(value=rec_msg),
889
+ pack_array(raw_image_np),
890
+ pack_array(base_masks),
891
+ )
892
+
893
+ except Exception as e:
894
+ import traceback
895
+ traceback.print_exc()
896
+ return (
897
+ 0,
898
+ None,
899
+ f"Error during segmentation: {str(e)}",
900
+ gr.update(visible=False),
901
+ None,
902
+ None,
903
+ 0.0,
904
+ gr.update(),
905
+ None,
906
+ None,
907
+ )
908
+
909
+
910
+ def run_viability(stored_masks, stored_image_np):
911
+ """Run model-based viability classification. Returns overlay + counts + label_map."""
912
+ if stored_masks is None or stored_image_np is None:
913
+ return None, 0, 0, 0.0, "Please run segmentation first.", {}
914
+ if VIABILITY_CLF is None:
915
+ return None, 0, 0, 0.0, "No viability model found. Add viability_clf.pkl and viability_scaler.pkl to the app directory.", {}
916
+
917
+ masks = unpack_array(stored_masks)
918
+ image_np = unpack_array(stored_image_np)
919
+
920
+ try:
921
+ dead, alive, overlay_np, label_map = classify_cells_by_model(image_np, masks)
922
+ total = alive + dead
923
+ viab_pct = (alive / total * 100) if total > 0 else 0.0
924
+ confluency = measure_confluency(masks, image_np)
925
+ info_msg = f"Total cells: {total}\nLive (green): {alive}\nDead (red): {dead}\n"
926
+ info_msg += f"Viability: {viab_pct:.1f}%\nConfluency: {confluency:.1f}%"
927
+ return Image.fromarray(overlay_np), alive, dead, viab_pct, info_msg, label_map
928
+ except Exception as e:
929
+ import traceback; traceback.print_exc()
930
+ return None, 0, 0, 0.0, f"Error: {str(e)}", {}
931
+
932
+
933
+ def pack_array(arr):
934
+ """
935
+ Serialise an array for gr.State.
936
+
937
+ Uses np.save rather than a PNG: label masks are int32 and routinely carry
938
+ more than 255 cell ids, which a uint8 PNG silently wraps (id 256 becomes
939
+ background). Dtype and values are preserved exactly.
940
+ """
941
+ buf = io.BytesIO()
942
+ np.save(buf, arr, allow_pickle=False)
943
+ return buf.getvalue()
944
+
945
+
946
+ def unpack_array(data):
947
+ buf = io.BytesIO(data)
948
+ try:
949
+ return np.load(buf, allow_pickle=False)
950
+ except ValueError:
951
+ # Legacy PNG-encoded state from an older session
952
+ buf.seek(0)
953
+ return np.array(Image.open(buf))
954
+
955
+
956
+ def save_tab_result(cell_count, confluency, viab_percent):
957
+ """Package per-tab results into a dict for Tab 5 summary."""
958
+ return {
959
+ "cell_count": float(cell_count) if cell_count is not None else None,
960
+ "confluency": float(confluency) if confluency is not None else None,
961
+ "viab_percent": float(viab_percent) if viab_percent is not None else None,
962
+ }
963
+
964
+
965
+ def compute_summary(r1, r2, r3, r4):
966
+ """Average cell count, confluency, and viability across tabs that have data."""
967
+ all_results = [r1, r2, r3, r4]
968
+ valid = [(i + 1, r) for i, r in enumerate(all_results)
969
+ if r is not None and r.get("cell_count") is not None]
970
+
971
+ if not valid:
972
+ return (
973
+ 0.0, 0.0, 0.0,
974
+ "No data yet β€” run segmentation in at least one tab, then click Refresh Summary."
975
+ )
976
+
977
+ avg_count = sum(r["cell_count"] for _, r in valid) / len(valid)
978
+ avg_conf = sum(r["confluency"] for _, r in valid) / len(valid)
979
+ avg_viab = sum(r["viab_percent"] for _, r in valid) / len(valid)
980
+
981
+ lines = [f"Tab {tab_num}: {r['cell_count']:.0f} cells | "
982
+ f"{r['confluency']:.1f}% confluency | "
983
+ f"{r['viab_percent']:.1f}% viability"
984
+ for tab_num, r in valid]
985
+ lines.append(f"\nAverages ({len(valid)} tab{'s' if len(valid) > 1 else ''}):")
986
+ lines.append(f" Cell count: {avg_count:.1f}")
987
+ lines.append(f" Confluency: {avg_conf:.1f}%")
988
+ lines.append(f" Viability: {avg_viab:.1f}%")
989
+
990
+ return avg_count, avg_conf, avg_viab, "\n".join(lines)
991
+
992
+
993
+ # ---------------------------------------------------------------------------
994
+ # Training data export β€” feature extraction per cell
995
+ # ---------------------------------------------------------------------------
996
+
997
+ def extract_cell_features(image_np, masks):
998
+ """
999
+ For every segmented cell, extract a fixed feature vector from the pixels
1000
+ inside its mask. Returns a list of dicts, one per cell.
1001
+
1002
+ Features:
1003
+ RGB channels β€” mean_r, mean_g, mean_b, std_r, std_g, std_b
1004
+ HSV channels β€” mean_h, mean_s, mean_v, std_s, std_v
1005
+ Ratios β€” blue_red_ratio, blue_green_ratio, rg_ratio
1006
+ Morphology β€” area_px, circularity
1007
+ Centre/edge profile β€” inner_brightness, peak_brightness,
1008
+ bright_spot_fraction, ring_darkness,
1009
+ centre_periphery_ratio, brightness_std_normalised
1010
+
1011
+ Profile zones are tuned to hemocytometer live-cell morphology:
1012
+ a small intense specular highlight at the centre surrounded by a dark
1013
+ navy membrane ring. Dead cells are pale blue-grey blobs with no ring
1014
+ and no bright spot.
1015
+ """
1016
+ if len(image_np.shape) == 2:
1017
+ image_np = cv2.cvtColor(image_np, cv2.COLOR_GRAY2RGB)
1018
+ elif image_np.shape[2] == 4:
1019
+ image_np = cv2.cvtColor(image_np, cv2.COLOR_RGBA2RGB)
1020
+
1021
+ hsv = cv2.cvtColor(image_np, cv2.COLOR_RGB2HSV).astype(np.float32)
1022
+
1023
+ h_img, w_img = image_np.shape[:2]
1024
+ grid_y, grid_x = np.mgrid[:h_img, :w_img]
1025
+
1026
+ cell_ids = np.unique(masks)
1027
+ cell_ids = cell_ids[cell_ids > 0]
1028
+ rows = []
1029
+
1030
+ for cid in cell_ids:
1031
+ cell_mask = (masks == cid)
1032
+ pixels_rgb = image_np[cell_mask].astype(np.float32)
1033
+ pixels_hsv = hsv[cell_mask]
1034
+
1035
+ r, g, b = pixels_rgb[:, 0], pixels_rgb[:, 1], pixels_rgb[:, 2]
1036
+ h, s, v = pixels_hsv[:, 0], pixels_hsv[:, 1], pixels_hsv[:, 2]
1037
+
1038
+ eps = 1e-6
1039
+ blue_red_ratio = b.mean() / (r.mean() + eps)
1040
+ blue_green_ratio = b.mean() / (g.mean() + eps)
1041
+ rg_ratio = r.mean() / (g.mean() + eps)
1042
+
1043
+ area_px = int(cell_mask.sum())
1044
+ contours, _ = cv2.findContours(
1045
+ cell_mask.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
1046
+ )
1047
+ perimeter = cv2.arcLength(contours[0], True) if contours else 1.0
1048
+ circularity = (4 * np.pi * area_px / (perimeter ** 2 + eps)) if perimeter > 0 else 0.0
1049
+
1050
+ ys_cell = grid_y[cell_mask].astype(np.float32)
1051
+ xs_cell = grid_x[cell_mask].astype(np.float32)
1052
+ centroid_y = ys_cell.mean()
1053
+ centroid_x = xs_cell.mean()
1054
+
1055
+ cell_radius = np.sqrt(area_px / np.pi) + eps
1056
+ dist_norm = np.sqrt((xs_cell - centroid_x)**2 + (ys_cell - centroid_y)**2) / cell_radius
1057
+
1058
+ v_all = hsv[:, :, 2][cell_mask]
1059
+
1060
+ # Tight inner core (15% radius) β€” captures specular highlight spot only
1061
+ inner_mask = dist_norm < 0.15
1062
+ # Membrane ring zone (20-60%) β€” dark navy ring on live cells
1063
+ ring_mask = (dist_norm >= 0.20) & (dist_norm <= 0.60)
1064
+ # Outer zone (>60%) β€” denominator for centre ratio
1065
+ outer_mask = dist_norm > 0.60
1066
+
1067
+ inner_brightness = float(v_all[inner_mask].mean()) if inner_mask.any() else float(v.mean())
1068
+ ring_brightness = float(v_all[ring_mask].mean()) if ring_mask.any() else float(v.mean())
1069
+ outer_brightness = float(v_all[outer_mask].mean()) if outer_mask.any() else float(v.mean())
1070
+
1071
+ # Peak V β€” specular spot is just a few pixels so mean dilutes it
1072
+ peak_brightness = float(v_all.max())
1073
+
1074
+ # Fraction of cell pixels with V > 200 (specular highlight region)
1075
+ bright_spot_fraction = float((v_all > 200).sum()) / (len(v_all) + eps)
1076
+
1077
+ # Ring darkness: ratio of ring zone to outer zone brightness
1078
+ # Live: ring << outer (dark membrane ring) -> ratio < 1
1079
+ # Dead: uniform blob -> ratio ~ 1
1080
+ ring_darkness = ring_brightness / (outer_brightness + eps)
1081
+
1082
+ centre_periphery_ratio = inner_brightness / (outer_brightness + eps)
1083
+
1084
+ brightness_std_normalised = float(v.std()) / (float(v.mean()) + eps)
1085
+
1086
+ rows.append({
1087
+ "cell_id": int(cid),
1088
+ "mean_r": float(r.mean()),
1089
+ "mean_g": float(g.mean()),
1090
+ "mean_b": float(b.mean()),
1091
+ "std_r": float(r.std()),
1092
+ "std_g": float(g.std()),
1093
+ "std_b": float(b.std()),
1094
+ "mean_h": float(h.mean()),
1095
+ "mean_s": float(s.mean()),
1096
+ "mean_v": float(v.mean()),
1097
+ "std_s": float(s.std()),
1098
+ "std_v": float(v.std()),
1099
+ "blue_red_ratio": round(blue_red_ratio, 5),
1100
+ "blue_green_ratio": round(blue_green_ratio, 5),
1101
+ "rg_ratio": round(rg_ratio, 5),
1102
+ "area_px": area_px,
1103
+ "circularity": round(float(circularity), 5),
1104
+ "inner_brightness": round(inner_brightness, 3),
1105
+ "peak_brightness": round(peak_brightness, 3),
1106
+ "bright_spot_fraction": round(bright_spot_fraction, 6),
1107
+ "ring_darkness": round(ring_darkness, 5),
1108
+ "centre_periphery_ratio": round(centre_periphery_ratio, 5),
1109
+ "brightness_std_normalised": round(brightness_std_normalised, 5),
1110
+ })
1111
+
1112
+ return rows
1113
+
1114
+ def attach_viability_labels(cell_features, masks, image_np, label_map=None):
1115
+ """
1116
+ Attach model predictions (from label_map) to each feature dict.
1117
+ label_map: {cell_id: 0=live, 1=dead} from classify_cells_by_model.
1118
+ If label_map is None, defaults all labels to 0 (live).
1119
+ """
1120
+ if not cell_features:
1121
+ return []
1122
+ labelled = []
1123
+ for feat in cell_features:
1124
+ row = dict(feat)
1125
+ cid = int(feat["cell_id"])
1126
+ row["label"] = int(label_map.get(cid, 0)) if label_map else 0
1127
+ row["corrected"] = False
1128
+ labelled.append(row)
1129
+ return labelled
1130
+
1131
+
1132
+ def export_cell_data_csv(cell_data):
1133
+ """Write cell_data list-of-dicts to a temp CSV and return the file path."""
1134
+ if not cell_data:
1135
+ return None
1136
+ tmp = tempfile.NamedTemporaryFile(
1137
+ mode="w", suffix=".csv", delete=False, newline=""
1138
+ )
1139
+ # Union of all keys across rows so any late-added keys (e.g. "corrected") are included
1140
+ fieldnames = list(dict.fromkeys(k for row in cell_data for k in row.keys()))
1141
+ writer = csv.DictWriter(tmp, fieldnames=fieldnames, extrasaction="ignore")
1142
+ writer.writeheader()
1143
+ writer.writerows(cell_data)
1144
+ tmp.close()
1145
+ return tmp.name
1146
+
1147
+
1148
+ def prepare_export(stored_masks, stored_image, threshold_bias):
1149
+ """
1150
+ Called by the Export button. Unpacks state, extracts features,
1151
+ attaches labels, writes CSV, returns (path, status_message).
1152
+ """
1153
+ if stored_masks is None or stored_image is None:
1154
+ return None, "Run segmentation first before exporting."
1155
+
1156
+ masks = unpack_array(stored_masks)
1157
+ image_np = unpack_array(stored_image)
1158
+
1159
+ features = extract_cell_features(image_np, masks)
1160
+ if not features:
1161
+ return None, "No cells found to export."
1162
+
1163
+ labelled = attach_viability_labels(features, masks, image_np, threshold_bias)
1164
+ path = export_cell_data_csv(labelled)
1165
+
1166
+ n = len(labelled)
1167
+ dead = sum(1 for r in labelled if r["label"] == 1)
1168
+ alive = n - dead
1169
+ msg = (f"Exported {n} cells ({alive} live, {dead} dead) β€” "
1170
+ f"threshold bias={threshold_bias:+d}.\n"
1171
+ f"Columns: {', '.join(list(labelled[0].keys())[:6])}… "
1172
+ f"({len(labelled[0])} total).")
1173
+ return path, msg
1174
+
1175
+
1176
+ # ---------------------------------------------------------------------------
1177
+ # Tab builder
1178
+ # ---------------------------------------------------------------------------
1179
+
1180
+ def draw_polygon_overlay(image_pil, points):
1181
+ """
1182
+ Draw numbered vertex dots and polygon edges onto a copy of image_pil.
1183
+ points: list of (x, y) tuples in original image pixel space.
1184
+ Returns a new PIL image.
1185
+ """
1186
+ img = image_pil.copy().convert("RGBA")
1187
+ overlay = Image.new("RGBA", img.size, (0, 0, 0, 0))
1188
+ draw = ImageDraw.Draw(overlay)
1189
+
1190
+ if len(points) >= 2:
1191
+ # Draw edges
1192
+ for i in range(len(points) - 1):
1193
+ draw.line([points[i], points[i + 1]], fill=(74, 170, 255, 220), width=3)
1194
+ if len(points) == 4:
1195
+ draw.line([points[-1], points[0]], fill=(74, 170, 255, 220), width=3)
1196
+ # Semi-transparent fill
1197
+ draw.polygon(points, fill=(74, 170, 255, 50))
1198
+
1199
+ # Draw vertex dots + numbers
1200
+ r = max(8, min(img.width, img.height) // 60)
1201
+ for i, (x, y) in enumerate(points):
1202
+ draw.ellipse([x - r, y - r, x + r, y + r],
1203
+ fill=(74, 170, 255, 255), outline=(255, 255, 255, 255))
1204
+ draw.text((x, y), str(i + 1), fill=(255, 255, 255, 255), anchor="mm")
1205
+
1206
+ combined = Image.alpha_composite(img, overlay)
1207
+ return combined.convert("RGB")
1208
+
1209
+
1210
+ def add_crop_point(image_pil, points, evt: gr.SelectData):
1211
+ """
1212
+ Called by gr.Image .select(). Appends the clicked coordinate,
1213
+ redraws the overlay, returns (updated_image, updated_points).
1214
+ Ignores clicks once 4 points are set.
1215
+ """
1216
+ if image_pil is None:
1217
+ return image_pil, points
1218
+ if points is None:
1219
+ points = []
1220
+ if len(points) >= 4:
1221
+ return draw_polygon_overlay(image_pil, points), points
1222
+
1223
+ x, y = int(evt.index[0]), int(evt.index[1])
1224
+ new_points = points + [(x, y)]
1225
+ return draw_polygon_overlay(image_pil, new_points), new_points
1226
+
1227
+
1228
+ def clear_crop_points(image_pil):
1229
+ """Reset polygon β€” return original image with no overlay and empty points."""
1230
+ return image_pil, []
1231
+
1232
+
1233
+
1234
+
1235
+
1236
+ # ---------------------------------------------------------------------------
1237
+ # Label correction grid
1238
+ # ---------------------------------------------------------------------------
1239
+
1240
+ THUMB_SIZE = 80 # each cell thumbnail is THUMB_SIZE Γ— THUMB_SIZE px
1241
+ GRID_COLS = 8 # thumbnails per row
1242
+ BORDER = 4 # coloured border thickness in px
1243
+ LABEL_H = 16 # height of the text label strip at the bottom of each thumb
1244
+
1245
+ def _crop_cell_thumb(image_np, masks, cid):
1246
+ """
1247
+ Return a tight square crop of the cell, padded to THUMB_SIZE Γ— THUMB_SIZE.
1248
+ """
1249
+ ys, xs = np.where(masks == cid)
1250
+ if len(ys) == 0:
1251
+ return Image.fromarray(np.zeros((THUMB_SIZE, THUMB_SIZE, 3), dtype=np.uint8))
1252
+
1253
+ y0, y1 = ys.min(), ys.max() + 1
1254
+ x0, x1 = xs.min(), xs.max() + 1
1255
+
1256
+ # add a small context border around the tight bounding box
1257
+ pad = max(4, int(max(y1 - y0, x1 - x0) * 0.15))
1258
+ h, w = image_np.shape[:2]
1259
+ y0c = max(0, y0 - pad)
1260
+ y1c = min(h, y1 + pad)
1261
+ x0c = max(0, x0 - pad)
1262
+ x1c = min(w, x1 + pad)
1263
+
1264
+ crop = image_np[y0c:y1c, x0c:x1c].copy()
1265
+
1266
+ # dim pixels that don't belong to this cell
1267
+ dim_mask = (masks[y0c:y1c, x0c:x1c] != cid)
1268
+ crop[dim_mask] = (crop[dim_mask] * 0.3).astype(np.uint8)
1269
+
1270
+ pil = Image.fromarray(crop).resize((THUMB_SIZE, THUMB_SIZE), Image.LANCZOS)
1271
+ return pil
1272
+
1273
+
1274
+ def build_correction_grid(image_np, masks, labelled_features, raw_image_np=None):
1275
+ """
1276
+ Render all cell thumbnails into a single PIL image grid.
1277
+ Each thumbnail has a coloured border: green=live(0), red=dead(1).
1278
+ A small number in the corner identifies the cell_id.
1279
+
1280
+ Returns the PIL grid image.
1281
+ Cell order in the grid matches the order of labelled_features.
1282
+ """
1283
+ if not labelled_features:
1284
+ placeholder = Image.fromarray(
1285
+ np.zeros((THUMB_SIZE, THUMB_SIZE, 3), dtype=np.uint8)
1286
+ )
1287
+ return placeholder
1288
+
1289
+ thumb_src = raw_image_np if raw_image_np is not None else image_np
1290
+
1291
+ n = len(labelled_features)
1292
+ n_cols = GRID_COLS
1293
+ n_rows = (n + n_cols - 1) // n_cols
1294
+
1295
+ cell_h = THUMB_SIZE + 2 * BORDER + LABEL_H
1296
+ cell_w = THUMB_SIZE + 2 * BORDER
1297
+
1298
+ grid_w = n_cols * cell_w
1299
+ grid_h = n_rows * cell_h
1300
+
1301
+ grid = Image.new("RGB", (grid_w, grid_h), (30, 30, 30))
1302
+ draw = ImageDraw.Draw(grid)
1303
+
1304
+ for idx, feat in enumerate(labelled_features):
1305
+ cid = feat["cell_id"]
1306
+ label = feat["label"] # 0=live, 1=dead (may have been corrected)
1307
+ color = (220, 50, 50) if label == 1 else (50, 200, 80)
1308
+
1309
+ thumb = _crop_cell_thumb(thumb_src, masks, cid)
1310
+
1311
+ col = idx % n_cols
1312
+ row = idx // n_cols
1313
+ x0 = col * cell_w
1314
+ y0 = row * cell_h
1315
+
1316
+ # coloured border rectangle
1317
+ draw.rectangle([x0, y0, x0 + cell_w - 1, y0 + cell_h - 1], outline=color, width=BORDER)
1318
+
1319
+ # paste thumbnail inside border
1320
+ grid.paste(thumb, (x0 + BORDER, y0 + BORDER))
1321
+
1322
+ # small cell-id label strip
1323
+ strip_y = y0 + BORDER + THUMB_SIZE
1324
+ draw.rectangle([x0, strip_y, x0 + cell_w - 1, y0 + cell_h - 1],
1325
+ fill=(20, 20, 20))
1326
+ draw.text((x0 + BORDER + 2, strip_y + 1),
1327
+ f"#{cid} {'D' if label == 1 else 'L'}",
1328
+ fill=color)
1329
+
1330
+ return grid
1331
+
1332
+
1333
+ def toggle_cell_label(labelled_features, image_np, masks, raw_image_np, evt: gr.SelectData):
1334
+ """
1335
+ Called when user taps the correction grid image.
1336
+ Maps the tap pixel coordinate back to which thumbnail was tapped,
1337
+ flips that cell's label, rebuilds and returns the updated grid.
1338
+ """
1339
+ if not labelled_features or image_np is None:
1340
+ return build_correction_grid(image_np, masks, labelled_features), labelled_features
1341
+
1342
+ cell_w = THUMB_SIZE + 2 * BORDER
1343
+ cell_h = THUMB_SIZE + 2 * BORDER + LABEL_H
1344
+
1345
+ px, py = int(evt.index[0]), int(evt.index[1])
1346
+ col = px // cell_w
1347
+ row = py // cell_h
1348
+ idx = row * GRID_COLS + col
1349
+
1350
+ if idx < 0 or idx >= len(labelled_features):
1351
+ return build_correction_grid(image_np, masks, labelled_features, raw_image_np), labelled_features
1352
+
1353
+ # Flip the label
1354
+ updated = list(labelled_features) # shallow copy of list
1355
+ cell = dict(updated[idx]) # copy the dict so we don't mutate in place
1356
+ cell["label"] = 1 - cell["label"] # 0β†’1 or 1β†’0
1357
+ cell["corrected"] = True
1358
+ updated[idx] = cell
1359
+
1360
+ grid = build_correction_grid(image_np, masks, updated, raw_image_np)
1361
+ n_corrected = sum(1 for f in updated if f.get("corrected"))
1362
+ return grid, updated, f"Tapped cell #{cell['cell_id']} β†’ {'Dead' if cell['label']==1 else 'Live'}. {n_corrected} correction(s) total."
1363
+
1364
+
1365
+ def prepare_export_corrected(stored_masks, stored_image, labelled_features, label_map):
1366
+ """Export CSV using labelled_features with any manual corrections applied."""
1367
+ if stored_masks is None or stored_image is None:
1368
+ return None, "Run segmentation first before exporting."
1369
+ masks = unpack_array(stored_masks)
1370
+ image_np = unpack_array(stored_image)
1371
+ if not labelled_features:
1372
+ features = extract_cell_features(image_np, masks)
1373
+ labelled_features = attach_viability_labels(features, masks, image_np, label_map)
1374
+ if not labelled_features:
1375
+ return None, "No cells found to export."
1376
+ path = export_cell_data_csv(labelled_features)
1377
+ n = len(labelled_features)
1378
+ dead = sum(1 for r in labelled_features if r["label"] == 1)
1379
+ alive = n - dead
1380
+ corrected = sum(1 for r in labelled_features if r.get("corrected"))
1381
+ msg = (f"Exported {n} cells ({alive} live, {dead} dead). "
1382
+ f"{corrected} label(s) manually corrected.")
1383
+ return path, msg
1384
+
1385
+ def build_tab(tab_index, masks_state, image_state, result_state):
1386
+ with gr.Tab(f"Tab {tab_index}"):
1387
+ gr.Markdown("Run segmentation")
1388
+
1389
+ # Per-tab state: list of (x,y) crop polygon points
1390
+ crop_points_state = gr.State(value=[])
1391
+ # Clean copy of the uploaded image (no polygon drawn on it)
1392
+ base_image_state = gr.State(value=None)
1393
+ #raw image state
1394
+ raw_image_state = gr.State(value=None)
1395
+ # Masks as segmented, before stereological exclusion β€” lets the zones
1396
+ # be re-chosen after segmentation without re-running the model.
1397
+ base_masks_state = gr.State(value=None)
1398
+
1399
+ with gr.Row():
1400
+ with gr.Column():
1401
+ img_input = gr.Image(
1402
+ type="pil",
1403
+ label="Upload image",
1404
+ image_mode="RGB",
1405
+ height=512
1406
+ )
1407
+
1408
+ gr.Markdown(
1409
+ "### Crop region (optional)\n"
1410
+ "Click/tap up to **4 points** on the image below to define the region "
1411
+ "to segment. The polygon will be drawn as you click. "
1412
+ "Leave empty to segment the full image."
1413
+ )
1414
+
1415
+ crop_display = gr.Image(
1416
+ type="pil",
1417
+ label="Click to set crop vertices (up to 4)",
1418
+ interactive=True,
1419
+ height=400,
1420
+ )
1421
+
1422
+ crop_status = gr.Markdown("*Upload an image to enable cropping*")
1423
+
1424
+ clear_crop_btn = gr.Button("βœ• Clear crop points", size="sm")
1425
+
1426
+ model_dropdown = gr.Dropdown(
1427
+ choices=list(MODEL_OPTIONS.keys()),
1428
+ label="Select Model",
1429
+ value="Hemocytometer Model"
1430
+ )
1431
+
1432
+ gr.Markdown("### Size Filters")
1433
+
1434
+ use_min_filter = gr.Checkbox(
1435
+ label="Enable minimum size filter",
1436
+ value=False,
1437
+ info="Remove objects smaller than the threshold below"
1438
+ )
1439
+ min_size_slider = gr.Slider(
1440
+ minimum=0,
1441
+ maximum=500,
1442
+ value=0,
1443
+ step=10,
1444
+ label="Minimum Cell Size (pixels)",
1445
+ interactive=False,
1446
+ )
1447
+ min_size_recommendation = gr.Markdown(
1448
+ value="*Run segmentation to see recommended minimum*",
1449
+ )
1450
+
1451
+ use_max_filter = gr.Checkbox(
1452
+ label="Enable maximum size filter",
1453
+ value=False,
1454
+ info="Remove objects larger than the threshold below"
1455
+ )
1456
+ max_size_slider = gr.Slider(
1457
+ minimum=0,
1458
+ maximum=10000,
1459
+ value=10000,
1460
+ step=10,
1461
+ label="Maximum Cell Size (pixels)",
1462
+ interactive=False,
1463
+ )
1464
+
1465
+ gr.Markdown("### Stereological Counting")
1466
+ use_stereo = gr.Checkbox(
1467
+ label="Enable Stereological Counting",
1468
+ value=False,
1469
+ info="Use unbiased stereological rules for cell counting"
1470
+ )
1471
+
1472
+ with gr.Group(visible=False) as stereo_controls:
1473
+ gr.Markdown("""
1474
+ **Stereological Counting Rules:**
1475
+ - Cells touching LEFT or TOP exclusion zones are EXCLUDED
1476
+ - Cells touching RIGHT or BOTTOM edges are INCLUDED
1477
+ - This provides unbiased counting for quantification
1478
+ """)
1479
+
1480
+ excl_preview = gr.Image(
1481
+ type="pil",
1482
+ label="Exclusion Zone Preview (Red = Excluded)",
1483
+ height=500
1484
+ )
1485
+
1486
+ left_excl = gr.Slider(
1487
+ minimum=0,
1488
+ maximum=50,
1489
+ value=10,
1490
+ step=1,
1491
+ label="Left Exclusion Width (%)",
1492
+ info="Width of left exclusion zone"
1493
+ )
1494
+
1495
+ top_excl = gr.Slider(
1496
+ minimum=0,
1497
+ maximum=50,
1498
+ value=10,
1499
+ step=1,
1500
+ label="Top Exclusion Width (%)",
1501
+ info="Width of top exclusion zone"
1502
+ )
1503
+
1504
+ segment_btn = gr.Button("πŸ”¬ Run Segmentation", variant="primary", size="lg")
1505
+
1506
+ with gr.Column():
1507
+ cell_count_out = gr.Number(label="Total Cells Detected", precision=0)
1508
+ confluency_out = gr.Number(label="Confluency (%)", precision=1)
1509
+ overlay_out = gr.Image(type="pil", label="Segmentation Result")
1510
+ info_out = gr.Textbox(label="Processing Info", lines=4)
1511
+
1512
+ with gr.Group(visible=False) as viability_section:
1513
+ gr.Markdown("### Viability Assessment (Trypan Blue)")
1514
+
1515
+ viab_run_btn = gr.Button("Run Viability Analysis", variant="primary")
1516
+
1517
+ with gr.Row():
1518
+ live_count_out = gr.Number(label="Live Cells (Green)", precision=0)
1519
+ dead_count_out = gr.Number(label="Dead Cells (Red)", precision=0)
1520
+
1521
+ viab_overlay = gr.Image(type="pil", label="Viability (Green=Live Β· Red=Dead)")
1522
+ viab_percent_out = gr.Number(label="Viability (%)", precision=1)
1523
+ viab_info = gr.Textbox(label="Analysis Results", lines=4)
1524
+
1525
+ gr.Markdown("### Label Correction & Export")
1526
+ gr.Markdown(
1527
+ "After running viability, click **Build correction grid** to review every cell. "
1528
+ "**Green border = Live, Red border = Dead** (model predictions). "
1529
+ "Tap any thumbnail to flip its label β€” the counts and overlay update instantly. "
1530
+ "Export the corrected CSV for retraining."
1531
+ )
1532
+
1533
+ build_grid_btn = gr.Button("πŸ”² Build correction grid", variant="secondary")
1534
+ labelled_state = gr.State(value=[])
1535
+ label_map_state = gr.State(value={})
1536
+
1537
+ correction_grid = gr.Image(
1538
+ type="pil",
1539
+ label="Tap a cell to flip its label (green=live Β· red=dead)",
1540
+ interactive=True,
1541
+ visible=False,
1542
+ )
1543
+ correction_status = gr.Markdown(visible=False)
1544
+
1545
+ with gr.Row():
1546
+ export_btn = gr.Button("⬇️ Export corrected CSV", variant="secondary")
1547
+ export_info = gr.Textbox(label="Export status", lines=2, interactive=False)
1548
+ export_file = gr.File(label="Download CSV", visible=False)
1549
+
1550
+ # ---- Event handlers ------------------------------------------------
1551
+
1552
+ use_stereo.change(
1553
+ fn=toggle_stereological_mode,
1554
+ inputs=[use_stereo],
1555
+ outputs=[stereo_controls]
1556
+ )
1557
+
1558
+ def on_image_upload(img):
1559
+ if img is None:
1560
+ return None, None, "*Upload an image to enable cropping*"
1561
+ return img, img, "*Image loaded β€” click up to 4 points to define crop region*"
1562
+
1563
+ img_input.change(
1564
+ fn=on_image_upload,
1565
+ inputs=[img_input],
1566
+ outputs=[crop_display, base_image_state, crop_status]
1567
+ ).then(fn=lambda: [], outputs=[crop_points_state])
1568
+
1569
+ # Grey the size sliders out while their filter is disabled, so the
1570
+ # checkbox state is visible rather than silently inferred.
1571
+ use_min_filter.change(
1572
+ fn=lambda on: gr.update(interactive=bool(on)),
1573
+ inputs=[use_min_filter], outputs=[min_size_slider]
1574
+ )
1575
+ use_max_filter.change(
1576
+ fn=lambda on: gr.update(interactive=bool(on)),
1577
+ inputs=[use_max_filter], outputs=[max_size_slider]
1578
+ )
1579
+
1580
+ def on_exclusion_change(stored_base, stored_image, use_stereo_on,
1581
+ left_pct, top_pct):
1582
+ """Re-apply the exclusion zones to existing masks β€” no re-segmentation."""
1583
+ if stored_base is None or stored_image is None:
1584
+ return (gr.update(), gr.update(), gr.update(),
1585
+ gr.update(), gr.update())
1586
+ base_masks = unpack_array(stored_base)
1587
+ image_np = unpack_array(stored_image)
1588
+ masks, cell_count, confluency, overlay_pil, excluded = render_segmentation(
1589
+ base_masks, image_np, use_stereo_on, left_pct, top_pct
1590
+ )
1591
+ if use_stereo_on:
1592
+ msg = (f"Stereological exclusion (Left: {left_pct}%, Top: {top_pct}%): "
1593
+ f"{excluded} cells excluded, {cell_count} counted.\n"
1594
+ f"Confluency: {confluency:.1f}%\n"
1595
+ f"Re-run viability classification to update live/dead counts.")
1596
+ else:
1597
+ msg = (f"Stereological exclusion off β€” {cell_count} cells counted.\n"
1598
+ f"Confluency: {confluency:.1f}%")
1599
+ return cell_count, overlay_pil, confluency, pack_array(masks), msg
1600
+
1601
+ exclusion_inputs = [base_masks_state, image_state, use_stereo,
1602
+ left_excl, top_excl]
1603
+ exclusion_outputs = [cell_count_out, overlay_out, confluency_out,
1604
+ masks_state, info_out]
1605
+ for _component in (use_stereo, left_excl, top_excl):
1606
+ _component.change(fn=on_exclusion_change,
1607
+ inputs=exclusion_inputs, outputs=exclusion_outputs)
1608
+
1609
+ excl_preview_inputs = [img_input, left_excl, top_excl, crop_points_state]
1610
+ img_input.change(fn=update_exclusion_preview,
1611
+ inputs=excl_preview_inputs, outputs=[excl_preview])
1612
+ left_excl.change(fn=update_exclusion_preview,
1613
+ inputs=excl_preview_inputs, outputs=[excl_preview])
1614
+ top_excl.change(fn=update_exclusion_preview,
1615
+ inputs=excl_preview_inputs, outputs=[excl_preview])
1616
+
1617
+ def on_crop_click(base_img, points, evt: gr.SelectData):
1618
+ updated_img, updated_pts = add_crop_point(base_img, points, evt)
1619
+ n = len(updated_pts)
1620
+ status = (f"*{n} / 4 points set β€” keep clicking*" if n < 4
1621
+ else "*4 points set βœ“ β€” click **βœ• Clear** to redo, or run segmentation*")
1622
+ return updated_img, updated_pts, status
1623
+
1624
+ crop_display.select(fn=on_crop_click,
1625
+ inputs=[base_image_state, crop_points_state],
1626
+ outputs=[crop_display, crop_points_state, crop_status]
1627
+ ).then(fn=update_exclusion_preview,
1628
+ inputs=excl_preview_inputs, outputs=[excl_preview])
1629
+
1630
+ def on_clear_crop(base_img):
1631
+ img, pts = clear_crop_points(base_img)
1632
+ return img, pts, "*Points cleared β€” click to set new vertices*"
1633
+
1634
+ clear_crop_btn.click(fn=on_clear_crop,
1635
+ inputs=[base_image_state],
1636
+ outputs=[crop_display, crop_points_state, crop_status]
1637
+ ).then(fn=update_exclusion_preview,
1638
+ inputs=excl_preview_inputs, outputs=[excl_preview])
1639
+
1640
+ segment_btn.click(
1641
+ fn=run_segmentation,
1642
+ inputs=[img_input, model_dropdown, min_size_slider, max_size_slider,
1643
+ use_stereo, left_excl, top_excl, crop_points_state,
1644
+ use_min_filter, use_max_filter],
1645
+ outputs=[cell_count_out, overlay_out, info_out, viability_section,
1646
+ masks_state, image_state, confluency_out,
1647
+ min_size_recommendation, raw_image_state, base_masks_state]
1648
+ )
1649
+
1650
+ # ---- Run Viability button -------------------------------------------
1651
+ def on_run_viability(stored_masks, stored_image):
1652
+ overlay, alive, dead, viab_pct, info, label_map = run_viability(stored_masks, stored_image)
1653
+ return overlay, alive, dead, viab_pct, info, label_map
1654
+
1655
+ viab_run_btn.click(
1656
+ fn=on_run_viability,
1657
+ inputs=[masks_state, image_state],
1658
+ outputs=[viab_overlay, live_count_out, dead_count_out,
1659
+ viab_percent_out, viab_info, label_map_state]
1660
+ ).then(
1661
+ fn=save_tab_result,
1662
+ inputs=[cell_count_out, confluency_out, viab_percent_out],
1663
+ outputs=[result_state]
1664
+ )
1665
+
1666
+ # ---- Build correction grid -----------------------------------------
1667
+ def on_build_grid(stored_masks, stored_image, label_map, stored_raw_image):
1668
+ if stored_masks is None or stored_image is None or not label_map:
1669
+ return (gr.update(visible=False), [],
1670
+ gr.update(value="*Run viability analysis first.*", visible=True))
1671
+ masks = unpack_array(stored_masks)
1672
+ image_np = unpack_array(stored_image)
1673
+ raw_image_np = unpack_array(stored_raw_image) if stored_raw_image is not None else None
1674
+ features = extract_cell_features(image_np, masks)
1675
+ labelled = attach_viability_labels(features, masks, image_np, label_map)
1676
+ if not labelled:
1677
+ return (gr.update(visible=False), [],
1678
+ gr.update(value="*No cells found.*", visible=True))
1679
+ grid = build_correction_grid(image_np, masks, labelled, raw_image_np)
1680
+ n = len(labelled)
1681
+ dead = sum(1 for r in labelled if r["label"] == 1)
1682
+ msg = (f"*{n} cells β€” {n-dead} live (green), {dead} dead (red). "
1683
+ f"Tap any thumbnail to flip its label.*")
1684
+ return gr.update(value=grid, visible=True), labelled, gr.update(value=msg, visible=True)
1685
+
1686
+ build_grid_btn.click(
1687
+ fn=on_build_grid,
1688
+ inputs=[masks_state, image_state, label_map_state, raw_image_state],
1689
+ outputs=[correction_grid, labelled_state, correction_status]
1690
+ )
1691
+
1692
+ # ---- Grid tap β€” flip label, update overlay + counts ----------------
1693
+ def on_grid_tap(labelled, stored_masks, stored_image, stored_raw_image, evt: gr.SelectData):
1694
+ if not labelled or stored_masks is None:
1695
+ return None, labelled, "", 0, 0, 0.0, None, {}
1696
+ masks = unpack_array(stored_masks)
1697
+ image_np = unpack_array(stored_image)
1698
+ raw_image_np = unpack_array(stored_raw_image) if stored_raw_image is not None else None
1699
+ grid, updated, msg = toggle_cell_label(labelled, image_np, masks, raw_image_np, evt)
1700
+
1701
+ # Rebuild label_map from corrected labelled list
1702
+ new_label_map = {int(f["cell_id"]): int(f["label"]) for f in updated}
1703
+ overlay_np = draw_viability_overlay(image_np, masks, new_label_map)
1704
+ dead = sum(1 for f in updated if f["label"] == 1)
1705
+ alive = len(updated) - dead
1706
+ total = alive + dead
1707
+ viab_pct = (alive / total * 100) if total > 0 else 0.0
1708
+
1709
+ return (grid, updated, f"*{msg}*",
1710
+ alive, dead, viab_pct,
1711
+ Image.fromarray(overlay_np), new_label_map)
1712
+
1713
+ correction_grid.select(
1714
+ fn=on_grid_tap,
1715
+ inputs=[labelled_state, masks_state, image_state, raw_image_state],
1716
+ outputs=[correction_grid, labelled_state, correction_status,
1717
+ live_count_out, dead_count_out, viab_percent_out,
1718
+ viab_overlay, label_map_state]
1719
+ )
1720
+
1721
+ # ---- Export --------------------------------------------------------
1722
+ def on_export(stored_masks, stored_image, labelled, label_map):
1723
+ path, msg = prepare_export_corrected(stored_masks, stored_image, labelled, label_map)
1724
+ if path is None:
1725
+ return gr.update(visible=False), msg
1726
+ return gr.update(value=path, visible=True), msg
1727
+
1728
+ export_btn.click(
1729
+ fn=on_export,
1730
+ inputs=[masks_state, image_state, labelled_state, label_map_state],
1731
+ outputs=[export_file, export_info]
1732
+ )
1733
+
1734
+
1735
+
1736
+ # ---------------------------------------------------------------------------
1737
+ # Gradio interface
1738
+ # ---------------------------------------------------------------------------
1739
+ with gr.Blocks(
1740
+ title="CellposeCellCounter",
1741
+ theme=gr.themes.Soft(),
1742
+ ) as demo:
1743
+ gr.Markdown("# CellposeCellCounter")
1744
+ gr.Markdown("For accurate cell confluency, crop the image to display only desired area. Note that some image file types are not yet supported. PNG and JPEG are preferred.")
1745
+
1746
+ # Shared mask/image state (one pair per tab so tabs don't clobber each other)
1747
+ masks_states = [gr.State(value=None) for _ in range(4)]
1748
+ image_states = [gr.State(value=None) for _ in range(4)]
1749
+ result_states = [gr.State(value=None) for _ in range(4)]
1750
+
1751
+ # Build Tabs 1–4 with a loop
1752
+ for i in range(4):
1753
+ build_tab(i + 1, masks_states[i], image_states[i], result_states[i])
1754
+
1755
+ # -------------------------------------------------------------------------
1756
+ # Tab 5 β€” Summary
1757
+ # -------------------------------------------------------------------------
1758
+ with gr.Tab("Tab 5 β€” Summary"):
1759
+ gr.Markdown("## Average Results Across All Tabs")
1760
+ gr.Markdown(
1761
+ "Run segmentation in one or more tabs, "
1762
+ "then click **Refresh Summary** to see the averages."
1763
+ )
1764
+
1765
+ refresh_btn = gr.Button("πŸ”„ Refresh Summary", variant="primary", size="lg")
1766
+
1767
+ with gr.Row():
1768
+ avg_count_out = gr.Number(label="Avg Cell Count", precision=1)
1769
+ avg_conf_out = gr.Number(label="Avg Confluency (%)", precision=1)
1770
+ avg_viab_out = gr.Number(label="Avg Viability (%)", precision=1)
1771
+
1772
+ summary_box = gr.Textbox(label="Per-Tab Breakdown", lines=10)
1773
+
1774
+ refresh_btn.click(
1775
+ fn=compute_summary,
1776
+ inputs=result_states, # list of 4 gr.State components
1777
+ outputs=[avg_count_out, avg_conf_out, avg_viab_out, summary_box]
1778
+ )
1779
+
1780
+
1781
+
1782
+ if __name__ == "__main__":
1783
+ demo.launch()