Blame for app.py

"""Hemocytometer ground-truth annotator.

Tab "Single image": upload one photo, click the 4 corners of the 4x4 block,
segment, then correct the model by clicking.

Tab "4 squares": upload the four corner squares, set each block's corners
(click them, or Auto-detect from a click in the block centre), segment all
four with one button, then correct each square. Records the end-to-end
segmentation time and gives a copyable TSV row (model vs corrected counts and
times) for a spreadsheet.

Speed is the point of this design. After segmentation every click is handled
in the browser by gt_canvas.js -- hit-testing, drawing, undo, the min-size
filter and the count -- and corner placement by corner_canvas.js, so nothing
round-trips to the server and Gradio never replaces the image with a progress
bar. The server only decodes uploads, detects grids, warps and segments.
"""
import base64
import io
import json
import os
import time
import uuid

import cv2
import gradio as gr
import numpy as np
from huggingface_hub import hf_hub_download
from PIL import Image, ImageOps

from grid_seeded import detect_from_seed

try:                       # ZeroGPU on Spaces; a no-op decorator locally
    import spaces
    GPU = spaces.GPU(duration=30)
except Exception:          # pragma: no cover - local runs without `spaces`
    def GPU(fn):
        return fn

HF_REPO_ID = "LiangLabUMB/cellposecellcounter"
MODEL_FILE = "hemocytometer_retrained_20260825.npy"
LOCAL_MODEL_ENV = "CPCC_LOCAL_MODEL"   # same override as the main app

MAX_SIDE = 1024        # segmented square is capped here, as in the main app
PREVIEW_MAX = 2048     # corner-picking view; corners are scaled back to full res

HERE = os.path.dirname(os.path.abspath(__file__))
_model = None


def model_path():
    local = os.environ.get(LOCAL_MODEL_ENV, "").strip()
    if local and os.path.exists(local):
        return local
    return hf_hub_download(repo_id=HF_REPO_ID, filename=MODEL_FILE)


# ---------------------------------------------------------------------------
# Geometry -- identical to the main app, so a ground truth made here is
# measured on exactly the square the counter sees.
# ---------------------------------------------------------------------------
def order_quad(pts):
    """(top-left, top-right, bottom-right, bottom-left), robust to rotation."""
    pts = np.asarray(pts, dtype=np.float32)
    centre = pts.mean(axis=0)
    angles = np.arctan2(pts[:, 1] - centre[1], pts[:, 0] - centre[0])
    ordered = pts[np.argsort(angles)]
    start = np.argmin(ordered.sum(axis=1))
    return np.roll(ordered, -start, axis=0)


def warp_to_square(image_np, points):
    tl, tr, br, bl = src = order_quad(points)
    out_w = max(1, int(round(max(np.linalg.norm(br - bl), np.linalg.norm(tr - tl)))))
    out_h = max(1, int(round(max(np.linalg.norm(tr - br), np.linalg.norm(tl - bl)))))
    dst = np.array([[0, 0], [out_w - 1, 0], [out_w - 1, out_h - 1],
                    [0, out_h - 1]], dtype=np.float32)
    M = cv2.getPerspectiveTransform(src, dst)
    return cv2.warpPerspective(image_np, M, (out_w, out_h))


def safe_resize(image_np):
    h, w = image_np.shape[:2]
    if max(h, w) <= MAX_SIDE:
        return image_np
    s = MAX_SIDE / max(h, w)
    return cv2.resize(image_np, (max(1, int(w * s)), max(1, int(h * s))),
                      interpolation=cv2.INTER_AREA)


# ---------------------------------------------------------------------------
# Segmentation
# ---------------------------------------------------------------------------
@GPU
def segment_np(image_np):
    """Cellpose on one warped square. Only this function holds the GPU.

    One eval on the whole <=1024 px square: Cellpose tiles internally with
    overlap and stitches flows before computing masks, so there are no seams.
    """
    global _model
    from cellpose import models
    if _model is None:
        _model = models.CellposeModel(gpu=True, pretrained_model=model_path())
    masks, _, _ = _model.eval(image_np, diameter=None)
    return masks


