LiangLabUMB commited on
Commit
e3ada61
·
verified ·
1 Parent(s): 826178c

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +243 -141
app.py CHANGED
@@ -9,80 +9,42 @@ from PIL import Image, ImageDraw
9
  import io
10
  from huggingface_hub import hf_hub_download
11
  import base64
 
12
  import csv
13
  import joblib
14
  import os
15
 
16
- HF_REPO_ID = "myang4218/cellposemodel"
17
- HF_REPO_ID2 = "LiangLabUMB/viability_model"
18
- HF_REPO_CPSAM = "mouseland/cellpose-sam"
19
  MODEL_OPTIONS = {
20
  "Hemocytometer Model": "hemocytometermodel.npy",
21
- "General Model": "generalmodel.npy",
22
- "Cellpose SAMv2": "cpsam_v2",
23
  }
24
- MODEL_REPOS = {
25
- "hemocytometermodel.npy": HF_REPO_ID,
26
- "generalmodel.npy": HF_REPO_ID,
27
- "cpsam_v2": HF_REPO_CPSAM,
28
- }
29
-
30
-
31
- # Viability classifier is loaded LAZILY (only when viability is first requested).
32
- # Loading XGBoost / the pickled model at module level initialises a CUDA context
33
- # in the main process, which conflicts with ZeroGPU's per-call GPU allocation and
34
- # leaves segmentation requests stuck in the GPU queue. Keeping the main process
35
- # CUDA-clean until viability is explicitly called avoids this.
36
- VIABILITY_CLF = None
37
- VIABILITY_SCALER = None
38
- _VIABILITY_LOADED = False
39
- _VIABILITY_ERROR = None
40
-
41
-
42
- def ensure_viability_loaded():
43
- """
44
- Load the XGBoost classifier + scaler on first use, pinned to CPU.
45
 
46
- XGBoost initialises a CUDA context if it thinks a GPU is available. Once
47
- that context exists in the main process, every subsequent ZeroGPU fork for
48
- run_segmentation inherits a dirty CUDA state and hangs in the GPU queue.
49
- We force the loaded booster onto CPU so it never touches CUDA — viability
50
- inference on ~20 features per cell is trivially fast on CPU anyway.
51
- """
52
- global VIABILITY_CLF, VIABILITY_SCALER, _VIABILITY_LOADED, _VIABILITY_ERROR
53
- if _VIABILITY_LOADED:
54
- return
55
- try:
56
- import xgboost
57
- _clf_path = hf_hub_download(repo_id=HF_REPO_ID2, filename="viability_xgb_clf.pkl")
58
- _scaler_path = hf_hub_download(repo_id=HF_REPO_ID2, filename="viability_xgb_scaler.pkl")
59
- VIABILITY_CLF = joblib.load(_clf_path)
60
- VIABILITY_SCALER = joblib.load(_scaler_path)
61
-
62
- # Force CPU inference so XGBoost never initialises a CUDA context.
63
- try:
64
- VIABILITY_CLF.set_params(device="cpu", predictor="cpu_predictor", tree_method="hist")
65
- except Exception:
66
- pass
67
- try:
68
- booster = VIABILITY_CLF.get_booster()
69
- booster.set_param({"device": "cpu", "predictor": "cpu_predictor"})
70
- except Exception:
71
- pass
72
-
73
- _VIABILITY_LOADED = True
74
- print("✓ Viability classifier loaded (lazy, CPU-pinned).")
75
- except Exception as e:
76
- _VIABILITY_ERROR = str(e)
77
- print(f"Viability classifier failed to load: {e}")
78
-
79
- # mobile safe resize limits
80
  MAX_SIDE = 1024
81
  MAX_PIXELS = 1024 * 1024
82
 
83
 
84
  def safe_resize(image_np):
85
-
 
 
 
86
  h, w = image_np.shape[:2]
87
  total = h * w
88
 
@@ -190,7 +152,11 @@ FEATURE_COLS_INFERENCE = [
190
 
191
 
192
  def classify_cells_by_model(image_np, masks):
193
-
 
 
 
 
194
  import numpy as np
195
  cell_ids = np.unique(masks)
196
  cell_ids = cell_ids[cell_ids > 0]
@@ -222,7 +188,11 @@ def classify_cells_by_model(image_np, masks):
222
 
223
 
224
  def draw_viability_overlay(image_np, masks, label_map):
225
-
 
 
 
 
226
  overlay = image_np.copy()
227
  cell_ids = np.unique(masks)
228
  cell_ids = cell_ids[cell_ids > 0]
@@ -250,10 +220,114 @@ def draw_viability_overlay(image_np, masks, label_map):
250
  (0, 0, 0), -1)
