LiangLabUMB commited on
Commit
ceb77c3
·
verified ·
1 Parent(s): f4a4d5c

Delete app.py

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