def encode_labels(masks):
    """Label PNG for the browser: R,G = 16-bit id, B = distance to edge (px).

    The distance lets the canvas draw outlines of any thickness without
    recomputing anything when the zoom changes. Edges count against background
    AND against the image border, so cells cut by the crop still get a ring.
    """
    m = masks.astype(np.int32)
    if m.max() >= 65536:
        raise ValueError("more than 65535 objects")
    mp = cv2.copyMakeBorder(m.astype(np.float32), 1, 1, 1, 1,
                            cv2.BORDER_CONSTANT, value=0)
    k = np.ones((3, 3), np.uint8)
    edge = ((cv2.dilate(mp, k) != mp) | (cv2.erode(mp, k) != mp))[1:-1, 1:-1]
    dist = cv2.distanceTransform((~edge).astype(np.uint8), cv2.DIST_L2, 3)
    d = np.clip(np.round(dist) + 1, 1, 255).astype(np.uint8)
    d[m == 0] = 0
    rgb = np.dstack([(m >> 8).astype(np.uint8), (m & 255).astype(np.uint8), d])
    ok, buf = cv2.imencode(".png", rgb[..., ::-1])       # cv2 wants BGR
    return "data:image/png;base64," + base64.b64encode(buf).decode()


def jpeg_url(image_np, quality=92):
    ok, buf = cv2.imencode(".jpg", cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR),
                           [cv2.IMWRITE_JPEG_QUALITY, quality])
    return "data:image/jpeg;base64," + base64.b64encode(buf).decode()


def segment_square(full, pts, name=""):
    """Warp + segment one square -> annotation dict for the canvas."""
    square = np.ascontiguousarray(safe_resize(warp_to_square(full, pts)))
    t0 = time.time()
    masks = segment_np(square)
    secs = time.time() - t0
    return {"name": name, "img": jpeg_url(square), "lab": encode_labels(masks),
            "w": square.shape[1], "h": square.shape[0],
            "nobj": int(masks.max()), "secs": secs}


def open_upload(path):
    try:
        im = Image.open(path)
        im = ImageOps.exif_transpose(im) or im
        return np.asarray(im.convert("RGB"))
    except Exception:
        return None


def parse_corners(value):
    try:
        pts = json.loads(value) if value else None
    except ValueError:
        return None
    return pts if pts and len(pts) == 4 else None


# ---------------------------------------------------------------------------
# Tab 1: single image -- the annotation canvas also does the corner step
# ---------------------------------------------------------------------------
def on_upload(path):
    if not path:
        return None, gr.update(), "*Upload an image to start.*"
    full = open_upload(path)
    if full is None:
        return None, gr.update(), "⚠️ Could not read that file."
    h, w = full.shape[:2]
    s = min(1.0, PREVIEW_MAX / max(h, w))
    prev = cv2.resize(full, (int(w * s), int(h * s)),
                      interpolation=cv2.INTER_AREA) if s < 1 else full
    payload = json.dumps({"seq": uuid.uuid4().hex, "key": "single",
                          "mode": "corners", "img": jpeg_url(prev, 88),
                          "w": prev.shape[1], "h": prev.shape[0], "scale": s})
    return (full, payload,
            "**%s** (%d × %d) — click the 4 corners of the 4×4 block, then "
            "press **Segment**." % (os.path.basename(path), w, h))


def on_segment(corners_json, full):
    if corners_json == "CANCEL":
        return gr.update(), "*Kept your edits.*"
    if full is None:
        return gr.update(), "⚠️ Upload an image first."
    pts = parse_corners(corners_json)
    if pts is None:
        return gr.update(), "⚠️ Place all 4 corners first."
    t0 = time.time()
    try:
        sq = segment_square(full, pts)
    except Exception as e:
        return gr.update(), "⚠️ Segmentation failed: %s" % e
    payload = json.dumps(dict(sq, seq=uuid.uuid4().hex, key="single",
                              mode="annotate"))
    return (payload,
            "Segmented %d × %d px in %.1f s — %d objects detected. Correct the "
            "count by clicking below." % (sq["w"], sq["h"], time.time() - t0,
                                          sq["nobj"]))