251
  cv2.putText(overlay, label_str,
252
  (cx - tw//2, cy + th//2),
253
- font, font_scale, color, thickness, cv2.LINE_AA)
254
  return overlay
255
 
256
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
257
 
258
 
259
  def measure_confluency(masks, image_np):
@@ -377,10 +451,12 @@ def warp_polygon_to_square(image_np, points):
377
 
378
 
379
  def toggle_stereological_mode(use_stereology):
 
380
  return gr.update(visible=use_stereology)
381
 
382
 
383
  def update_exclusion_preview(image, left_width, top_width):
 
384
  if image is None:
385
  return None
386
 
@@ -389,8 +465,9 @@ def update_exclusion_preview(image, left_width, top_width):
389
  return Image.fromarray(overlay)
390
 
391
 
 
392
  # Patch segmentation
393
-
394
  PATCH_SIZE = 512 # target patch side length
395
  PATCH_OVERLAP = 64 # overlap border on each edge (pixels)
396
  MIN_PATCH_DIM = 256 # don't bother patching if image fits comfortably
@@ -478,42 +555,62 @@ def _merge_patch_masks(patch_results, full_h, full_w, overlap=PATCH_OVERLAP):
478
  return renumbered
479
 
480
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
481
  def run_segmentation_patched(image_np, model_filename):
482
  """
483
- Split image into overlapping patches, run Cellpose sequentially on each,
484
  then stitch back into a single full-resolution mask.
485
-
486
- NOTE: ThreadPoolExecutor was removed because ZeroGPU allocates the GPU
487
- only to the main thread. Worker threads spawned by ThreadPoolExecutor run
488
- outside the GPU context, causing a 3-minute timeout waiting for a GPU
489
- grant that never arrives for the worker threads.
490
- Sequential patching within the @spaces.GPU context is correct and fast.
491
  """
492
  h, w = image_np.shape[:2]
493
- repo = MODEL_REPOS.get(model_filename, HF_REPO_ID)
494
- # hf_hub_download caches the weights file on disk, so this only hits the
495
- # network once — subsequent calls read from the local cache.
496
- model_path = hf_hub_download(repo_id=repo, filename=model_filename)
497
-
498
- # IMPORTANT: build the CellposeModel fresh on every call. ZeroGPU allocates
499
- # a DIFFERENT physical GPU for each @spaces.GPU invocation, so a model cached
500
- # from a previous call holds a CUDA context bound to a now-deallocated GPU.
501
- # Reusing it makes the next request hang in the GPU queue. Constructing the
502
- # model is cheap because the weights are already on local disk.
503
- model = models.CellposeModel(gpu=True, pretrained_model=model_path)
504
 
505
  # Small images: no benefit from patching
506
  if max(h, w) <= MIN_PATCH_DIM * 2:
507
- mask, _, _ = model.eval(image_np, diameter=None)
508
- return mask, 1
509
 
510
  patches = _split_patches(image_np)
511
  n_patches = len(patches)
512
 
513
- patch_results = []
514
- for patch, row_start, col_start in patches:
515
- mask_patch, _, _ = model.eval(patch, diameter=None)
516
- patch_results.append((mask_patch, row_start, col_start))
 
 
 
 
 
 
 
 
 
 
 
 
 
517
 
518
  full_mask = _merge_patch_masks(patch_results, h, w)
519
  return full_mask, n_patches
@@ -521,7 +618,6 @@ def run_segmentation_patched(image_np, model_filename):
521
 
522
  @spaces.GPU
523
  def run_segmentation(image, model_choice, min_cell_size, max_cell_size,
524
- use_min_filter, use_max_filter,
525
  use_stereology, left_exclusion, top_exclusion,
526
  crop_points=None):
527
  image_np = np.array(image)
@@ -565,18 +661,21 @@ def run_segmentation(image, model_choice, min_cell_size, max_cell_size,
565
  print("p90:", np.percentile(sizes, 90) if len(sizes) > 0 else 0)
566
  print("max:", sizes.max() if len(sizes) > 0 else 0)
567
 
568
- # Compute recommendation from RAW masks (always shown, never auto-applied)
569
  recommend_min = rec_min_size(masks_raw)
570
 
571
- # Apply filters only if their checkboxes are enabled
 
 
 
572
  masks = masks_raw.copy()
573
  removed_small = 0
574
  removed_large = 0
575
 
576
- if use_min_filter and int(min_cell_size) > 0:
577
- masks, removed_small = filter_mask_by_size(masks, int(min_cell_size))
578
 
579
- if use_max_filter and max_cell_size > 0:
580
  masks, removed_large = filter_mask_by_maxsize(masks, int(max_cell_size))
581
 
582
  # Apply stereological exclusion if enabled
@@ -588,7 +687,7 @@ def run_segmentation(image, model_choice, min_cell_size, max_cell_size,
588
 
589
  filter_msg = ""
590
  if removed_small:
591
- filter_msg += f"Removed {removed_small} small objects (< {int(min_cell_size)} pixels).\n"
592
  if removed_large:
593
  filter_msg += f"Removed {removed_large} large objects (> {int(max_cell_size)} pixels).\n"
594
  if use_stereology and excluded_count > 0:
@@ -630,7 +729,7 @@ def run_segmentation(image, model_choice, min_cell_size, max_cell_size,
630
  pack_array(masks),
631
  pack_array(processed_image_np),
632
  confluency,
633
- f"Recommended minimum: **{recommend_min} px** (25th percentile of detected cell sizes)",
634
  pack_array(raw_image_np),
635
  )
636
 
@@ -645,20 +744,17 @@ def run_segmentation(image, model_choice, min_cell_size, max_cell_size,
645
  None,
646
  None,
647
  0.0,
648
- "",
649
  None,
650
  )
651
 
652
 
653
  def run_viability(stored_masks, stored_image_np):
 
654
  if stored_masks is None or stored_image_np is None:
655
  return None, 0, 0, 0.0, "Please run segmentation first.", {}
656
-
657
- # Lazy-load the classifier the first time viability is requested
658
- ensure_viability_loaded()
659
  if VIABILITY_CLF is None:
660
- err = _VIABILITY_ERROR or "unknown error"
661
- return None, 0, 0, 0.0, f"Viability model failed to load: {err}", {}
662
 
663
  masks = unpack_array(stored_masks)
664
  image_np = unpack_array(stored_image_np)
@@ -677,20 +773,14 @@ def run_viability(stored_masks, stored_image_np):
677
 
678
 
679
  def pack_array(arr):
680
- """
681
- Serialise a numpy array to a base64 string for gr.State storage.
682
- Uses numpy .npy format (preserves int32 exactly, no 255 truncation)
683
- encoded as base64 so it is a plain Python string — safe for ZeroGPU
684
- state serialisation which cannot handle raw bytes objects.
685
- """
686
  buf = io.BytesIO()
687
- np.save(buf, arr)
688
- return base64.b64encode(buf.getvalue()).decode("ascii")
689
 
690
 
691
  def unpack_array(data):
692
- buf = io.BytesIO(base64.b64decode(data))
693
- return np.load(buf, allow_pickle=False)
694
 
695
 
696
  def save_tab_result(cell_count, confluency, viab_percent):
@@ -730,12 +820,29 @@ def compute_summary(r1, r2, r3, r4):
730
  return avg_count, avg_conf, avg_viab, "\n".join(lines)
731
 
732
 
733
-
734
  # Training data export — feature extraction per cell
735
-
736
 
737
  def extract_cell_features(image_np, masks):
738
-
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
739
  if len(image_np.shape) == 2:
740
  image_np = cv2.cvtColor(image_np, cv2.COLOR_GRAY2RGB)
741
  elif image_np.shape[2] == 4:
@@ -896,8 +1003,9 @@ def prepare_export(stored_masks, stored_image, threshold_bias):
896
  return path, msg
897
 
898
 
899
-
900
  # Tab builder
 
901
 
902
  def draw_polygon_overlay(image_pil, points):
903
  """
@@ -955,12 +1063,14 @@ def clear_crop_points(image_pil):
955
 
956
 
957
 
 
958
  # Label correction grid
 
959
 
960
- THUMB_SIZE = 80
961
- GRID_COLS = 10
962
- BORDER = 4
963
- LABEL_H = 16
964
 
965
  def _crop_cell_thumb(image_np, masks, cid):
966
  """
@@ -992,7 +1102,14 @@ def _crop_cell_thumb(image_np, masks, cid):
992
 
993
 
994
  def build_correction_grid(image_np, masks, labelled_features, raw_image_np=None):
995
-
 
 
 
 
 
 
 
996
  if not labelled_features:
997
  placeholder = Image.fromarray(
998
  np.zeros((THUMB_SIZE, THUMB_SIZE, 3), dtype=np.uint8)
@@ -1139,29 +1256,14 @@ def build_tab(tab_index, masks_state, image_state, result_state):
1139
  value="Hemocytometer Model"
1140
  )
1141
 
1142
- gr.Markdown("### Size Filters")
1143
-
1144
- use_min_filter = gr.Checkbox(
1145
- label="Enable minimum size filter",
1146
- value=False,
1147
- info="Remove objects smaller than the threshold below"
1148
- )
1149
  min_size_slider = gr.Slider(
1150
  minimum=0,
1151
  maximum=500,
1152
  value=0,
1153
  step=10,
1154
- label="Minimum Cell Size (pixels)",
1155
- )
1156
- min_size_recommendation = gr.Markdown(
1157
- value="*Run segmentation to see recommended minimum*",
1158
  )
1159
 
1160
- use_max_filter = gr.Checkbox(
1161
- label="Enable maximum size filter",
1162
- value=False,
1163
- info="Remove objects larger than the threshold below"
1164
- )
1165
  max_size_slider = gr.Slider(
1166
  minimum=0,
1167
  maximum=10000,
@@ -1303,10 +1405,9 @@ def build_tab(tab_index, masks_state, image_state, result_state):
1303
  segment_btn.click(
1304
  fn=run_segmentation,
1305
  inputs=[img_input, model_dropdown, min_size_slider, max_size_slider,
1306
- use_min_filter, use_max_filter,
1307
  use_stereo, left_excl, top_excl, crop_points_state],
1308
  outputs=[cell_count_out, overlay_out, info_out, viability_section,
1309
- masks_state, image_state, confluency_out, min_size_recommendation, raw_image_state]
1310
  )
1311
 
1312
  # ---- Run Viability button -------------------------------------------
@@ -1395,8 +1496,9 @@ def build_tab(tab_index, masks_state, image_state, result_state):
1395
 
1396
 
1397
 
 
1398
  # Gradio interface
1399
-
1400
  with gr.Blocks(
1401
  title="CellposeCellCounter",
1402
  theme=gr.themes.Soft(),
 
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
 
 
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]
 
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]
 
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):
 
451
 
452
 
453
  def toggle_stereological_mode(use_stereology):
454
+ """Show/hide stereological controls based on checkbox"""
455
  return gr.update(visible=use_stereology)
456
 
457
 
458
  def update_exclusion_preview(image, left_width, top_width):
459
+ """Update the preview image with exclusion zone overlay"""
460
  if image is None:
461
  return None
462
 
 
465
  return Image.fromarray(overlay)
466
 
467
 
468
+ # ---------------------------------------------------------------------------
469
  # Patch segmentation
470
+ # ---------------------------------------------------------------------------
471
  PATCH_SIZE = 512 # target patch side length
472
  PATCH_OVERLAP = 64 # overlap border on each edge (pixels)
473
  MIN_PATCH_DIM = 256 # don't bother patching if image fits comfortably
 
555
  return renumbered
556
 
557
 
558
+ def _segment_patch(args):
559
+ """Worker: run cellpose on a single patch. Called from a thread pool."""
560
+ patch_np, row_start, col_start, model_filename, hf_repo = args
561
+ # Each thread uses the shared loaded_models cache (GIL-safe for reads;
562
+ # model.eval() releases the GIL during GPU work so threads overlap.)
563
+ model_path = hf_hub_download(repo_id=hf_repo, filename=model_filename)
564
+ if model_filename in loaded_models:
565
+ model = loaded_models[model_filename]
566
+ else:
567
+ model = models.CellposeModel(gpu=True, pretrained_model=model_path)
568
+ loaded_models[model_filename] = model
569
+
570
+ mask, _, _ = model.eval(patch_np, diameter=None, channels=[0, 0])
571
+ return mask, row_start, col_start
572
+
573
+
574
  def run_segmentation_patched(image_np, model_filename):
575
  """
576
+ Split image into overlapping patches, run Cellpose on each in parallel,
577
  then stitch back into a single full-resolution mask.
578
+ Falls back to whole-image segmentation if the image is small enough
579
+ that patching adds overhead without benefit.
 
 
 
 
580
  """
581
  h, w = image_np.shape[:2]
582
+ model_path = hf_hub_download(repo_id=HF_REPO_ID, filename=model_filename)
583
+ if model_filename in loaded_models:
584
+ model = loaded_models[model_filename]
585
+ else:
586
+ model = models.CellposeModel(gpu=True, pretrained_model=model_path)
587
+ loaded_models[model_filename] = model
 
 
 
 
 
588
 
589
  # Small images: no benefit from patching
590
  if max(h, w) <= MIN_PATCH_DIM * 2:
591
+ mask, _, _ = model.eval(image_np, diameter=None, channels=[0, 0])
592
+ return mask, 1 # 1 patch
593
 
594
  patches = _split_patches(image_np)
595
  n_patches = len(patches)
596
 
597
+ # Build argument list for the thread pool
598
+ args_list = [
599
+ (patch, r, c, model_filename, HF_REPO_ID)
600
+ for patch, r, c in patches
601
+ ]
602
+
603
+ patch_results = [] # (mask, row_start, col_start) in submission order
604
+
605
+ # ThreadPoolExecutor: GPU kernels release the GIL so threads overlap on GPU
606
+ with ThreadPoolExecutor(max_workers=min(n_patches, 4)) as pool:
607
+ futures = {pool.submit(_segment_patch, a): a for a in args_list}
608
+ for future in as_completed(futures):
609
+ mask_patch, row_start, col_start = future.result()
610
+ patch_results.append((mask_patch, row_start, col_start))
611
+
612
+ # Re-sort by (row, col) so stitching is deterministic
613
+ patch_results.sort(key=lambda x: (x[1], x[2]))
614
 
615
  full_mask = _merge_patch_masks(patch_results, h, w)
616
  return full_mask, n_patches
 
618
 
619
  @spaces.GPU
620
  def run_segmentation(image, model_choice, min_cell_size, max_cell_size,
 
621
  use_stereology, left_exclusion, top_exclusion,
622
  crop_points=None):
623
  image_np = np.array(image)
 
661
  print("p90:", np.percentile(sizes, 90) if len(sizes) > 0 else 0)
662
  print("max:", sizes.max() if len(sizes) > 0 else 0)
663
 
664
+ # Compute recommendation from RAW masks
665
  recommend_min = rec_min_size(masks_raw)
666
 
667
+ # If user sets slider to 0, use the recommendation
668
+ min_used = recommend_min if (min_cell_size == 0) else int(min_cell_size)
669
+
670
+ # Apply filters
671
  masks = masks_raw.copy()
672
  removed_small = 0
673
  removed_large = 0
674
 
675
+ if min_used > 0:
676
+ masks, removed_small = filter_mask_by_size(masks, min_used)
677
 
678
+ if max_cell_size > 0:
679
  masks, removed_large = filter_mask_by_maxsize(masks, int(max_cell_size))
680
 
681
  # Apply stereological exclusion if enabled
 
687
 
688
  filter_msg = ""
689
  if removed_small:
690
+ filter_msg += f"Removed {removed_small} small objects (< {min_used} pixels).\n"
691
  if removed_large:
692
  filter_msg += f"Removed {removed_large} large objects (> {int(max_cell_size)} pixels).\n"
693
  if use_stereology and excluded_count > 0:
 
729
  pack_array(masks),
730
  pack_array(processed_image_np),
731
  confluency,
732
+ gr.update(value=recommend_min),
733
  pack_array(raw_image_np),
734
  )
735
 
 
744
  None,
745
  None,
746
  0.0,
747
+ gr.update(),
748
  None,
749
  )
