Upload 4 files

README.md

@@ -1,5 +1,5 @@
 ---
-title: CellposeCellCounter_Mobile_v3c
+title: CellposeCellCounter_Mobile
 emoji: 🚀
 colorFrom: purple
 colorTo: pink

app.py

@@ -548,53 +548,77 @@ def _merge_patch_masks(patch_results, full_h, full_w, overlap=PATCH_OVERLAP):
         id_offset += max_id + 1
 
     # --- Renumber to compact sequential IDs ---
+    # Lookup table, not a loop over ids. The old form was
+    #     for new_id, old_id in enumerate(unique_ids, 1):
+    #         renumbered[full_mask == old_id] = new_id
+    # i.e. one full-array scan PER CELL -- O(cells x pixels). At ~450 cells on a
+    # 1024 px mask that is ~950 M element ops (measured 158 ms vs 12 ms here),
+    # and it ran inside the @spaces.GPU window, so it burnt ZeroGPU quota doing
+    # CPU integer relabelling. Output is bit-identical.
     unique_ids = np.unique(full_mask)
     unique_ids = unique_ids[unique_ids > 0]
-    renumbered = np.zeros_like(full_mask)
-    for new_id, old_id in enumerate(unique_ids, start=1):
-        renumbered[full_mask == old_id] = new_id
+    if unique_ids.size == 0:
+        return np.zeros_like(full_mask)
+    lut = np.zeros(int(full_mask.max()) + 1, dtype=full_mask.dtype)
+    lut[unique_ids] = np.arange(1, unique_ids.size + 1, dtype=full_mask.dtype)
+    renumbered = lut[full_mask]
 
     return renumbered
 
 
 def _segment_patch(args):
-    """Worker: run cellpose on a single patch. Called from a thread pool."""
-    patch_np, row_start, col_start, model_filename, hf_repo = args
-    # Each thread uses the shared loaded_models cache (GIL-safe for reads;
-    # model.eval() releases the GIL during GPU work so threads overlap.)
-    model_path = hf_hub_download(repo_id=hf_repo, filename=model_filename)
-    if model_filename in loaded_models:
-        model = loaded_models[model_filename]
-    else:
-        model = models.CellposeModel(gpu=True, pretrained_model=model_path)
-        loaded_models[model_filename] = model
+    """Worker: run cellpose on a single patch. Called from a thread pool.
 
+    The model is built ONCE by the caller and handed in. It used to be looked up
+    here, which meant an hf_hub_download cache round-trip per patch (nine of them
+    on a 1024 px image) plus a check-then-set race on loaded_models between the
+    pool threads. Both were pure overhead inside the GPU window.
+    """
+    patch_np, row_start, col_start, model = args
     mask, _, _ = model.eval(patch_np, diameter=None, channels=[0, 0])
     return mask, row_start, col_start
 
 