# ---------------------------------------------------------------------------
# Tab 2: 4 squares -- corner canvases (corner_canvas.js) feed one annotation
# canvas with a square switcher
# ---------------------------------------------------------------------------
QUAD_KEYS = ["q0", "q1", "q2", "q3"]


def cp_payload(key, full):
    """Image payload for CP.load. full=None clears the canvas."""
    p = {"seq": uuid.uuid4().hex, "key": key, "mode": "image", "img": None}
    if full is not None:
        h, w = full.shape[:2]
        s = min(1.0, PREVIEW_MAX / max(h, w))
        prev = cv2.resize(full, (int(w * s), int(h * s)),
                          interpolation=cv2.INTER_AREA) if s < 1 else full
        p.update(img=jpeg_url(prev, 88), w=prev.shape[1], h=prev.shape[0],
                 scale=s)
    return json.dumps(p)


def cp_detect(key, seed_json, full):
    """Grid auto-detection from a seed click. CPU only, NOT behind the GPU."""
    out = {"seq": uuid.uuid4().hex, "key": key, "mode": "corners",
           "corners": [], "msg": ""}
    try:
        seed = json.loads(seed_json) if seed_json else None
    except ValueError:
        seed = None
    if full is None or not seed:
        out["msg"] = "Upload an image first."
        return json.dumps(out)
    res = detect_from_seed(full, (float(seed[0]), float(seed[1])))
    if res.get("ok"):
        qc = res.get("qc") or {}
        out["corners"] = [[float(x), float(y)] for x, y in res["corners"]]
        out["msg"] = ("✓ Auto-detected — confidence %.0f%%, edges %.0f px "
                      "off the rulings. Drag a handle to adjust."
                      % (res["confidence"] * 100, qc.get("line_offset_mean_px", 0)))
    else:
        out["msg"] = ("⚠ Auto-detect failed (%s) — click the 4 corners "
                      "instead." % res.get("reason", "no fit"))
    return json.dumps(out)


def on_upload4(files):
    """Sorted by filename into Square 1..4. Returns fulls, names, payloads, msg."""
    fulls, names = [None] * 4, [""] * 4
    if files:
        paths = sorted((str(getattr(f, "name", f)) for f in files),
                       key=lambda p: os.path.basename(p).lower())
        bad, k = [], 0
        for p in paths:
            if k >= 4:
                break
            im = open_upload(p)
            if im is None:
                bad.append(os.path.basename(p))
                continue
            fulls[k], names[k] = im, os.path.basename(p)
            k += 1
        msg = "  ·  ".join("Square %d ← %s" % (i + 1, n)
                           for i, n in enumerate(names) if n)
        msg = ("**%s**" % msg) if k else "⚠️ *None of those files could be read.*"
        if bad:
            msg += "  \n⚠️ *Unreadable, skipped: %s*" % ", ".join(bad)
        if len(paths) - k - len(bad) > 0:
            msg += "  \n⚠️ *Only 4 squares — extra files ignored.*"
        if k:
            msg += ("  \n*Set each block: click its 4 corners, or press ⌖ Auto "
                    "and click the block centre. Then **Segment all 4**.*")
    else:
        msg = "*Upload the four square images.*"
    return (*fulls, names, *[cp_payload(k, f) for k, f in zip(QUAD_KEYS, fulls)],
            msg)