750
 
751
 
752
  def run_viability(stored_masks, stored_image_np):
753
+ """Run model-based viability classification. Returns overlay + counts + label_map."""
754
  if stored_masks is None or stored_image_np is None:
755
  return None, 0, 0, 0.0, "Please run segmentation first.", {}
 
 
 
756
  if VIABILITY_CLF is None:
757
+ return None, 0, 0, 0.0, "No viability model found. Add viability_clf.pkl and viability_scaler.pkl to the app directory.", {}
 
758
 
759
  masks = unpack_array(stored_masks)
760
  image_np = unpack_array(stored_image_np)
 
773
 
774
 
775
  def pack_array(arr):
776
+ pil = Image.fromarray(arr.astype(np.uint8))
 
 
 
 
 
777
  buf = io.BytesIO()
778
+ pil.save(buf, format="PNG")
779
+ return buf.getvalue()
780
 
781
 
782
  def unpack_array(data):
783
+ return np.array(Image.open(io.BytesIO(data)))
 
784
 
785
 
786
  def save_tab_result(cell_count, confluency, viab_percent):
 
820
  return avg_count, avg_conf, avg_viab, "\n".join(lines)
821
 
822
 
823
+ # ---------------------------------------------------------------------------
824
  # Training data export — feature extraction per cell
