Deploy CellQuant

README.md

@@ -1,17 +1,18 @@
 ---
-title: CellQuant
+title: AI Cell Quantification
 emoji: "⬜"
 colorFrom: gray
 colorTo: gray
 sdk: gradio
 sdk_version: 5.49.1
-python_version: "3.11"
+python_version: "3.12"
+suggested_hardware: zero-a10g
 app_file: app.py
 pinned: false
 license: mit
 ---
 
-# CellQuant: nucleus, antibody and puncta counting for fluorescence sections
+# AI Cell Quantification (CellQuant): nucleus, antibody and puncta counting for fluorescence sections
 
 CellQuant counts the six categories that were annotated in CVAT on 35-day spinal-cord
 sections stained with DAPI (blue), an antibody (green: Iba1 or TMEM119) and a red punctum
@@ -131,9 +132,25 @@ huggingface-cli login
 huggingface-cli upload <your-user>/cellquant . . --repo-type space --exclude "runs/*" --exclude "tests/*" --exclude ".pytest_cache/*"
 ```
 
-Create the Space first (SDK: Gradio, hardware: free CPU) or add `--create`. Add the optional
+Create the Space first (SDK: Gradio) or add `--create`. Add the optional
 `GEMINI_API_KEY` secret under *Settings → Variables and secrets*.
 
+The Space runs on **ZeroGPU** (a shared NVIDIA GPU that is attached only while a call runs):
+`app.py` imports `spaces` before torch, puts the model on `cuda` at start-up and wraps the two
+inference functions in `@spaces.GPU` with a duration estimated from the number of images and the
+TTA setting. Visitors need no account: an unauthenticated visitor has 2 GPU minutes a day (about
+60 full images at 8× TTA), a free account 5, a PRO account 40. Hosting a ZeroGPU Space needs a
+Hugging Face account in good standing (two Spaces free; PRO accounts up to ten), Python 3.12 and
+torch ≥ 2.8 (`README.md` front matter and `requirements.txt`). The deploy script does the whole
+round trip and times a full image on the Space:
+
+```bash
+python scripts/deploy_hf.py --repo <your-user>/cellquant --hardware zero-a10g --bench
+```
+
+Without `spaces` installed, or on CPU hardware, the same code runs unchanged on the CPU
+(`CELLQUANT_DEVICE` overrides the device).
+
 ## Running locally
 
 ```bash

app.py

@@ -17,8 +17,17 @@ import zipfile
 
 os.environ.setdefault("GRADIO_ANALYTICS_ENABLED", "False")
 
+# ZeroGPU (Hugging Face Spaces): `spaces` must be imported before torch; its GPU decorator
+# borrows a GPU for the duration of a call and is a no-op everywhere else.
+try:
+    import spaces  # noqa: E402
+except ImportError:  # local runs and tests without the package
+    spaces = None
+ZERO_GPU = spaces is not None and str(os.environ.get("SPACES_ZERO_GPU", "")).lower() in ("1", "true")
+
 import gradio as gr  # noqa: E402
 import numpy as np  # noqa: E402
+import time  # noqa: E402
 from PIL import Image  # noqa: E402
 
 sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))  # run from any cwd / python -I
@@ -29,28 +38,56 @@ from cellquant.gemini_summary import enabled as gemini_enabled
 from cellquant.gemini_summary import summarize as gemini_summarize
 from cellquant.infer import Predictor
 from cellquant.labels import ALL_LABELS, LABEL_COLORS
+from cellquant.model import load_checkpoint
 from cellquant.postprocess import PostConfig, analyze, render_overlay, result_to_json
 from cellquant.validation import COMPARISON_HEADERS, COMPARISON_NOTE, DIFF_HEADERS, ValidationSet, compare, write_bundle
 
 WEIGHTS = os.environ.get("CELLQUANT_WEIGHTS", os.path.join(os.path.dirname(os.path.abspath(__file__)), "weights", "cellquant.pt"))