def on_segment4(c0, c1, c2, c3, f0, f1, f2, f3, names):
    """Four separate GPU calls, one per square, like the main app."""
    if c0 == "CANCEL":
        return gr.update(), "*Kept your edits.*"
    t0 = time.time()
    squares, notes = [], []
    names = names or [""] * 4
    for k, (c, f) in enumerate(zip((c0, c1, c2, c3), (f0, f1, f2, f3))):
        pts = parse_corners(c)
        if f is None:
            squares.append(None)
            notes.append("Square %d: no image" % (k + 1))
            continue
        if pts is None:
            squares.append(None)
            notes.append("Square %d: corners not set" % (k + 1))
            continue
        try:
            squares.append(segment_square(f, pts, names[k]))
        except Exception as e:
            squares.append(None)
            notes.append("Square %d failed: %s" % (k + 1, e))
    server = time.time() - t0
    n = sum(1 for s in squares if s)
    payload = json.dumps({"seq": uuid.uuid4().hex, "key": "quad",
                          "mode": "annotate4", "squares": squares,
                          "server_secs": server, "note": "; ".join(notes)})
    if not n:     # still sent, so the canvas resets its timer display
        return payload, "⚠️ Nothing segmented — " + "; ".join(notes)
    msg = "Segmented %d of 4 squares (server %.1f s)." % (n, server)
    if notes:
        msg += "  \n⚠️ " + "; ".join(notes)
    return payload, msg


# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------
def _read(name):
    with open(os.path.join(HERE, name), encoding="utf-8") as f:
        return f.read()


HEAD = ("<script>\n" + _read("gt_canvas.js") + "\n</script>\n"
        "<script>\n" + _read("corner_canvas.js") + "\n</script>")

CSS = """
.gt-hidden, .cp-hidden { display: none !important; }
.gt-root { outline: none; display: flex; flex-direction: column; gap: 6px; }
.gt-bar { display: flex; flex-wrap: wrap; align-items: center; gap: 10px 16px; }
.gt-count { font-size: 34px; font-weight: 700; line-height: 1;
            font-variant-numeric: tabular-nums; min-width: 2.5em; }
.gt-count small { font-size: 13px; font-weight: 400; margin-left: 6px;
                  opacity: .7; }
.gt-breakdown { font-size: 13px; opacity: .85; flex: 1 1 260px; }
.gt-r { color: #e5484d; } .gt-a { color: #c9a400; }
.gt-minl { display: flex; align-items: center; gap: 6px; font-size: 13px; }
.gt-minl input[type=number] { width: 76px; }
.gt-bar button, .gt-sq button, .gt-copy button {
    padding: 4px 10px; border-radius: 6px; font-size: 13px;
    border: 1px solid #8884; background: transparent; color: inherit;
    cursor: pointer; }
.gt-bar button:disabled, .gt-sq button:disabled { opacity: .4; cursor: default; }
.gt-sq { display: flex; gap: 6px; flex-wrap: wrap; }
.gt-sq:empty { display: none; }
.gt-sq button.gt-cur { background: #4aaaff; color: #fff; border-color: #4aaaff; }
.gt-stage { position: relative; height: 74vh; min-height: 420px;
            border-radius: 8px; overflow: hidden; background: #1b1d22; }
.gt-canvas { width: 100%; height: 100%; display: block; cursor: crosshair;
             touch-action: none; }
.gt-empty { position: absolute; inset: 0; display: flex; align-items: center;
            justify-content: center; color: #aaa; pointer-events: none; }
.gt-hint { font-size: 12px; opacity: .75; min-height: 1.3em; }
.gt-results { font-size: 13px; }
.gt-results table { border-collapse: collapse; margin: 6px 0; }
.gt-results th, .gt-results td { padding: 3px 10px; border-bottom: 1px solid #8883;
                                 text-align: right; }
.gt-results td.gt-name { text-align: left; max-width: 280px; overflow: hidden;
                         text-overflow: ellipsis; white-space: nowrap; }
.gt-results tr.gt-cur td { background: #4aaaff22; }
.gt-results tr.gt-mean td { font-weight: 600; }
.gt-copy { display: flex; gap: 6px; align-items: center; }
.gt-copy input { flex: 1; font-family: monospace; font-size: 12px;
                 padding: 4px 6px; }
/* corner canvases */
.cp-root { display: flex; flex-direction: column; gap: 4px; }
.cp-bar { display: flex; align-items: center; gap: 8px; flex-wrap: wrap;
          font-size: 13px; }
.cp-title { font-weight: 600; }
.cp-info { flex: 1 1 160px; opacity: .8; }
.cp-info.cp-ok { color: #2e9e5b; opacity: 1; }
.cp-bar button { padding: 2px 9px; border-radius: 6px; font-size: 12px;
                 border: 1px solid #8884; background: transparent;
                 color: inherit; cursor: pointer; }
.cp-bar button:disabled { opacity: .4; cursor: default; }
.cp-bar button.cp-on { background: #ffcc00; color: #000; border-color: #ffcc00; }
.cp-stage { position: relative; border-radius: 8px; overflow: hidden;
            background: #1b1d22; }
.cp-stage canvas { width: 100%; height: 100%; display: block;
                   cursor: crosshair; touch-action: none; }
.cp-empty { position: absolute; inset: 0; display: flex; align-items: center;
            justify-content: center; color: #aaa; pointer-events: none;
            font-size: 14px; }
"""