825
+ # ---------------------------------------------------------------------------
826
 
827
  def extract_cell_features(image_np, masks):
828
+ """
829
+ For every segmented cell, extract a fixed feature vector from the pixels
830
+ inside its mask. Returns a list of dicts, one per cell.
831
+
832
+ Features:
833
+ RGB channels — mean_r, mean_g, mean_b, std_r, std_g, std_b
834
+ HSV channels — mean_h, mean_s, mean_v, std_s, std_v
835
+ Ratios — blue_red_ratio, blue_green_ratio, rg_ratio
836
+ Morphology — area_px, circularity
837
+ Centre/edge profile — inner_brightness, peak_brightness,
838
+ bright_spot_fraction, ring_darkness,
839
+ centre_periphery_ratio, brightness_std_normalised
840
+
841
+ Profile zones are tuned to hemocytometer live-cell morphology:
842
+ a small intense specular highlight at the centre surrounded by a dark
843
+ navy membrane ring. Dead cells are pale blue-grey blobs with no ring
844
+ and no bright spot.
845
+ """
846
  if len(image_np.shape) == 2:
847
  image_np = cv2.cvtColor(image_np, cv2.COLOR_GRAY2RGB)
848
  elif image_np.shape[2] == 4:
 
1003
  return path, msg
1004
 
