feat accept correct reject

README.md

@@ -27,7 +27,7 @@ Its `"table"` AnnData must contain:
 | obs column | content |
 | --- | --- |
 | `cell_type` (configurable) | cell-type label for each cell |
-| `spatial_group` (configurable) | binary spatial region label (e.g. `0` = outside tumour, `1` = inside tumour) |
+| `spatial_group` (configurable) | spatial region label with two values to compare (e.g. `0`/`1`, or `"out_of_tumor"`/`"in_tumor"`). Which value is the target is chosen in Step 3 with `--spatial-group-target` / `--spatial-group-reference`, defaulting to `1` / `0` |
 | `center_x`, `center_y` | cell centroid in microns |
 
 The zarr must also expose the following SpatialData elements, used to render the per-cell annotation panels (Step 1):
@@ -90,6 +90,12 @@ python scripts/export.py \
 --layer counts
 ```
 
+`--annotation-mode` selects the actions offered in Step 2, and defaults to
+`accept_correct_reject`. Use `--annotation-mode accept_reject` to drop the relabelling
+option. The value is saved to `config.json`; because the cell-type vocabulary is written
+there too (always, whatever the mode), you can switch modes afterwards by editing
+`config.json`, without re-exporting the panels.
+
 `--target-proportion` controls the fraction of cells of interest in the subsample. Cells of interest are upweighted accordingly (importance sampling); the unnormalized weight for each sampled cell is stored in `metadata.csv` for downstream use.
 
 `--layer` selects which expression matrix to read: the named `.layers` entry holding the raw counts (e.g. `counts`), or `.X` when omitted. The value is saved to `config.json` and reused throughout the workflow — the same layer feeds the top-gene panels here in Step 1 and the CSDE model in Step 3, so set it once at export time. **It must point at raw counts**, since the noise model (Poisson / negative binomial) assumes integer counts; pointing it at normalised or log-transformed values will produce invalid results.
@@ -101,9 +107,9 @@ The script writes:
 ├── images/
 │   ├── cell_<id>.png   # one panel per cell
 │   └── ...
-├── config.json         # all export arguments (read by annotate.py)
+├── config.json         # export arguments + cell_type_vocabulary (read by annotate.py)
 ├── metadata.csv        # cell_id, cell_type, image_path, sampling_weight, center_x, center_y
-└── annotations.json    # {cell_id: true/false} — written by annotate.py
+└── annotations.json    # {cell_id: {action, label}} — written by annotate.py
 ```
 
 Each panel contains:
@@ -127,12 +133,27 @@ A simple JSON mapping gene names to colours:
 
 ## Step 2 — Manual validation (`scripts/annotate.py`)
 
-For each exported image, an annotator decides whether the cell is **correct** — meaning it is both properly **segmented** and properly **labelled**. A cell should be rejected (marked incorrect) when either check fails:
+For each exported image, the annotator runs two checks in order:
+
+1. **Segmentation** — is the cell boundary (left panel) consistent with the nuclei / membrane staining, or does it merge two cells or clip part of one?
+2. **Cell-type label** — are the top expressed genes (right panel) consistent with the assigned label?
+
+which lead to one of three actions:
+
+| action | when | effect |
+| --- | --- | --- |
+| **accept** | segmentation fine, label fine | the cell keeps its automated label |
+| **correct** | segmentation fine, label wrong | the annotator picks the right cell type |
+| **reject** | segmentation inadequate | the cell is excluded from both compared groups |
 
-- **Segmentation** — the cell boundary (left panel) is not consistent with the nuclei / membrane staining, e.g. it merges two cells or clips part of one.
-- **Cell-type label** — the top expressed genes (right panel) include genes unlikely to be expressed by the cell type of interest, suggesting the automated label is wrong.
+Correcting a cell revises only its **cell type**; its spatial region is treated as reliable
+and is always taken from the automated pipeline. So correcting a cell *into* the cell type
+of interest is what places it in the target or reference group, according to the region it
+already sits in — this is the case an accept/reject workflow cannot express.
 