def gt_html(key, multi):
    return """
<div class="gt-root" data-key="{key}" data-multi="{multi}" tabindex="0">
  <div class="gt-bar">
    <div class="gt-count"><span class="gt-total">&mdash;</span><small>cells</small></div>
    <div class="gt-breakdown"></div>
    <label class="gt-minl">Min size
      <input type="range" class="gt-min" min="0" max="500" step="1" value="0">
      <input type="number" class="gt-minnum" min="0" step="1" value="0"> px
    </label>
    <button class="gt-undo" title="Ctrl+Z">&#8630; Undo</button>
    <button class="gt-redo" title="Ctrl+Shift+Z">&#8631; Redo</button>
    <button class="gt-fit" title="F">Fit</button>
    <button class="gt-resetc">Reset corners</button>
    <button class="gt-editc">Edit corners</button>
  </div>
  <div class="gt-sq"></div>
  <div class="gt-stage">
    <canvas class="gt-canvas"></canvas>
    <div class="gt-empty">{empty}</div>
  </div>
  <div class="gt-hint"></div>
  <div class="gt-results"></div>
</div>
""".format(key=key, multi="1" if multi else "0",
           empty="Segment the 4 squares above to start" if multi
           else "Upload an image to start")


CP_HTML = """
<div class="cp-root" data-key="{key}">
  <div class="cp-bar">
    <span class="cp-title">{title}</span><span class="cp-info"></span>
    <button class="cp-auto" title="Auto-detect the 4x4 block: press, then click the block centre">&#8982; Auto</button>
    <button class="cp-reset">Reset</button>
    <button class="cp-fit">Fit</button>
  </div>
  <div class="cp-stage" style="height:{height}px">
    <canvas></canvas><div class="cp-empty">{title} — upload above</div>
  </div>
</div>
"""

LOAD_JS = "(...ps) => { for (const p of ps) if (window.GT) window.GT.load(p); }"
CP_LOAD_JS = "(...ps) => { for (const p of ps) if (window.CP) window.CP.load(p); }"

LEGEND = (
    "<small>Counted cells: coloured outline · removed: red with ✕ · below the "
    "min size: grey · touching the left/bottom 1% exclusion zones (red "
    "bands): faint white, not counted · added dots: yellow. Min size is an "
    "area in pixels of the segmented square (max 1024 px side), the same "
    "units as the main counter.</small>")