1005
 
1006
+ # ---------------------------------------------------------------------------
1007
  # Tab builder
1008
+ # ---------------------------------------------------------------------------
1009
 
1010
  def draw_polygon_overlay(image_pil, points):
1011
  """
 
1063
 
1064
 
1065
 
1066
+ # ---------------------------------------------------------------------------
1067
  # Label correction grid
1068
+ # ---------------------------------------------------------------------------
1069
 
1070
+ THUMB_SIZE = 80 # each cell thumbnail is THUMB_SIZE × THUMB_SIZE px
1071
+ GRID_COLS = 8 # thumbnails per row
1072
+ BORDER = 4 # coloured border thickness in px
1073
+ LABEL_H = 16 # height of the text label strip at the bottom of each thumb
1074
 
1075
  def _crop_cell_thumb(image_np, masks, cid):
1076
  """
 
1102
 
1103
 
1104
  def build_correction_grid(image_np, masks, labelled_features, raw_image_np=None):
1105
+ """
1106
+ Render all cell thumbnails into a single PIL image grid.
1107
+ Each thumbnail has a coloured border: green=live(0), red=dead(1).
1108
+ A small number in the corner identifies the cell_id.
1109
+
1110
+ Returns the PIL grid image.
1111
+ Cell order in the grid matches the order of labelled_features.
1112
+ """
1113
  if not labelled_features:
