Spaces:
Running on Zero
Running on Zero
| import numpy as np | |
| from cellquant.labels import ALL_LABELS | |
| from cellquant.postprocess import PostConfig, analyze, logits_to_probs, render_overlay, segment_cells | |
| def _disc(h, w, cy, cx, r): | |
| yy, xx = np.mgrid[:h, :w] | |
| return (yy - cy) ** 2 + (xx - cx) ** 2 <= r * r | |
| def _probs(h=96, w=96): | |
| """Two touching discs; the left one antibody+ with a punctum inside; one external punctum.""" | |
| d1, d2 = _disc(h, w, 48, 30, 14), _disc(h, w, 48, 56, 14) | |
| fg = d1 | d2 | |
| inner = _disc(h, w, 48, 30, 11) | _disc(h, w, 48, 56, 11) | |
| cell = np.zeros((3, h, w), np.float32) | |
| cell[1][inner] = 0.9 | |
| cell[2][fg & ~inner] = 0.9 | |
| cell[0] = 1 - cell[1] - cell[2] | |
| antibody = np.where(d1, 0.95, 0.05).astype(np.float32) | |
| cellnos = np.full((h, w), 0.1, np.float32) | |
| nos = np.zeros((3, h, w), np.float32) | |
| p_in = _disc(h, w, 48, 30, 2) | |
| p_ext = _disc(h, w, 10, 80, 2) | |
| nos[1][p_in] = 0.9 | |
| nos[2][p_ext] = 0.9 | |
| nos[0] = 1 - nos[1] - nos[2] | |
| return {"cell": cell, "antibody": antibody, "cellnos": cellnos, "nos": nos} | |
| def test_touching_discs_split_into_two_instances(): | |
| inst = segment_cells(_probs()["cell"], PostConfig()) | |
| assert inst.max() == 2 | |
| areas = np.bincount(inst.ravel())[1:] | |
| assert all(500 <= a <= 700 for a in areas) # pi*14^2 ~ 616 | |
| def test_analyze_labels_and_counts(): | |
| res = analyze(_probs(), PostConfig()) | |
| assert len(res.cells) == 2 and len(res.nostrils) == 2 | |
| labels = sorted(c["label"] for c in res.cells) | |
| assert labels == ["Cell, Antibody -", "Cell, Nostril, Antibody +"] | |
| assert sorted(n["label"] for n in res.nostrils) == ["Nostril External of Nucleus", "Nostril in Nucleus"] | |
| assert sum(res.counts.values()) == 4 and set(res.counts) == set(ALL_LABELS) | |
| inside = [n for n in res.nostrils if n["label"] == "Nostril in Nucleus"][0] | |
| assert inside["cell_id"] == [c for c in res.cells if c["label"].startswith("Cell, Nostril")][0]["id"] | |
| def test_cellnos_rules(): | |
| p = _probs() | |
| head_only = analyze(p, PostConfig(cellnos_rule="head")) | |
| assert all(not c["label"].startswith("Cell, Nostril") for c in head_only.cells) | |
| contains = analyze(p, PostConfig(cellnos_rule="contains")) | |
| assert sum(c["label"].startswith("Cell, Nostril") for c in contains.cells) == 1 | |
| def test_geometric_nostril_rule_overrides_head(): | |
| p = _probs() | |
| # make the network call the inside punctum "external"; geometric rule should still say inside | |
| p["nos"][1], p["nos"][2] = p["nos"][2].copy(), p["nos"][1].copy() | |
| head = analyze(p, PostConfig(nostril_rule="head")) | |
| geo = analyze(p, PostConfig(nostril_rule="geometric")) | |
| assert head.counts["Nostril in Nucleus"] == 1 and geo.counts["Nostril in Nucleus"] == 1 | |
| # with head rule the punctum inside the disc is now labelled external | |
| inside_disc = [n for n in head.nostrils if n["cell_id"] > 0][0] | |
| assert inside_disc["label"] == "Nostril External of Nucleus" | |
| assert [n for n in geo.nostrils if n["cell_id"] > 0][0]["label"] == "Nostril in Nucleus" | |
| def test_min_area_filters_small_instances(): | |
| p = _probs() | |
| assert len(analyze(p, PostConfig(min_cell_area=1000)).cells) == 0 | |
| def test_logits_to_probs_and_overlay(): | |
| logits = np.random.default_rng(0).normal(size=(8, 40, 40)).astype(np.float32) | |
| pr = logits_to_probs(logits) | |
| assert np.allclose(pr["cell"].sum(0), 1, atol=1e-5) and np.allclose(pr["nos"].sum(0), 1, atol=1e-5) | |
| assert pr["antibody"].min() >= 0 and pr["antibody"].max() <= 1 | |
| res = analyze(_probs(), PostConfig()) | |
| img = np.zeros((96, 96, 3), np.uint8) | |
| ov = render_overlay(img, res) | |
| assert ov.shape == img.shape and ov.any() | |
| def test_empty_image_gives_no_objects(): | |
| h = w = 64 | |
| p = {"cell": np.stack([np.ones((h, w)), np.zeros((h, w)), np.zeros((h, w))]).astype(np.float32), | |
| "antibody": np.zeros((h, w), np.float32), "cellnos": np.zeros((h, w), np.float32), | |
| "nos": np.stack([np.ones((h, w)), np.zeros((h, w)), np.zeros((h, w))]).astype(np.float32)} | |
| res = analyze(p, PostConfig()) | |
| assert res.cells == [] and res.nostrils == [] and sum(res.counts.values()) == 0 | |