-DEVICE = os.environ.get("CELLQUANT_DEVICE", "cpu")
+# On ZeroGPU the model must sit on cuda from start-up (CUDA is emulated outside GPU calls); elsewhere CPU unless told otherwise.
+DEVICE = os.environ.get("CELLQUANT_DEVICE", "cuda" if ZERO_GPU else "cpu")
+TTA_OF = {"none (fastest)": 1, "4x flips": 4, "8x rotations+flips": 8}
 
 EX_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "examples")
 VAL_DIR = os.environ.get("CELLQUANT_VALIDATION_DIR", os.path.join(os.path.dirname(os.path.abspath(__file__)), "validation"))
 VALIDATION = ValidationSet(VAL_DIR)
 
+MODEL = None
+if os.path.exists(WEIGHTS):
+    MODEL, _META = load_checkpoint(WEIGHTS)
+    MODEL.to(DEVICE).eval()
 _predictors: dict[int, Predictor] = {}
 _val_probs: dict[tuple[str, int], dict] = {}   # (image key, tta) -> network probabilities, a few entries
 
 
 def get_predictor(tta: int) -> Predictor:
     if tta not in _predictors:
-        if not os.path.exists(WEIGHTS):
+        if MODEL is None:
             raise gr.Error(f"Model weights not found at {WEIGHTS}. Train with train.py or set CELLQUANT_WEIGHTS.")
-        _predictors[tta] = Predictor(WEIGHTS, device=DEVICE, tta=tta)
+        _predictors[tta] = Predictor(MODEL, device=DEVICE, tta=tta)
     return _predictors[tta]
 
 