-@spaces.GPU(duration=30)
-def run_segmentation_patched(image_np, model_filename):
-    """
-    Split image into overlapping patches, run Cellpose on each in parallel,
-    then stitch back into a single full-resolution mask.
-
-    The @spaces.GPU decorator sits HERE rather than on run_segmentation because
-    ZeroGPU quota is charged for the whole time the GPU is attached, and the
-    caller does a lot of CPU-only work -- decoding a 12 MP photo, the perspective
-    warp, size filtering, overlay rendering. Holding an A10G through all of that
-    burnt visitors' daily quota on NumPy.
-    Falls back to whole-image segmentation if the image is small enough
-    that patching adds overhead without benefit.
+# ---------------------------------------------------------------------------
+# ZeroGPU duration budget
+#
+# Quota is charged on the duration you REQUEST, not the time you actually use,
+# and remaining quota is what sets your place in the queue -- so over-requesting
+# is paid for twice, once in quota and again in waiting. Tune these two numbers
+# from the "GPU x.xx s" figure the app already prints in the processing info box.
+# ---------------------------------------------------------------------------
+GPU_SETUP_S     = 20    # model construction + GPU attach: paid ONCE per call
+GPU_PER_IMAGE_S = 18    # marginal cost of each image in the batch
+GPU_MAX_S       = 180   # ceiling: ZeroGPU rejects very long requests outright
+
+
+def _gpu_duration(images, model_filename):
+    """Dynamic duration, so a 1-image call does not reserve a 4-image budget."""
+    try:
+        n = max(1, len(images))
+    except TypeError:
+        n = 1
+    return int(min(GPU_MAX_S, GPU_SETUP_S + GPU_PER_IMAGE_S * n))
+
+
+def _get_model(model_filename):
+    """Build (or reuse) the Cellpose model. Called once per GPU acquisition."""
+    if model_filename in loaded_models:
+        return loaded_models[model_filename]
+    model_path = hf_hub_download(repo_id=HF_REPO_ID, filename=model_filename)
+    model = models.CellposeModel(gpu=True, pretrained_model=model_path)
+    loaded_models[model_filename] = model
+    return model
+
+
+def _segment_image_with_model(image_np, model):
+    """Patch, segment and stitch ONE image using an already-built model.
+
+    Falls back to whole-image segmentation if the image is small enough that
+    patching adds overhead without benefit.
     """
     _t_gpu = time.perf_counter()
     h, w = image_np.shape[:2]
-    model_path = hf_hub_download(repo_id=HF_REPO_ID, filename=model_filename)
-    if model_filename in loaded_models:
-        model = loaded_models[model_filename]
-    else:
-        model = models.CellposeModel(gpu=True, pretrained_model=model_path)
-        loaded_models[model_filename] = model
 
     # Small images: no benefit from patching
     if max(h, w) <= MIN_PATCH_DIM * 2:
@@ -604,13 +628,10 @@ def run_segmentation_patched(image_np, model_filename):
     patches = _split_patches(image_np)
     n_patches = len(patches)
 
-    # Build argument list for the thread pool
-    args_list = [
-        (patch, r, c, model_filename, HF_REPO_ID)
-        for patch, r, c in patches
-    ]
+    # The model goes in the args now -- see _segment_patch.
+    args_list = [(patch, r, c, model) for patch, r, c in patches]
 
-    patch_results = []  # (mask, row_start, col_start) in submission order
+    patch_results = []  # (mask, row_start, col_start) in completion order
 
     # ThreadPoolExecutor: GPU kernels release the GIL so threads overlap on GPU
     with ThreadPoolExecutor(max_workers=min(n_patches, 4)) as pool:
@@ -626,6 +647,34 @@ def run_segmentation_patched(image_np, model_filename):
     return full_mask, n_patches, time.perf_counter() - _t_gpu
 
 
+@spaces.GPU(duration=_gpu_duration)
+def run_segmentation_patched_many(images, model_filename):
+    """Segment a LIST of images inside ONE ZeroGPU acquisition.
+
+    This is the point of the batching change. The four hemocytometer squares
+    used to take four separate GPU allocations, which meant four waits in the
+    ZeroGPU queue and four constructions of the Cellpose model -- module state
+    does not survive from one allocation to the next, so `loaded_models` never
+    actually helped across squares.
+
+    The decorator sits HERE rather than on run_segmentation because ZeroGPU
+    quota is charged for the whole time the GPU is attached, and the caller does
+    a lot of CPU-only work -- decoding a 12 MP photo, the perspective warp, size
+    filtering, overlay rendering. Holding an A10G through all of that burnt
+    visitors' daily quota on NumPy.
+
+    Returns a list of (mask, n_patches, gpu_seconds), one per input image.
+    """
+    model = _get_model(model_filename)
+    return [_segment_image_with_model(im, model) for im in images]
+
+
+def run_segmentation_patched(image_np, model_filename):
+    """Single-image wrapper. One image, one GPU acquisition -- unchanged for the
+    comprehensive and confluency tabs, which segment one image at a time."""
+    return run_segmentation_patched_many([image_np], model_filename)[0]
+
+
 def render_segmentation(base_masks, processed_image_np,
                         use_stereology, left_exclusion, top_exclusion):
     """