1114
  placeholder = Image.fromarray(
1115
  np.zeros((THUMB_SIZE, THUMB_SIZE, 3), dtype=np.uint8)
 
1256
  value="Hemocytometer Model"
1257
  )
1258
 
 
 
 
 
 
 
 
1259
  min_size_slider = gr.Slider(
1260
  minimum=0,
1261
  maximum=500,
1262
  value=0,
1263
  step=10,
1264
+ label="Minimum Cell Size (pixels). Leave at zero for automated recommendation",
 
 
 
1265
  )
1266
 
 
 
 
 
 
1267
  max_size_slider = gr.Slider(
1268
  minimum=0,
1269
  maximum=10000,
 
1405
  segment_btn.click(
1406
  fn=run_segmentation,
1407
  inputs=[img_input, model_dropdown, min_size_slider, max_size_slider,
 
1408
  use_stereo, left_excl, top_excl, crop_points_state],
1409
  outputs=[cell_count_out, overlay_out, info_out, viability_section,
1410
+ masks_state, image_state, confluency_out, min_size_slider, raw_image_state]
1411
  )
1412
 
1413
  # ---- Run Viability button -------------------------------------------
 
1496
 
1497
 
1498
 
1499
+ # ---------------------------------------------------------------------------
1500
  # Gradio interface
1501
+ # ---------------------------------------------------------------------------
1502
  with gr.Blocks(
1503
  title="CellposeCellCounter",
1504
  theme=gr.themes.Soft(),