+def gpu(duration):
+    """`@spaces.GPU(duration=...)` on ZeroGPU, identity elsewhere. ``duration`` may be a
+    callable of the decorated function's arguments (seconds of GPU time to reserve;
+    shorter reservations get a better place in the queue)."""
+    if spaces is None:
+        return lambda f: f
+    return spaces.GPU(duration=duration)
+
+
+def _analyze_duration(files, tta_choice, *args, **kwargs) -> int:
+    n = max(1, len(files or []))
+    per_image = {1: 4, 4: 6, 8: 10}.get(TTA_OF.get(tta_choice, 1), 10)
+    return int(min(600, 15 + per_image * n))
+
+
+def _device_name() -> str:
+    if DEVICE.startswith("cuda"):
+        return "GPU (ZeroGPU)" if ZERO_GPU else "GPU"
+    return "Apple GPU" if DEVICE == "mps" else "CPU"
+
+
 def _legend_html() -> str:
     items = "".join(
         f'<span style="display:inline-block;margin:2px 10px 2px 0"><span style="display:inline-block;width:14px;height:14px;'
@@ -60,11 +97,13 @@ def _legend_html() -> str:
     return f"<div>{items}</div>"
 
 
+@gpu(duration=_analyze_duration)
 def analyze_files(files, tta_choice, antibody_thr, cellnos_thr, min_cell_area, nostril_rule, cellnos_rule,
                   want_summary, progress=gr.Progress()):
     if not files:
         raise gr.Error("Upload at least one image.")
-    tta = {"none (fastest)": 1, "4x flips": 4, "8x rotations+flips": 8}[tta_choice]
+    t_start = time.time()
+    tta = TTA_OF[tta_choice]
     cfg = PostConfig(antibody_thr=float(antibody_thr), cellnos_thr=float(cellnos_thr),
                      min_cell_area=int(min_cell_area), nostril_rule=nostril_rule, cellnos_rule=cellnos_rule)
     pred = get_predictor(tta)
@@ -121,7 +160,9 @@ def analyze_files(files, tta_choice, antibody_thr, cellnos_thr, min_cell_area, n
             zf.write(p, os.path.basename(p))
         for p, _ in gallery:
             zf.write(p, os.path.basename(p))
-    summary = "\n\n".join(f"Warning: {w}" for w in warnings_)
+    summary = (f"Analysed {len(rows)} image{'s' if len(rows) != 1 else ''} in {time.time() - t_start:.1f} s "
+               f"on the {_device_name()} ({tta_choice}).")
+    summary += "".join(f"\n\nWarning: {w}" for w in warnings_)
     if want_summary:
         if gemini_enabled():
             first = next(iter(overlays.values()))
@@ -132,6 +173,7 @@ def analyze_files(files, tta_choice, antibody_thr, cellnos_thr, min_cell_area, n
     return rows, gallery, zip_path, summary
 
 
+@gpu(duration=lambda key, tta_choice="none (fastest)", *a, **k: 15 + {1: 4, 4: 6, 8: 10}.get(TTA_OF.get(tta_choice, 1), 10))
 def validate_image(key, tta_choice="none (fastest)", progress=gr.Progress()):
     """Run the model on one bundled annotated image and compare it with the annotation.
 
@@ -139,7 +181,8 @@ def validate_image(key, tta_choice="none (fastest)", progress=gr.Progress()):
     their captions), the per-label table, the per-object table and the download."""
     if not key or key not in VALIDATION.keys:
         raise gr.Error("Pick one of the annotated images.")
-    tta = {"none (fastest)": 1, "4x flips": 4, "8x rotations+flips": 8}[tta_choice]
+    t_start = time.time()
+    tta = TTA_OF[tta_choice]
     img = VALIDATION.image(key)
     R = VALIDATION.raster(key)
     progress(0.1, desc=f"running the model on {key}")
@@ -162,7 +205,8 @@ def validate_image(key, tta_choice="none (fastest)", progress=gr.Progress()):
     missed = gr.Image(value=os.path.join(tmpdir, f"{key}_missed_only.png"),
                       label=f"3. Missed only: {n_missed_nuc} nuclei and {n_missed_pun} puncta annotated but not predicted "
                             f"(red, numbered as in the list below)")
-    return cmp.summary, manual, predicted, missed, cmp.rows, cmp.diffs, zip_path
+    summary = cmp.summary + f"\n* **Run time:** {time.time() - t_start:.1f} s on the {_device_name()} ({tta_choice})."
+    return summary, manual, predicted, missed, cmp.rows, cmp.diffs, zip_path
 
 
 def pick_validation_image(tta_choice, evt: gr.SelectData):
@@ -197,9 +241,10 @@ def build_demo() -> gr.Blocks:
     tta_default = os.environ.get("CELLQUANT_TTA_DEFAULT", "8x rotations+flips")
     if tta_default not in tta_choices:
         tta_default = tta_choices[-1]
-    with gr.Blocks(title="CellQuant", theme=theme, css=CSS, js=FORCE_LIGHT_JS) as demo:
+    with gr.Blocks(title="AI Cell Quantification", theme=theme, css=CSS, js=FORCE_LIGHT_JS) as demo:
         gr.Markdown(
-            f"# CellQuant {__version__}\n"
+            f"# AI Cell Quantification\n"
+            f"CellQuant {__version__}. "
             "Counts nuclei (antibody positive or negative), nuclei containing a red punctum (\"Nostril\") and puncta "
             "inside or outside nuclei in RGB fluorescence images (blue = DAPI, green = antibody, red = puncta). "
             "Trained on CVAT outlines of 35-day spinal-cord sections stained for Iba1 or TMEM119."
@@ -264,13 +309,15 @@ def build_demo() -> gr.Blocks:
             "Nucleus outlines come from a watershed over the network's distance and class maps. "
             "A nucleus is labelled Nostril when a punctum classified as in-nucleus lies inside it. "
             "Images larger than 1536 px are processed in overlapping tiles. Nothing leaves the Space except, "
-            "optionally, the counts and an overlay thumbnail sent to Gemini when the summary is requested."
+            "optionally, the counts and an overlay thumbnail sent to Gemini when the summary is requested. "
+            + ("Inference runs on a shared GPU (Hugging Face ZeroGPU): no account is needed, each visitor has a daily "
+               "allowance of GPU minutes and one image takes a few seconds." if ZERO_GPU else f"Inference runs on the {_device_name()}.")
         )
     return demo
 
 
 if __name__ == "__main__":
-    if not os.path.exists(WEIGHTS):
+    if MODEL is None:
         raise SystemExit(f"CellQuant weights not found at {WEIGHTS}: run train.py --out weights/cellquant.pt or set CELLQUANT_WEIGHTS")
     build_demo().launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)),
                         allowed_paths=[EX_DIR], show_error=True)

cellquant/__init__.py