@@ -661,16 +710,19 @@ def render_segmentation(base_masks, processed_image_np,
     return masks, cell_count, confluency, Image.fromarray(overlay), excluded_count
 
 
-def run_segmentation(image, model_choice, min_cell_size, max_cell_size,
-                     use_stereology, left_exclusion, top_exclusion,
-                     crop_points=None, use_min_filter=False, use_max_filter=False):
-    _t_start = time.perf_counter()
+def _prep_for_segmentation(image, model_choice, crop_points):
+    """CPU-only preparation: crop / perspective-warp, downscale, RGB-normalise.
+
+    Split out of run_segmentation so a whole batch of images can be prepared
+    BEFORE a single GPU acquisition, instead of one acquisition per image.
+    Returns (processed_image_np, raw_image_np, model_filename).
+    """
     image_np = np.array(image)
 
     # Crop BEFORE downscaling, so the ROI is sampled from the original pixels
-    # rather than upscaled out of a ≤MAX_SIDE working copy. The crop points were
+    # rather than upscaled out of a <=MAX_SIDE working copy. The crop points were
     # clicked on the full-resolution upload, so they are already in this space.
-    # (Need ≥3 points for a polygon.)
+    # (Need >=3 points for a polygon.)
     if crop_points and len(crop_points) >= 3:
         import json
         if len(crop_points) == 4:
@@ -682,7 +734,7 @@ def run_segmentation(image, model_choice, min_cell_size, max_cell_size,
             image_pil_masked = apply_polygon_mask(Image.fromarray(image_np), pts_json)
             image_np = np.array(image_pil_masked)
 
-    # Cap the working image only after cropping — a small ROI now keeps its
+    # Cap the working image only after cropping -- a small ROI now keeps its
     # native detail, and a large one is still bounded for segmentation.
     image_np = safe_resize(image_np)
 
@@ -690,21 +742,48 @@ def run_segmentation(image, model_choice, min_cell_size, max_cell_size,
     # Must be taken after the crop so it stays index-compatible with the masks.
     raw_image_np = image_np.copy()
 
+    model_filename = MODEL_OPTIONS[model_choice]
 
-    try:
-        model_filename = MODEL_OPTIONS[model_choice]
+    # Process image format to RGB
+    if len(image_np.shape) == 2:
+        processed_image_np = cv2.cvtColor(image_np, cv2.COLOR_GRAY2RGB)
+    elif len(image_np.shape) == 3 and image_np.shape[2] == 4:
+        processed_image_np = cv2.cvtColor(image_np, cv2.COLOR_RGBA2RGB)
+    else:
+        processed_image_np = image_np
+
+    return processed_image_np, raw_image_np, model_filename
+
+
+def run_segmentation(image, model_choice, min_cell_size, max_cell_size,
+                     use_stereology, left_exclusion, top_exclusion,
+                     crop_points=None, use_min_filter=False, use_max_filter=False,
+                     _prepped=None, _gpu_result=None, _t_start_override=None):
+    """Segment one image and build the full result tuple.
+
+    _prepped / _gpu_result / _t_start_override are used by the batched path
+    (_segment_many): they let the same post-processing tail run against work
+    that has ALREADY been prepared and segmented, so the four squares can share
+    one GPU acquisition. Left unset, behaviour is exactly as before.
+    """
+    _t_start = time.perf_counter() if _t_start_override is None else _t_start_override
 
-        # Process image format to RGB
-        if len(image_np.shape) == 2:
-            processed_image_np = cv2.cvtColor(image_np, cv2.COLOR_GRAY2RGB)
-        elif len(image_np.shape) == 3 and image_np.shape[2] == 4:
-            processed_image_np = cv2.cvtColor(image_np, cv2.COLOR_RGBA2RGB)
+    try:
+        if _prepped is not None:
+            processed_image_np, raw_image_np, model_filename = _prepped
         else:
-            processed_image_np = image_np
+            # Now inside the try: a warp failure used to propagate out of
+            # run_segmentation and take the whole four-square run down with it.
+            # It now fails that one square and leaves the others countable.
+            processed_image_np, raw_image_np, model_filename = _prep_for_segmentation(
+                image, model_choice, crop_points)
 
         # Run patch-parallel Cellpose segmentation
-        masks_raw, n_patches, gpu_seconds = run_segmentation_patched(
-            processed_image_np, model_filename)
+        if _gpu_result is not None:
+            masks_raw, n_patches, gpu_seconds = _gpu_result
+        else:
+            masks_raw, n_patches, gpu_seconds = run_segmentation_patched(
+                processed_image_np, model_filename)
 
         # Same single pass feeds the log line and the recommendation below;
         # this used to be computed from scratch twice.
@@ -1807,14 +1886,8 @@ def _crop_picker(label):
     return st
 
 
-def _segment_one(image, points, model_name, use_min, min_px,
-                 use_stereo, left_pct, top_pct, want_viability):
-    """Segment (and optionally classify) a single image. Returns a dict."""
-    if image is None:
-        return None
-    out = run_segmentation(image, model_name, min_px, 10000,
-                           use_stereo, left_pct, top_pct, points,
-                           use_min, False)
+def _result_from_out(out, want_viability):
+    """Turn the run_segmentation tuple into the dict the tabs consume."""
     (cell_count, seg_overlay, info_msg, _vis, packed_masks, packed_img,
      confluency, _rec, packed_raw, _base, seg_seconds, _ov) = out
     if packed_masks is None:
@@ -1827,6 +1900,7 @@ def _segment_one(image, points, model_name, use_min, min_px,
            "alive": None, "dead": None, "viab": None, "viab_overlay": None}
 
     if want_viability:
+        # CPU-only (sklearn), so it stays outside the GPU window on purpose.
         v_overlay, alive, dead, viab_pct, _vinfo, lm = run_viability(
             packed_masks, packed_img)
         res.update(alive=alive, dead=dead, viab=viab_pct,
@@ -1834,6 +1908,75 @@ def _segment_one(image, points, model_name, use_min, min_px,
     return res
 
 
+def _segment_one(image, points, model_name, use_min, min_px,
+                 use_stereo, left_pct, top_pct, want_viability):
+    """Segment (and optionally classify) a single image. Returns a dict."""
+    if image is None:
+        return None
+    out = run_segmentation(image, model_name, min_px, 10000,
+                           use_stereo, left_pct, top_pct, points,
+                           use_min, False)
+    return _result_from_out(out, want_viability)
+
+
+def _segment_many(items, model_name, use_min, min_px,
+                  use_stereo, left_pct, top_pct, want_viability):
+    """Segment several (image, points) pairs with ONE ZeroGPU acquisition.
+
+    Produces the same per-image dicts as calling _segment_one in a loop, but the
+    squares share a single trip through the ZeroGPU queue and a single Cellpose
+    model construction instead of one each. All the CPU work -- warping,
+    filtering, overlays, viability -- stays outside the GPU window as before.
+
+    Returns a list the same length as `items`; an entry is None if that image
+    was None, or {"error": ...} if that one square failed. One square failing
+    never takes the others down.
+    """
+    t0 = time.perf_counter()
+    n = len(items)
+    results = [None] * n
+    prepped = [None] * n
+
+    # --- CPU: prepare everything first, no GPU held ---
+    for i, (image, points) in enumerate(items):
+        if image is None:
+            continue
+        try:
+            prepped[i] = _prep_for_segmentation(image, model_name, points)
+        except Exception as e:
+            import traceback; traceback.print_exc()
+            results[i] = {"error": "Could not prepare image: %s" % e}
+
+    live = [i for i in range(n) if prepped[i] is not None]
+    if not live:
+        return results
+
+    # --- GPU: one acquisition for the whole batch ---
+    model_filename = prepped[live[0]][2]
+    try:
+        gpu_out = run_segmentation_patched_many(
+            [prepped[i][0] for i in live], model_filename)
+    except Exception as e:
+        import traceback; traceback.print_exc()
+        for i in live:
+            results[i] = {"error": "Segmentation failed: %s" % e}
+        return results
+
+    # --- CPU: the unchanged post-processing tail, per image ---
+    for k, i in enumerate(live):
+        try:
+            out = run_segmentation(
+                None, model_name, min_px, 10000,
+                use_stereo, left_pct, top_pct, items[i][1], use_min, False,
+                _prepped=prepped[i], _gpu_result=gpu_out[k], _t_start_override=t0)
+            results[i] = _result_from_out(out, want_viability)
+        except Exception as e:
+            import traceback; traceback.print_exc()
+            results[i] = {"error": "Post-processing failed: %s" % e}
+
+    return results
+
+
 def _build_correction(res):
     """Per-cell thumbnails for manual live/dead correction.
 
@@ -2021,20 +2164,43 @@ with gr.Blocks(title="CellposeCellCounter", theme=gr.themes.Soft()) as demo:
             imgs, ptss = [i1, i2, i3, i4], [p1, p2, p3, p4]
             t0 = time.time()
             rows, results, notes = [], [], []
+
+            # Work out which squares are eligible BEFORE touching the GPU, then
+            # segment all of them in ONE acquisition. Previously this loop
+            # called _segment_one per square, so a four-square run queued for a
+            # GPU four separate times and rebuilt the Cellpose model each time.
+            skipped, eligible = {}, []
             for k, (im, pts) in enumerate(zip(imgs, ptss), start=1):
-                blank = ["Square %d" % k, "—", "—", "—", "—"]
                 if im is None:
-                    rows.append(blank); results.append(None)
                     continue
                 if not (pts and len(pts) >= 3):
                     # Deliberately skipped, not segmented whole-frame: the whole
                     # frame is a much larger area than one block, so counting it
                     # would silently corrupt the concentration.
+                    skipped[k - 1] = "Square %d: no block tapped, skipped" % k
+                    continue
+                eligible.append(k - 1)
+
+            seg = [None, None, None, None]
+            if eligible:
+                batch = _segment_many([(imgs[i], ptss[i]) for i in eligible],
+                                      HEMO_MODEL, True, 0, True, 1, 1, True)
+                for slot, r in zip(eligible, batch):
+                    seg[slot] = r
+
+            # Rebuild rows in square order so the table and the notes read 1..4.
+            for k in range(1, 5):
+                slot = k - 1
+                blank = ["Square %d" % k, "\u2014", "\u2014", "\u2014", "\u2014"]
+                if slot in skipped:
                     rows.append(blank); results.append(None)
-                    notes.append("Square %d: no block tapped, skipped" % k)
+                    notes.append(skipped[slot])
                     continue
-                r = _segment_one(im, pts, HEMO_MODEL, True, 0, True, 1, 1, True)
-                if r is None or "error" in r:
+                r = seg[slot]
+                if r is None:
+                    rows.append(blank); results.append(None)
+                    continue
+                if "error" in r:
                     rows.append(blank); results.append(None)
                     notes.append("Square %d failed" % k)
                     continue