-The result is a boolean column `is_correct` added to `metadata.csv`, which becomes `adata_gt` in Step 3.
+Segmentation is never edited: an accepted or corrected cell keeps the automated expression
+counts. Rejection therefore doubles as a quality-control filter for cells whose
+quantification cannot be trusted at all.
 
 ```bash
 streamlit run scripts/annotate.py -- --dir /path/to/annotation_dir
@@ -142,10 +163,19 @@ The `--` is required: it tells Streamlit to pass everything after it to the scri
 
 VS Code Remote forwards the Streamlit port automatically. Open the URL printed in the terminal, then use:
 
-- **`1`** — label as correct
-- **`2`** — label as incorrect
+| key | `accept_correct_reject` (default) | `accept_reject` |
+| --- | --- | --- |
+| **`1`** | accept | accept |
+| **`2`** | correct | reject |
+| **`3`** | reject | — |
 
-Progress is saved after every keypress to `annotations.json`. Re-running the command resumes from where you left off. You can also start annotating while `export.py` is still running — the UI picks up newly exported cells automatically.
+Pressing **`2`** in `accept_correct_reject` mode opens a cell-type selector below the panel
+— type a few characters to filter, then pick the label. Nothing is written until you
+choose one, so pressing `2` by mistake is harmless: hit Cancel and the cell stays
+unannotated.
+
+Progress is saved after every keypress to `annotations.json`, as
+`{cell_id: {"action": ..., "label": ...}}` (`label` is set only for corrections). Re-running the command resumes from where you left off. You can also start annotating while `export.py` is still running — the UI picks up newly exported cells automatically.
 
 ---
 
@@ -157,11 +187,27 @@ python scripts/differential_expression.py --dir /path/to/annotation_dir
 
 Reads all export settings from `config.json` and writes gene-level results to `<dir>/results.csv`.
 
+The three-way comparison is built here: cells of interest in spatial group `0` (reference)
+and group `1` (target) form the two compared populations, and everything else — including
+rejected cells — is collapsed into a third group. Both the automated labels and the manual
+ones are built the same way; only the cell type differs between them. The script prints a
+summary of the annotations first (counts per action, plus how many cells the curation moved
+into and out of the compared groups), which is the quickest check that the annotations were
+read as intended.
+
+If your region column is not encoded as `1` / `0`, set `--spatial-group-target` and
+`--spatial-group-reference` to the two values you want to compare; the script reports the
+values it found if they don't match. The target region is the one positive log-fold changes
+refer to, so swapping the two flips the sign of every result — this is deliberately not
+inferred for you, even when the column has exactly two values.
+
 | option | default | description |
 |---|---|---|
 | `--dir` | *(required)* | annotation directory (output of steps 1 & 2) |
 | `--out` | `<dir>/results.csv` | output CSV path |
 | `--spatial-group-key` | `spatial_group` | obs column encoding the two spatial populations |
+| `--spatial-group-target` | `1` | value of that column identifying the target region (group 1) |
+| `--spatial-group-reference` | `0` | value of that column identifying the reference region (group 0) |
 | `--n-cells-expressed-threshold` | `10` | min annotated cells expressing a gene for it to be tested |
 | `--noise-model` | `poisson` | `poisson` or `nb` (negative binomial) |
 

pyproject.toml

@@ -33,7 +33,7 @@ dev = [
     "isort",
     "flake8",
 ]
-annotate = ["streamlit"]
+annotate = ["streamlit>=1.30"]
 cuda12 = ["jax[cuda12]"]
 cuda13 = ["jax[cuda13]"]
 

scripts/annotate.py

@@ -4,6 +4,18 @@
 Usage
 -----
 streamlit run scripts/annotate.py -- --dir /path/to/annotations/R2_macrophages
+
+The available actions are set by ``annotation_mode`` in config.json (written by
+export.py):
+
+- ``accept_reject``          1 = accept, 2 = reject
+- ``accept_correct_reject``  1 = accept, 2 = correct, 3 = reject
+
+Writes ``annotations.json`` as
+``{cell_id: {"action": "accept" | "correct" | "reject", "label": str | None}}``.
+``label`` is set only for corrections; the manual cell type of an accepted cell is
+resolved from the automated label downstream. This schema is validated by
+``csde.read_annotations`` — keep the two in step.
 """
 
 import argparse
@@ -14,6 +26,20 @@
 import streamlit as st
 import streamlit.components.v1 as components
 
+ACCEPT, CORRECT, REJECT = "accept", "correct", "reject"
+
+# Action -> (button glyph, button label). The glyph is what the keyboard handler
+# matches on, so each must be unique.
+ACTION_STYLE = {
+    ACCEPT: ("✓", "Accept"),
+    CORRECT: ("✎", "Correct"),
+    REJECT: ("✗", "Reject"),
+}
+MODE_ACTIONS = {
+    "accept_reject": [ACCEPT, REJECT],
+    "accept_correct_reject": [ACCEPT, CORRECT, REJECT],
+}
+
 
 def parse_args():
     p = argparse.ArgumentParser()
@@ -34,20 +60,43 @@ def save_annotations(annotations: dict, annotation_dir: Path) -> None:
         json.dump(annotations, f, indent=2)
 
 
+def format_status(record: dict | None) -> str:
+    if record is None:
+        return "not annotated"
+    glyph, label = ACTION_STYLE[record["action"]]
+    if record["action"] == CORRECT:
+        return f"{glyph} corrected → **{record['label']}**"
+    return f"{glyph} {label.lower()}ed"
+
+
 def main():
     args = parse_args()
     annotation_dir = Path(args.dir)
 
     cell_type_of_interest = "cell of interest"
+    annotation_mode = "accept_correct_reject"
+    vocabulary = []
     config_path = annotation_dir / "config.json"
     if config_path.exists():
         with open(config_path) as f:
             config = json.load(f)
         cell_type_of_interest = config.get("cell_type_of_interest", cell_type_of_interest)
+        annotation_mode = config.get("annotation_mode", annotation_mode)
+        vocabulary = config.get("cell_type_vocabulary", [])
+
+    actions = MODE_ACTIONS[annotation_mode]
 
     st.set_page_config(layout="wide", page_title="Cell Annotator")
     st.title(f"Cell Annotation — {cell_type_of_interest}")
 
+    if CORRECT in actions and not vocabulary:
+        st.error(
+            "`cell_type_vocabulary` is missing from config.json, so corrections "
+            "cannot offer any cell type. Re-run scripts/export.py, or set "
+            '`"annotation_mode": "accept_reject"` in config.json.'
+        )
+        st.stop()
+
     metadata_path = annotation_dir / "metadata.csv"
     if not metadata_path.exists():
         st.warning("metadata.csv not found — waiting for export.py to write the first cell.")
@@ -103,13 +152,14 @@ def main():
     row = metadata.iloc[idx]
     cell_id = row["cell_id"]
     existing = annotations.get(cell_id)
-    status = "✓ correct" if existing is True else ("✗ incorrect" if existing is False else "not annotated")
 
-    st.subheader(f"Cell `{cell_id}` — predicted: **{row['cell_type']}** — {status}")
+    st.subheader(
+        f"Cell `{cell_id}` — predicted: **{row['cell_type']}** — {format_status(existing)}"
+    )
     st.image(str(row["image_path"]), use_container_width=True)
 
-    def annotate(is_correct: bool) -> None:
-        annotations[cell_id] = is_correct
+    def annotate(action: str, label: str | None = None) -> None:
+        annotations[cell_id] = {"action": action, "label": label}
         save_annotations(annotations, annotation_dir)
         # Advance to next unannotated after annotating
         remaining = metadata.index[~metadata["cell_id"].isin(annotations)]
@@ -119,32 +169,74 @@ def annotate(is_correct: bool) -> None:
         elif remaining.any():
             st.session_state.current_idx = int(remaining[0])
 
-    col1, col2, _ = st.columns([1, 1, 4])
-    with col1:
-        if st.button("✓ Correct [1]", type="primary", use_container_width=True):
-            annotate(True)
-            st.rerun()
-    with col2:
-        if st.button("✗ Incorrect [2]", use_container_width=True):
-            annotate(False)
+    select_key = f"correction_choice_{cell_id}"
+    pending = st.session_state.get("pending_correction") == cell_id
+
+    if pending:
+        # The action buttons are removed while a correction is pending, so the
+        # digit shortcuts below match nothing and cannot fire. The selector is
+        # rendered under the pan
+        # el to keep the image visible while choosing.
+        st.markdown(f"**Correcting cell `{cell_id}`** — pick the right cell type:")
+        sel_col, cancel_col, _ = st.columns([2, 1, 3])
+        with sel_col:
+            choice = st.selectbox(
+                "Corrected cell type",
+                [t for t in vocabulary if t != str(row["cell_type"])],
+                index=None,
+                placeholder="Type to filter…",
+                label_visibility="collapsed",
+                key=select_key,
+            )
+        with cancel_col:
+            if st.button("Cancel", use_container_width=True):
+                st.session_state.pending_correction = None
+                st.rerun()
+        if choice is not None:
+            annotate(CORRECT, label=choice)
+            st.session_state.pending_correction = None
             st.rerun()
+    else:
+        cols = st.columns([1] * len(actions) + [6 - len(actions)])
+        for col, action in zip(cols, actions):
+            glyph, label = ACTION_STYLE[action]
+            key_hint = actions.index(action) + 1
+            with col:
+                clicked = st.button(
+                    f"{glyph} {label} [{key_hint}]",
+                    type="primary" if action == ACCEPT else "secondary",
+                    use_container_width=True,
+                )
+            if clicked:
+                if action == CORRECT:
+                    # Nothing is written yet: the cell stays unannotated until a
+                    # label is picked.
+                    st.session_state.pending_correction = cell_id
+                    st.session_state.pop(select_key, None)
+                else:
+                    annotate(action)
+                st.rerun()
 
-    # Keyboard shortcuts: 1/2 annotate, ←/→ navigate
-    components.html("""
+    # Keyboard shortcuts: digits annotate, ←/→ navigate
+    key_to_glyph = {
+        str(i + 1): ACTION_STYLE[action][0] for i, action in enumerate(actions)
+    }
+    components.html(f"""
     <script>
-    (function() {
+    (function() {{
+        const keyToGlyph = {json.dumps(key_to_glyph)};
         const doc = window.parent.document;
-        doc.addEventListener('keydown', function(e) {
+        doc.addEventListener('keydown', function(e) {{
             if (e.target.tagName === 'INPUT' || e.target.tagName === 'TEXTAREA') return;
-            doc.querySelectorAll('button').forEach(function(btn) {
+            const glyph = keyToGlyph[e.key];
+            doc.querySelectorAll('button').forEach(function(btn) {{
                 const t = btn.textContent.trim();
-                if (e.key === '1' && t.startsWith('✓')) btn.click();
-                if (e.key === '2' && t.startsWith('✗')) btn.click();
+                if (glyph && t.startsWith(glyph)) btn.click();
                 if (e.key === 'ArrowLeft'  && t.startsWith('←')) btn.click();
                 if (e.key === 'ArrowRight' && t.startsWith('Next')) btn.click();
-            });
-        }, true);
-    })();
+            }});
+        }}, true);
+    }})();
     </script>
     """, height=0)
 

scripts/differential_expression.py

@@ -24,6 +24,13 @@ def parse_args():
                    help="Output CSV path (default: <dir>/results.csv).")
     p.add_argument("--spatial-group-key", default="spatial_group",
                    help="obs column encoding the two spatial populations.")
+    p.add_argument("--spatial-group-target", default="1",
+                   help="Value of --spatial-group-key identifying the target region "
+                        "(group 1). Positive log-fold changes mean higher expression "
+                        "here.")
+    p.add_argument("--spatial-group-reference", default="0",
+                   help="Value of --spatial-group-key identifying the reference "
+                        "region (group 0).")
     p.add_argument("--n-cells-expressed-threshold", type=int, default=10,
                    help="Min cells expressing a gene for it to be tested.")
     p.add_argument("--noise-model", default="poisson", choices=["poisson", "nb"])
@@ -42,21 +49,28 @@ def main():
     inputs = prepare_csde_inputs(
         annotation_dir=annotation_dir,
         spatial_group_key=args.spatial_group_key,
+        spatial_group_target=args.spatial_group_target,
+        spatial_group_reference=args.spatial_group_reference,
         layer=layer,
         n_cells_expressed_threshold=args.n_cells_expressed_threshold,
     )
     adata_gt = inputs["adata_gt"]
     adata_other = inputs["adata_other"]
+    summary = inputs["summary"]
+
+    print("Manual annotations:")
+    for key, value in summary.items():
+        print(f"  {key}: {value}")
 
     results = run_csde(
         adata_pred=adata_other,
         adata_gt=adata_gt,
         pred_cell_pop_key="prediction",
+        gt_cell_pop_key="annotation",
         cell_pop_a=0,
         cell_pop_b=1,
-        gt_key="is_correct",
         layer_name=layer,
-        importance_weights=adata_gt.obs["sampling_weight"].values,
+        importance_weights=1.0 / adata_gt.obs["sampling_weight"].values,
         noise_model=args.noise_model,
     )
 

scripts/export.py

@@ -35,6 +35,11 @@ def parse_args():
                    help="JSON file mapping gene name → colour.")
     p.add_argument("--image-channel", default="DAPI")
     p.add_argument("--n-cells", type=int, default=600)
+    p.add_argument("--annotation-mode",
+                   choices=["accept_correct_reject", "accept_reject"],
+                   default="accept_correct_reject",
+                   help="Actions offered by annotate.py. accept_correct_reject "
+                        "additionally lets the annotator relabel a cell.")
     p.add_argument("--delta", type=float, default=50.0,
                    help="Half-width of the spatial crop around each cell (microns).")
     p.add_argument("--n-top-genes", type=int, default=15)
@@ -58,8 +63,10 @@ def main():
     annotation_dir = Path(args.out)
     annotation_dir.mkdir(parents=True, exist_ok=True)
 
+    cell_types = sdata["table"].obs[args.cell_type_key].dropna().unique()
+    config = vars(args) | {"cell_type_vocabulary": sorted(map(str, cell_types))}
     with open(annotation_dir / "config.json", "w") as f:
-        json.dump(vars(args), f, indent=2)
+        json.dump(config, f, indent=2)
 
     metadata = export_cell_panels(
         sdata=sdata,