with gr.Blocks(title="Hemocytometer Ground Truth", head=HEAD, css=CSS,
               theme=gr.themes.Soft()) as demo:
    gr.Markdown("## Hemocytometer ground truth")

    # ================================================================ tab 1
    with gr.Tab("Single image"):
        with gr.Row(equal_height=True):
            upload = gr.UploadButton("📂 Upload image", file_types=["image"],
                                     variant="secondary", scale=0)
            seg_btn = gr.Button("🔬 Segment", variant="primary", scale=0)
            status = gr.Markdown("*Upload an image to start.*")
        gr.HTML(gt_html("single", False))
        gr.Markdown(LEGEND)

        full_state = gr.State(None)
        # Hidden with CSS rather than visible=False so their values are always
        # present in the page for the js steps.
        with gr.Column(elem_classes="gt-hidden"):
            payload = gr.Textbox(show_label=False, container=False, max_lines=1)
            corners = gr.Textbox(show_label=False, container=False, max_lines=1)

        upload.upload(fn=on_upload, inputs=[upload],
                      outputs=[full_state, payload, status],
                      show_progress="minimal"
                      ).then(fn=None, inputs=[payload], js=LOAD_JS)

        # The js runs first in the browser and its return value replaces the
        # inputs sent to on_segment: the corners come straight from the canvas.
        seg_btn.click(fn=on_segment, inputs=[corners, full_state],
                      outputs=[payload, status], show_progress="minimal",
                      js="(c, s) => [window.GT ? window.GT.getCorners('single') : '', s]"
                      ).then(fn=None, inputs=[payload], js=LOAD_JS)

    # ================================================================ tab 2
    with gr.Tab("4 squares"):
        gr.Markdown(
            "**1 ·** Upload the four square images (sorted by filename into "
            "Square 1–4). **2 ·** Set each block: click its 4 corners, or press "
            "**⌖ Auto** and click the block centre; drag any handle to adjust. "
            "**3 ·** Press **Segment all 4** — the time from that click until "
            "all four squares are ready is recorded. **4 ·** Correct each square "
            "below (buttons or keys 1–4 switch squares), then copy the results "
            "row into your spreadsheet.")
        up4 = gr.Files(label="Upload the four square images",
                       file_count="multiple", file_types=["image"], height=95)
        status4 = gr.Markdown("*Upload the four square images.*")

        fulls4 = [gr.State(None) for _ in range(4)]
        names4 = gr.State([""] * 4)
        cp_payloads, cp_corners = [], []
        for _pair in range(2):
            with gr.Row():
                for _k in (2 * _pair, 2 * _pair + 1):
                    with gr.Column():
                        key = QUAD_KEYS[_k]
                        gr.HTML(CP_HTML.format(key=key, title="Square %d" % (_k + 1),
                                               height=380))
                        with gr.Column(elem_classes="cp-hidden"):
                            pl = gr.Textbox(show_label=False, container=False,
                                            max_lines=1)
                            cr = gr.Textbox(show_label=False, container=False,
                                            max_lines=1)
                            sd = gr.Textbox(show_label=False, container=False,
                                            max_lines=1)
                            det = gr.Button("detect", elem_id="cp-detect-%s" % key)
                        cp_payloads.append(pl)
                        cp_corners.append(cr)
                        det.click(fn=lambda s, f, _key=key: cp_detect(_key, s, f),
                                  inputs=[sd, fulls4[_k]], outputs=[pl],
                                  js="(s, f) => [window.CP.seed('%s'), f]" % key,
                                  show_progress="hidden"
                                  ).then(fn=None, inputs=[pl], js=CP_LOAD_JS)

        seg4_btn = gr.Button("🔬 Segment all 4", variant="primary", size="lg")
        gr.HTML(gt_html("quad", True))
        gr.Markdown(LEGEND)
        with gr.Column(elem_classes="gt-hidden"):
            payload4 = gr.Textbox(show_label=False, container=False, max_lines=1)

        up4.change(fn=on_upload4, inputs=[up4],
                   outputs=[*fulls4, names4, *cp_payloads, status4],
                   show_progress="minimal"
                   ).then(fn=None, inputs=cp_payloads, js=CP_LOAD_JS)

        # js pre-step: start the end-to-end timer, then read the 4 corner sets
        # from the canvases. A cancelled re-segment sends "CANCEL".
        seg4_btn.click(
            fn=on_segment4,
            inputs=[*cp_corners, *fulls4, names4],
            outputs=[payload4, status4], show_progress="minimal",
            js="(...a) => { if (window.GT && !window.GT.markStart('quad')) {"
               " a[0] = 'CANCEL'; return a; }"
               " for (let i = 0; i < 4; i++) a[i] = window.CP ? "
               "window.CP.get('q' + i) : ''; return a; }"
        ).then(fn=None, inputs=[payload4], js=LOAD_JS)


if __name__ == "__main__":
    try:
        model_path()             # warm the weights cache before the first run
    except Exception as e:
        print("Could not prefetch weights:", e)
    demo.launch()