@@ -3,4 +3,4 @@
 Trained on CVAT polygon annotations of 1024x1024 RGB images (blue = DAPI nuclei,
 green = antibody stain such as Iba1 or Tmem119, red = puncta).
 """
-__version__ = "1.2.1"
+__version__ = "1.2.2"

requirements.txt

@@ -1,6 +1,7 @@
---extra-index-url https://download.pytorch.org/whl/cpu
-torch>=2.2
-torchvision>=0.17
+# CUDA build of torch from PyPI (ZeroGPU needs torch >= 2.8); locally any build works.
+torch>=2.8
+torchvision>=0.23
+spaces>=0.40
 timm>=0.9.7
 numpy>=1.26
 scipy>=1.11

scripts/deploy_hf.py

@@ -29,6 +29,8 @@ def main(argv=None):
     ap.add_argument("--public", action="store_true")
     ap.add_argument("--no-wait", action="store_true")
     ap.add_argument("--timeout", type=int, default=1800)
+    ap.add_argument("--hardware", default="", help="Space hardware to request after the upload, e.g. zero-a10g (ZeroGPU) or cpu-basic")
+    ap.add_argument("--bench", action="store_true", help="after the smoke test, time a full validation image with no TTA and with 8x TTA")
     a = ap.parse_args(argv)
     api = HfApi()
     me = whoami()["name"]
@@ -46,6 +48,13 @@ def main(argv=None):
     api.upload_folder(folder_path=ROOT, repo_id=repo, repo_type="space", ignore_patterns=IGNORE,
                       commit_message="Deploy CellQuant")
     print("uploaded")
+    if a.hardware:
+        rt = api.get_space_runtime(repo)
+        if str(rt.hardware) != a.hardware or (rt.requested_hardware and str(rt.requested_hardware) != a.hardware):
+            api.request_space_hardware(repo, a.hardware)
+            print(f"requested hardware {a.hardware} (was {rt.hardware})")
+        else:
+            print(f"hardware already {a.hardware}")
     if a.no_wait:
         return
     t0 = time.time()
@@ -83,6 +92,19 @@ def main(argv=None):
     t = time.time()
     out = client.predict([handle_file(served)], "none (fastest)", 0.5, 0.5, 50, "head", "contains", False, api_name="/analyze_files")
     print(f"API call ok in {time.time()-t:.1f}s; counts row: {out[0]['data'][0] if isinstance(out[0], dict) else out[0]}")
+    print(f"  app reports: {str(out[3]).splitlines()[0] if out[3] else '(no summary line)'}")
+    if a.bench:
+        val_dir = os.path.join(ROOT, "validation")
+        pngs = sorted(f for f in os.listdir(val_dir) if f.lower().endswith(".png")) if os.path.isdir(val_dir) else []
+        if not pngs:
+            print("bench: no validation image to time")
+            return
+        full = os.path.join(val_dir, pngs[0])
+        for tta in ("none (fastest)", "8x rotations+flips"):
+            t = time.time()
+            out = client.predict([handle_file(full)], tta, 0.5, 0.5, 50, "head", "contains", False, api_name="/analyze_files")
+            row = out[0]["data"][0] if isinstance(out[0], dict) else out[0]
+            print(f"bench {os.path.basename(full)} {tta}: round trip {time.time()-t:.1f}s; {str(out[3]).splitlines()[0] if out[3] else ''}; nuclei {sum(row[1:5])}")
 
 
 if __name__ == "__main__":

tests/test_app.py

@@ -20,7 +20,8 @@ def test_analyze_files_end_to_end(tmp_path):
     rows, gallery, zip_path, summary = app.analyze_files([str(p1), str(p2)], "none (fastest)", 0.5, 0.5, 50, "head", "contains", False)
     assert len(rows) == 2 and rows[0][0] == "a.png" and len(rows[0]) == 7
     assert len(gallery) == 2 and os.path.exists(gallery[0][0])
-    assert os.path.exists(zip_path) and summary == ""
+    assert os.path.exists(zip_path) and summary.startswith("Analysed 2 images in ") and "on the CPU" in summary
+    assert app._analyze_duration(["a", "b", "c"], "8x rotations+flips") == 45 and app._analyze_duration([], "none (fastest)") == 19
     import zipfile
     with zipfile.ZipFile(zip_path) as zf:
         names = zf.namelist()