RootScope v4: manuscript model (LightGBM per round, 3 seeds, layer + soft-neighbor context, radial prior)

README.md

@@ -1,5 +1,5 @@
 ---
-title: 'RootScope: Cross-species Root Cell-Type Classification from Confocal Microscopy Images'
+title: 'RootScope v4: Cross-species Root Cell-Type Classification from Confocal Microscopy Images'
 emoji: 🌱
 colorFrom: green
 colorTo: blue
@@ -14,9 +14,16 @@ short_description: Classify root cell types across species from confocal TIFFs
 # RootScope web app
 
 Gradio front end for [RootScope](https://github.com/ct-tranchau/Rootscope):
-upload a confocal root-tip cross-section TIFF, get back a labeled overlay and a
+upload a confocal root cross-section TIFF, get back a labeled overlay and a
 per-cell CSV. Runs on a Hugging Face Space with ZeroGPU.
 
+**Since 4.0.0 the app runs the v4 model** (the one in the manuscript):
+LightGBM per refinement round, three seeds averaged, soft-neighbor and
+layer-distance context rebuilt each round, radial-order prior at the end.
+The model dropdown therefore shows a single entry. Set
+`ROOTSCOPE_MODEL_VERSION=v2` in the Space's variables to run the old ensemble.
+v4 keeps every segmented cell (no debris filter), as the manuscript did.
+
 ## Do not load the classifiers in the main process
 
 Unpickling XGBoost builds a Booster that probes for CUDA devices. On ZeroGPU

app.py

@@ -31,7 +31,7 @@ from PIL import Image
 from rootscope import predict as rs
 from rootscope.cnn_embeddings import load_dinov2
 from rootscope.extract_features import load_cellpose_model
-from rootscope.weights import resolve_cnn_weights, resolve_model_dir
+from rootscope.weights import is_v4, model_version, resolve_cnn_weights, resolve_model_dir
 
 # ── ZeroGPU shim ─────────────────────────────────────────────────────────────
 # `spaces` only exists on Hugging Face. Locally the decorator becomes a no-op
@@ -68,6 +68,8 @@ MAX_PIXELS = 40_000_000  # ~6300x6300; refuse anything larger
 # Settings → Variables and restart; no code change needed.
 SHOW_DIAGNOSTICS = os.environ.get("ROOTSCOPE_DIAGNOSTICS") == "1"
 MODEL_CHOICES = ["Ensemble", "RandomForest", "LightGBM", "XGBoost"]
+if is_v4():
+    MODEL_CHOICES = ["LightGBM"]   # v4 ships one classifier (3 seeds averaged inside)
 
 # ── fetch weights at startup, load nothing ───────────────────────────────────
 # Downloading here means the first visitor does not wait for ~550 MB. Note that
@@ -161,13 +163,8 @@ def _segment_duration(img_rgb):
 
 
 def _embed_duration(masks, img_rgb, df_base):
-    # v2's backbone is DINOv2 ViT-B/14 (86.6 M parameters) where v1 used
-    # ViT-S/14 (22.1 M), about 3.5x the work per cell. Measured on the live
-    # Space, v1 embedded 214 cells in ~2 s, so v2 should need ~7 s; the floor
-    # and slope below are roughly 5x that, keeping the same deliberate margin
-    # the v1 numbers had. Still capped under the signed-out quota.
     n = int(masks.max())
-    return int(min(MAX_DURATION, max(30, 12 + 0.12 * n)))
+    return int(min(MAX_DURATION, max(25, 10 + 0.06 * n)))
 
 
 def _gpu_call(fn, *args):
@@ -509,7 +506,8 @@ def run_rootscope(tif_path, tif_file, um_per_px, which_model, label_cells,
     # ---- CPU: tissue layers + handcrafted features ----
     progress(0.35, desc=f"{n_cells} cells, layer index and features")
     masks, df_base, layer_lookup, adjacency, n_layers = rs.stage_features(
-        masks, img_rgb, um_per_px=float(um_per_px))
+        masks, img_rgb, um_per_px=float(um_per_px),
+        remove_debris=not is_v4())   # v4 keeps every cell, as in the manuscript
     if df_base is None:
         raise gr.Error("Every detected object was filtered out as debris "
                        "outside the tissue body.")
@@ -618,9 +616,11 @@ with gr.Blocks(title=TITLE, **_blocks_kwargs) as demo:
                      "feature.",
             )
             model_pick = gr.Dropdown(
-                MODEL_CHOICES, value="Ensemble", label="Model to display",
-                info="All four are computed in one run; this only picks which "
-                     "one you see. Ensemble is the published result.",
+                MODEL_CHOICES, value=MODEL_CHOICES[0], label="Model to display",
+                info=("v4: one LightGBM (three training seeds averaged inside); "
+                      "this is the published result." if is_v4() else
+                      "All four are computed in one run; this only picks which "
+                      "one you see. Ensemble is the published result."),
             )
             with gr.Accordion("Advanced", open=False):
                 label_cells = gr.Checkbox(
@@ -629,9 +629,12 @@ with gr.Blocks(title=TITLE, **_blocks_kwargs) as demo:
                          "section the text overlaps and hides the image.",
                 )
                 max_rounds = gr.Slider(
-                    1, 20, value=10, step=1, label="Max refinement rounds",
-                    info="The classifier re-predicts using neighbours' types "
-                         "until predictions stop changing.",
+                    1, 20, value=(6 if is_v4() else 10), step=1,
+                    label="Max refinement rounds",
+                    info=("v4 was evaluated at 6; going to 10 changes accuracy "
+                          "by under 0.1 pp." if is_v4() else
+                          "The classifier re-predicts using neighbours' types "
+                          "until predictions stop changing."),
                 )
             run_btn = gr.Button("Run RootScope", variant="primary")
             # Hugging Face meters free GPU per visitor, not per Space, so the

classify_worker.py

@@ -40,6 +40,26 @@ def run(job):
     # Imported here, not at module scope, so that even this process does no
     # model work until it has a job to do.
     from rootscope import predict as rs
+    from rootscope.weights import is_v4
+
+    if is_v4():
+        # v4: one LightGBM per round x 3 seeds, soft-neighbor + layer features
+        # rebuilt each round, radial-order prior at the end. No XGBoost anywhere,
+        # but the subprocess isolation is kept: it costs ~2 s and guards the
+        # parent's CUDA state against the same class of bug.
+        df = rs.stage_classify_v4(
+            job["df_base"], job["masks"], job["img_rgb"],
+            model_dir=job["model_dir"],
+            out_dir=job["out_dir"],
+            stem=job["stem"],
+            source_name=job["source_name"],
+            um_per_px=job["um_per_px"],
+            max_rounds=job["max_rounds"],
+            label_cells=job["label_cells"],
+        )
+        from rootscope.classify_v4 import CLASSES
+        return {"ok": True, "df": df, "classes": list(CLASSES),
+                "models": ["LightGBM"]}
 
     models, scalers, feature_cols, le = rs.load_models(job["model_dir"])
     if not models:

requirements.txt

@@ -1,26 +1,38 @@
-# torch / torchvision are deliberately absent: ZeroGPU supplies its own CUDA
-# build (2.8+) and pinning them here would fight it. Same pattern the
-# Cellpose-SAM Space uses.
-gradio
+# Dependencies for the RootScope v4 Hugging Face Space.
+#
+# The Space vendors its own copy of rootscope/ (uploaded next to app.py), so
+# RootScope is NOT installed from GitHub here. Doing both would put two
+# packages of the same name on sys.path and whichever imports first wins --
+# a bug that shows up only as wrong predictions.
+#
+# torch is NOT pinned, deliberately: ZeroGPU supports only torch >= 2.8 and
+# picks the version. torch 2.6+ defaults torch.load(weights_only=True), which
+# backbone.pt and the Cellpose-SAM checkpoint must load under -- test the two
+# example images on the Space before announcing a change.
+#
+# v4 needs no xgboost: the classifier is LightGBM only. It is not listed on
+# purpose (unpickling an XGBoost Booster is what poisons ZeroGPU's CUDA state;
+# see classify_worker.py). Do not add it back for "completeness".
 
-numpy
-pandas
-scipy
-scikit-image
-# Hard pin: the released .joblib models were pickled with 1.7.2 and other
-# versions can fail to unpickle or silently change behaviour.
+gradio>=5.0
+spaces>=0.30
+
+# scikit-learn is a hard pin: the LightGBM bundles carry StandardScalers pickled
+# with 1.7.2, and other versions can fail to unpickle or change results.
 scikit-learn==1.7.2
-opencv-python-headless
-Pillow
-tifffile
-imagecodecs
-matplotlib
-seaborn
-tqdm
-joblib
 
-xgboost
-lightgbm
+numpy>=1.24
+pandas>=2.0
+scipy>=1.10
+scikit-image>=0.22
+opencv-python-headless>=4.8
+Pillow>=10.0
+tifffile>=2024.1.30
+imagecodecs>=2024.1.1
+joblib>=1.3
+tqdm>=4.66
+lightgbm>=4.0
 cellpose>=4.0
-
+torch
+torchvision
 huggingface_hub>=0.24

rootscope/__init__.py

@@ -1,5 +1,5 @@
 """
-Rootscope — cell-type prediction for confocal root-tip cross-section images.
+Rootscope: cell-type prediction for confocal root-tip cross-section images.
 
 Give it a TIFF, it segments the cells (Cellpose-SAM), extracts shape / size /
 intensity / layer features plus fine-tuned DINOv2 embeddings, and predicts the
@@ -16,7 +16,7 @@ Typical use (Python):
     df = predict_tif("my_image.tif", out_dir="results/", gpu=True)
 """
 
-__version__ = "0.1.0"
+__version__ = "4.0.0"
 
 from .api import predict_tif, predict_folder  # noqa: E402,F401
 

rootscope/_eval_honest.py

@@ -0,0 +1,491 @@
+"""
+BUOC 5 -- danh gia trung thuc, khop dung cach tool chay tren anh moi.
+
+Ba khac biet so voi train_iterative_cnn.py:
+
+  1. DINO backbone chi hoc 100 section train -> cnn_emb_* cua 26 section test sach.
+
+  2. Hang xom cua tap TEST khong lay tu BMP. Vong 1 = -1, vong sau lay tu chinh
+     PREDICTION, tra qua adjacency that (dung cell_id), y het
+     predict_cell_types_iterative.py. Khong dung
+     _fill_neighbor_celltypes_from_predictions() vi ham do doc GT tu CSV.
+
+  3. Hinh hoc hang xom tinh tren TAT CA te bao segment duoc (df_geom), khong chi
+     te bao co nhan -- vi tren anh moi khong he co khai niem "khong nhan" hay
+     "do tin cay thap"; moi te bao deu co mat trong mo va deu duoc du doan.
+     Chi tap con co nhan + du tin cay (df_model) moi dung de train va cham diem.
+
+Dong trung (source_file, cell_id) bi bo. Nhan BMP cua 26 section chi dung de cham diem.
+"""
+import argparse, json, sys
+from pathlib import Path
+
+import joblib
+import numpy as np
+import pandas as pd
+from sklearn.metrics import accuracy_score
+from sklearn.model_selection import StratifiedKFold, StratifiedGroupKFold
+from sklearn.preprocessing import LabelEncoder, StandardScaler
+from sklearn.utils.class_weight import compute_sample_weight
+
+ROOT = Path(__file__).resolve().parent   # khong dung khi suy dien; giu cho tuong thich
+
+# Copied verbatim from the research training module (train_iterative_cnn.py) so
+# that inference does not import it: that module pulls in xgboost, matplotlib and
+# seaborn at import time, none of which v4 inference uses. Only NEIGHBOR_CT_COLS
+# is read on the inference path (in _blank); FEATURE_COLUMNS is kept for
+# load_tables(), and the two training helpers are stubs that refuse to run.
+FEATURE_COLUMNS = [
+    "layer_index", "normalized_radius", "edt_distance_px", "edt_normalized",
+    "layer_fraction", "is_boundary_cell", "dist_from_centroid_um",
+    "angular_position", "n_layers_total",
+    "area_um2", "perimeter_um", "eccentricity", "solidity", "aspect_ratio",
+    "compactness", "major_axis_px", "minor_axis_px", "orientation_rad",
+    "extent", "area_perimeter_ratio", "equivalent_diameter",
+    "mean_intensity", "std_intensity", "min_intensity", "max_intensity",
+    "median_intensity", "intensity_range", "intensity_cv",
+    "intensity_skewness", "intensity_kurtosis", "intensity_p10", "intensity_p90",
+    "neighbors_count", "mean_neighbor_layer", "std_neighbor_layer",
+    "mean_neighbor_area", "layer_diff_from_neighbors", "area_ratio_to_neighbors",
+    "touches_background", "cells_in_same_layer", "frac_neighbors_same_layer",
+    "frac_neighbors_inner", "frac_neighbors_outer", "inner_neighbor_count",
+    "outer_neighbor_count", "layer_from_inside", "radial_intensity_gradient",
+    "neighbor_layer_range", "neighbors_boundary_cell", "min_neighbor_layer",
+    "max_neighbor_layer", "n_neighbors_touching_bg",
+    "area_zscore", "area_ratio_to_global_median", "area_percentile_in_layer",
+    "area_ratio_to_inner", "area_ratio_to_outer", "hexagonality", "n_vertices",
+    "layer_cell_count_ratio", "inner_layer_cell_count", "outer_layer_cell_count",
+    "layer_count_gradient", "layer_count_asymmetry", "adjacent_layer_area_ratio",
+    "sin_angular_position", "cos_angular_position", "radial_x_sin", "radial_x_cos",
+    "radial_inward_neighbor_area", "radial_outward_neighbor_area",
+    "cw_neighbor_area", "ccw_neighbor_area",
+    "radial_inward_neighbor_intensity", "radial_outward_neighbor_intensity",
+    "cw_neighbor_intensity", "ccw_neighbor_intensity",
+    "radial_inward_neighbor_celltype", "radial_outward_neighbor_celltype",
+    "tangential_cw_neighbor_celltype", "tangential_ccw_neighbor_celltype",
+    "wall_thickness_proxy", "wall_to_lumen_ratio", "local_area_rank_in_stele",
+    "neighbor_area_std", "cell_wall_contrast",
+    "lumen_darkness", "wall_lumen_gap", "interior_intensity_std",
+    "frac_dark_interior", "ring1_intensity", "ring2_intensity",
+    "wall_interior_gradient", "mean_neighbor_wall_thickness",
+    "wall_thickness_vs_neighbors", "area_ratio_to_stele_mean",
+]
+
+NEIGHBOR_CT_COLS = [
+    "radial_inward_neighbor_celltype", "radial_outward_neighbor_celltype",
+    "tangential_cw_neighbor_celltype", "tangential_ccw_neighbor_celltype",
+]
+
+
+def _build_model(name, n_est, seed):  # training-only; not shipped with the package
+    raise RuntimeError("_build_model is a training helper and is not part of the "
+                       "RootScope inference package.")
+
+
+def _compute_sample_weights(*a, **k):  # training-only; not shipped with the package
+    raise RuntimeError("_compute_sample_weights is a training helper and is not part "
+                       "of the RootScope inference package.")
+from .extract_features import (  # noqa: E402
+    build_cell_adjacency, compute_neighbor_celltypes, CELL_CLASSES,
+)
+
+OUT = ROOT / "Figure_groundtruth_predicted"
+
+_BASE_BUILD = _build_model          # giu ban goc truoc khi thay
+HP = {}                             # hyperparameter de de len tren, tu --hp
+
+
+def _build_model_hp(name, n_est, seed):
+    """Nhu _build_model nhung (a) nhan them 'Ensemble' = soft-vote ca 3 model,
+    (b) cho phep de hyperparameter tu --hp len tren mac dinh hardcode."""
+    if name == "Ensemble":
+        from sklearn.ensemble import VotingClassifier
+        return VotingClassifier(
+            [(n, _build_model_hp(n, n_est, seed))
+             for n in ("RandomForest", "XGBoost", "LightGBM")],
+            voting="soft")
+    m = _BASE_BUILD(name, n_est, seed)
+    if HP:
+        ok = {k: v for k, v in HP.items() if k in m.get_params()}
+        if ok:
+            m.set_params(**ok)
+    return m
+
+
+def dedup_best(df):
+    """
+    Bo dong trung (source_file, cell_id).
+
+    11 section co nhieu hon 1 file BMP (139 BMP / 126 section). Hai ban nhan
+    khong phai luc nao cung giong: do duoc 8.2% cap mau thuan. Nen giu ban co
+    cell_type_confidence CAO NHAT thay vi lay bua dong dau tien.
+    """
+    n0 = len(df)
+    if "cell_type_confidence" in df.columns:
+        df = df.sort_values("cell_type_confidence", ascending=False,
+                            na_position="last", kind="mergesort")
+    df = df.drop_duplicates(["source_file", "cell_id"], keep="first")
+    df = df.sort_values(["source_file", "cell_id"], kind="mergesort").reset_index(drop=True)
+    if len(df) != n0:
+        print(f"  bo {n0-len(df):,} dong trung (giu ban confidence cao nhat)")
+    return df
+
+
+def load_tables(csv_path, min_confidence):
+    df = pd.read_csv(csv_path, low_memory=False)
+    df = dedup_best(df)
+
+    available = [c for c in FEATURE_COLUMNS if c in df.columns]
+    if MT_KEEP is not None:
+        available = [c for c in available if c in MT_KEEP]
+        print(f"  chi giu {len(available)} feature hinh thai")
+    if USE_NB_AGG:
+        for c in NB_AGG_COLS:
+            df[c] = -1.0
+        available += NB_AGG_COLS
+        print(f"  THEM {len(NB_AGG_COLS)} feature thanh phan hang xom")
+    cnn = sorted([c for c in df.columns if c.startswith("cnn_emb_")],
+                 key=lambda x: int(x.split("_")[-1]))
+    sec = sorted([c for c in df.columns if c.startswith("sec_")])
+    if sec:
+        available += sec
+        print(f"  INCLUDING {len(sec)} feature muc section: {sec}")
+    if cnn:
+        available += cnn
+        print(f"  INCLUDING {len(cnn)} CNN embedding features")
+    else:
+        print("  WARNING: khong thay cot cnn_emb_*")
+
+    geom = df.dropna(subset=available).reset_index(drop=True)
+    m = (geom["cell_type_label"] >= 0) & geom["cell_type"].notna()
+    if "cell_type_confidence" in geom.columns:
+        m &= geom["cell_type_confidence"] >= min_confidence
+    model = geom[m].reset_index(drop=True)
+    print(f"  segment duoc: {len(geom):,} te bao;  dung train/cham diem: {len(model):,}")
+    return geom, model, available
+
+
+_ADJ = {}
+
+
+def _adjacency_for(source_file, mask_dir):
+    if source_file not in _ADJ:
+        p = Path(mask_dir) / f"{Path(source_file).stem}_masks.npy"
+        if not p.exists():
+            raise FileNotFoundError(f"thieu mask: {p}")
+        if len(_ADJ) > 200:   # du chua ca 126 section, khoi tinh lai moi vong
+            _ADJ.clear()
+        _ADJ[source_file] = build_cell_adjacency(np.load(p))
+    return _ADJ[source_file]
+
+
+def fill_neighbors(df_geom, pred_names, mask_dir):
+    """Dien 4 cot hang xom tu NHAN DU DOAN cua MOI te bao, qua adjacency that."""
+    pred_names = np.asarray(pred_names)
+    parts = []
+    for sf, idx in df_geom.groupby("source_file").groups.items():
+        sub = df_geom.loc[idx].copy()
+        pos = df_geom.index.get_indexer(idx)
+        labels = {int(c): CELL_CLASSES.get(p, -1)
+                  for c, p in zip(sub["cell_id"].values, pred_names[pos])}
+        adj = _adjacency_for(sf, mask_dir)
+        part = compute_neighbor_celltypes(sub, adj, cell_type_labels=labels)
+        if USE_NB_AGG:
+            names_by_cid = {int(c): pn for c, pn in zip(sub["cell_id"].values,
+                                                        pred_names[pos])}
+            for k, v in _nb_aggregates(sub, adj, names_by_cid).items():
+                part[k] = v
+        parts.append(part)
+    return pd.concat(parts).loc[df_geom.index]
+
+
+def _sample_weights(y, le, wp_boost):
+    """Nhu _compute_sample_weights goc, nhung he so cho xylem/phloem chinh duoc.
+
+    "balanced" da lo phan MAT CAN BANG (nghich dao tan suat). Moi he so nhan
+    them sau do la UU TIEN, khong phai chinh mat can bang. Ban goc hardcode
+    x15 cho xylem va phloem -- day la de do lai xem no con dang gia khong."""
+    w = compute_sample_weight("balanced", y)
+    for c in ("xylem", "phloem"):
+        if c in le.classes_:
+            w[y == le.transform([c])[0]] *= wp_boost
+    return w
+
+
+VASCULAR = ("stele", "xylem", "phloem")
+OUTER_TISSUE = ("epidermis", "exodermis", "root_cap")
+NB_AGG_COLS = ["frac_nb_vascular", "n_nb_vascular", "frac_nb_cortex",
+               "frac_nb_same", "frac_nb_outer_tissue"]
+ALL_CLASSES = ["root_cap", "epidermis", "exodermis", "cortex", "endodermis",
+               "pericycle", "stele", "xylem", "phloem"]
+NB_PER_CLASS = [f"frac_nb_{c}" for c in ALL_CLASSES]
+USE_NB_AGG = False
+MT_KEEP = None
+
+
+def _nb_aggregates(sub_df, adj, names_by_cid):
+    """Thanh phan lop cua TAT CA hang xom quanh moi cell.
+
+    4 cot NEIGHBOR_CT_COLS chi lay DUNG MOT hang xom gan nhat moi huong, nen
+    cell pericycle nam o lop trong bi pericycle khac che se mat han tin hieu
+    stele (do duoc: 393/843 cell roi vao truong hop nay). Cac cot nay tong hop
+    tren MOI hang xom nen khong bi che.
+
+    Tinh du 5 dai luong roi chi tra ve nhung cot dang duoc yeu cau.
+    """
+    allv = {c: [] for c in ("frac_nb_vascular", "n_nb_vascular", "frac_nb_cortex",
+                            "frac_nb_same", "frac_nb_outer_tissue")}
+    for cid in sub_df["cell_id"].values:
+        nb = [names_by_cid.get(b) for b in adj.get(int(cid), ())]
+        nb = [x for x in nb if x is not None and x != "__none__"]
+        n = len(nb)
+        me = names_by_cid.get(int(cid))
+        if n == 0:
+            for c in allv:
+                allv[c].append(-1.0)
+            continue
+        nv = sum(x in VASCULAR for x in nb)
+        allv["frac_nb_vascular"].append(nv / n)
+        allv["n_nb_vascular"].append(float(nv))
+        allv["frac_nb_cortex"].append(sum(x == "cortex" for x in nb) / n)
+        allv["frac_nb_same"].append(sum(x == me for x in nb) / n if me else -1.0)
+        allv["frac_nb_outer_tissue"].append(sum(x in OUTER_TISSUE for x in nb) / n)
+    for c in ALL_CLASSES:
+        allv[f"frac_nb_{c}"] = []
+    for cid in sub_df["cell_id"].values:
+        nb = [names_by_cid.get(b) for b in adj.get(int(cid), ())]
+        nb = [x for x in nb if x is not None and x != "__none__"]
+        n = len(nb)
+        for c in ALL_CLASSES:
+            allv[f"frac_nb_{c}"].append(sum(x == c for x in nb) / n if n else -1.0)
+    return {c: allv[c] for c in NB_AGG_COLS}
+
+
+def _blank(df):
+    d = df.copy()
+    for c in list(NEIGHBOR_CT_COLS) + list(NB_AGG_COLS):
+        if c in d.columns:
+            d[c] = -1.0
+    return d
+
+
+def _positions(df_geom, df_model):
+    """Vi tri cua tung dong df_model trong df_geom."""
+    g = pd.MultiIndex.from_frame(df_geom[["source_file", "cell_id"]])
+    m = pd.MultiIndex.from_frame(df_model[["source_file", "cell_id"]])
+    pos = g.get_indexer(m)
+    assert (pos >= 0).all(), "co dong model khong tim thay trong geom"
+    return pos
+
+
+def iterative_fit(name, geom, model_df, pos, fcols, y, le, cv, sw,
+                  mask_dir, n_est, seed, max_rounds, preset=False, groups=None):
+    df_cur, prev, hist = (geom if preset else _blank(geom)), None, []
+    mdl = scaler = None
+    for rnd in range(1, max_rounds + 1):
+        if rnd > 1:
+            df_cur = fill_neighbors(geom, prev, mask_dir)
+        Xg = df_cur[fcols].values.astype(np.float32)
+        Xm = Xg[pos]
+        scaler = StandardScaler().fit(Xm)
+        Xms = scaler.transform(Xm)
+
+        oof = np.zeros(len(y), dtype=int)
+        for tr, va in cv.split(Xms, y, groups):
+            f = _build_model(name, n_est, seed)
+            f.fit(Xms[tr], y[tr], sample_weight=sw[tr])
+            oof[va] = f.predict(Xms[va])
+        acc = accuracy_score(y, oof)
+        hist.append(round(float(acc), 4))
+
+        mdl = _build_model(name, n_est, seed)
+        mdl.fit(Xms, y, sample_weight=sw)
+
+        # nhan cho MOI te bao geom: co nhan -> oof (ngoai mau); con lai -> model vong nay
+        names = le.inverse_transform(mdl.predict(scaler.transform(Xg)))
+        names[pos] = le.inverse_transform(oof)
+
+        chg = -1 if prev is None else int((names != prev).sum())
+        print(f"      Round {rnd}: CV acc = {acc:.4f}" +
+              ("" if chg < 0 else f", {chg}/{len(names)} doi ({100*chg/len(names):.2f}%)"),
+              flush=True)
+        if prev is not None and chg == 0:
+            print("      Hoi tu!", flush=True)
+            break
+        prev = names
+    return mdl, scaler, hist, oof
+
+
+def iterative_predict(mdl, scaler, geom, fcols, le, mask_dir, max_rounds, preset=False):
+    df_cur, prev, hist = (geom if preset else _blank(geom)), None, []
+    names = None
+    for rnd in range(1, max_rounds + 1):
+        if rnd > 1:
+            df_cur = fill_neighbors(geom, prev, mask_dir)
+        names = le.inverse_transform(
+            mdl.predict(scaler.transform(df_cur[fcols].values.astype(np.float32))))
+        chg = -1 if prev is None else int((names != prev).sum())
+        hist.append(chg)
+        print(f"      Test round {rnd}" + ("" if chg < 0 else f": {chg}/{len(names)} doi"),
+              flush=True)
+        if prev is not None and chg == 0:
+            break
+        prev = names
+    return names, hist
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--features", required=True)
+    ap.add_argument("--mask-dir", required=True)
+    ap.add_argument("--tag", default="honest")
+    ap.add_argument("--model", default="LightGBM")
+    ap.add_argument("--n-estimators", type=int, default=500)
+    ap.add_argument("--max-rounds", type=int, default=10)
+    ap.add_argument("--seed", type=int, default=42)
+    ap.add_argument("--min-confidence", type=float, default=0.6)
+    ap.add_argument("--drop-neighbors", action="store_true")
+    ap.add_argument("--boost-outer", type=float, default=0.0,
+                    help="Nhan trong so mau len N lan cho cac class VONG NGOAI "
+                         "(epidermis, exodermis, endodermis, root_cap). Muc tieu: model "
+                         "uu tien dung o may vong don lop thay vi toi da accuracy tong.")
+    ap.add_argument("--boost-classes", default="epidermis,exodermis,endodermis,root_cap",
+                    help="danh sach class duoc boost, ngan cach bang dau phay")
+    ap.add_argument("--boost-spec", default="",
+                    help="He so RIENG cho tung class, vd "
+                         "'endodermis:20,exodermis:15,epidermis:10'. Neu dat thi "
+                         "no de len tren --boost-outer/--boost-classes.")
+    ap.add_argument("--nb-agg", action="store_true",
+                    help="Them feature tong hop tren TOAN BO hang xom.")
+    ap.add_argument("--nb-per-class", action="store_true",
+                    help="Mot feature cho MOI class: ty le hang xom thuoc tung "
+                         "loai trong 9 loai. Day du, khong gop nhom tuy tien.")
+    ap.add_argument("--nb-lean", action="store_true",
+                    help="Nhu --nb-agg nhung chi giu 2 cot thuc su moi.")
+    ap.add_argument("--mt-list", default="",
+                    help="File liet ke feature hinh thai duoc giu (1 ten/dong).")
+    ap.add_argument("--hp", default="",
+                    help="JSON hyperparameter de len tren mac dinh, vd "
+                         "'{\"num_leaves\": 127, \"learning_rate\": 0.03}'")
+    ap.add_argument("--wp-boost", type=float, default=15.0,
+                    help="He so nhan trong so cho xylem va phloem. Ban goc "
+                         "hardcode 15.0; dat 1.0 de bo han, chi giu 'balanced'.")
+    ap.add_argument("--group-cv", action="store_true",
+                    help="Chia CV theo SECTION (StratifiedGroupKFold) thay vi theo te bao. "
+                         "Quan trong: nhan hang xom bom vao vong lap se nhieu dung bang "
+                         "luc test (~0.81) thay vi qua sach (~0.96), nen train va test "
+                         "khop phan phoi.")
+    ap.add_argument("--neighbors-from-gt", action="store_true",
+                    help="CO CHU Y RO RI: dien hang xom tu NHAN THAT (kieu pipeline cu). "
+                         "Chi de do xem DINO dong gop bao nhieu phan ro ri, "
+                         "KHONG duoc dung lam so cong bo.")
+    a = ap.parse_args()
+    globals()["USE_NB_AGG"] = a.nb_agg or a.nb_lean
+    if a.nb_lean:
+        globals()["NB_AGG_COLS"] = ["frac_nb_vascular", "frac_nb_cortex"]
+    if a.nb_per_class:
+        # 9 ty le theo class + frac_nb_same. Cai thu 10 KHONG suy ra duoc tu 9
+        # cai kia: model khong nhan class cua chinh no lam dau vao, nen no
+        # khong biet cot nao la "cua minh".
+        globals()["NB_AGG_COLS"] = list(NB_PER_CLASS) + ["frac_nb_same"]
+        globals()["USE_NB_AGG"] = True
+    if a.mt_list:
+        globals()["MT_KEEP"] = {l.strip() for l in open(a.mt_list) if l.strip()}
+    if a.hp:
+        globals()["HP"] = json.loads(a.hp)
+        print(f"  hyperparameter de len tren: {HP}")
+    globals()["_build_model"] = _build_model_hp
+
+    geom, model_df, fcols = load_tables(a.features, a.min_confidence)
+    if a.drop_neighbors:
+        fcols = [c for c in fcols if c not in NEIGHBOR_CT_COLS]
+        a.max_rounds = 1
+        print(f"  bo 4 cot hang xom -> {len(fcols)} features (1 vong, khong lap)")
+
+    tr_s = {l.strip() for l in open(ROOT / "split_train.txt") if l.strip()}
+    te_s = {l.strip() for l in open(ROOT / "split_test.txt") if l.strip()}
+    assert not (tr_s & te_s), "split train/test chong nhau!"
+
+    g_tr = geom[geom.source_file.isin(tr_s)].reset_index(drop=True)
+    g_te = geom[geom.source_file.isin(te_s)].reset_index(drop=True)
+    m_tr = model_df[model_df.source_file.isin(tr_s)].reset_index(drop=True)
+    m_te = model_df[model_df.source_file.isin(te_s)].reset_index(drop=True)
+    p_tr, p_te = _positions(g_tr, m_tr), _positions(g_te, m_te)
+    print(f"  train {m_tr.source_file.nunique()} section: {len(g_tr):,} te bao "
+          f"({len(m_tr):,} co nhan)")
+    print(f"  test  {m_te.source_file.nunique()} section: {len(g_te):,} te bao "
+          f"({len(m_te):,} co nhan)", flush=True)
+
+    le = LabelEncoder().fit(sorted(model_df.cell_type.unique()))
+    y_tr = le.transform(m_tr.cell_type)
+    if a.group_cv:
+        cv = StratifiedGroupKFold(5, shuffle=True, random_state=a.seed)
+        groups = m_tr.source_file.values
+        print(f"  CV chia theo SECTION: {m_tr.source_file.nunique()} section -> 5 fold")
+    else:
+        cv = StratifiedKFold(5, shuffle=True, random_state=a.seed)
+        groups = None
+        print("  CV chia theo TE BAO (mac dinh)")
+    sw = _sample_weights(y_tr, le, a.wp_boost)
+    print(f"  trong so: balanced, xylem/phloem x{a.wp_boost}")
+    if a.boost_spec:
+        idx = {c: i for i, c in enumerate(le.classes_)}
+        sw = sw.copy()
+        for part in a.boost_spec.split(","):
+            cls, _, mult = part.partition(":")
+            cls, mult = cls.strip(), float(mult)
+            assert cls in idx, f"class khong co: {cls}"
+            m = y_tr == idx[cls]
+            sw[m] *= mult
+            print(f"  boost x{mult:g} cho {cls} ({int(m.sum()):,} te bao)")
+    elif a.boost_outer > 0:
+        OUTER = [c.strip() for c in a.boost_classes.split(",") if c.strip()]
+        idx = {c: i for i, c in enumerate(le.classes_)}
+        m = np.isin(y_tr, [idx[c] for c in OUTER if c in idx])
+        sw = sw.copy(); sw[m] *= a.boost_outer
+        print(f"  boost x{a.boost_outer} cho {OUTER} ({int(m.sum()):,} te bao)")
+
+    if a.neighbors_from_gt:
+        print("\n  *** CHE DO RO RI CO CHU Y: hang xom lay tu NHAN THAT ***", flush=True)
+        g_tr = fill_neighbors(g_tr, g_tr.cell_type.fillna("__none__").values, a.mask_dir)
+        g_te = fill_neighbors(g_te, g_te.cell_type.fillna("__none__").values, a.mask_dir)
+        a.max_rounds = 1
+
+    print("\n  --- train (hang xom = prediction, qua adjacency that) ---", flush=True)
+    mdl, scaler, hist, oof = iterative_fit(a.model, g_tr, m_tr, p_tr, fcols, y_tr, le,
+                                      cv, sw, a.mask_dir, a.n_estimators, a.seed,
+                                      a.max_rounds, preset=a.neighbors_from_gt,
+                                      groups=groups)
+
+    oof_df = m_tr[["source_file", "cell_id", "species", "stage", "cell_type"]].copy()
+    oof_df["predicted"] = le.inverse_transform(oof)
+    oof_df["correct"] = oof_df.cell_type.values == oof_df.predicted.values
+    oof_df.to_csv(OUT / f"oof_{a.tag}.csv", index=False)
+    print(f"  OOF train da ghi -> oof_{a.tag}.csv  (CV acc {oof_df.correct.mean():.4f})")
+
+    print("\n  --- cham diem 26 section (KHONG dung BMP lam feature) ---", flush=True)
+    names_all, thist = iterative_predict(mdl, scaler, g_te, fcols, le,
+                                         a.mask_dir, a.max_rounds,
+                                         preset=a.neighbors_from_gt)
+    pred = names_all[p_te]
+    acc = accuracy_score(m_te.cell_type.values, pred)
+
+    print(f"\n  >>> {a.tag}: TEST = {acc:.4f}", flush=True)
+    print(f"  >>> vong TRAIN den hoi tu: {len(hist)}  (CV moi vong: {hist})", flush=True)
+    print(f"  >>> vong TEST  den hoi tu: {len(thist)} (so cell doi: {thist})", flush=True)
+
+    joblib.dump(mdl, OUT / f"model_{a.tag}.joblib")
+    joblib.dump(scaler, OUT / f"scaler_{a.tag}.joblib")
+    joblib.dump(le, OUT / f"labelencoder_{a.tag}.joblib")
+    o = m_te[["source_file", "cell_id", "species", "stage", "cell_type"]].copy()
+    o["predicted"], o["correct"] = pred, m_te.cell_type.values == pred
+    o.to_csv(OUT / f"test_predictions_{a.tag}.csv", index=False)
+    (o.groupby("source_file").agg(n_cells=("correct", "size"), accuracy=("correct", "mean"))
+       .round(4).to_csv(OUT / f"per_image_accuracy_{a.tag}.csv"))
+    pc = o.groupby("cell_type")["correct"].agg(["size", "mean"]).round(4)
+    print("\n  Accuracy tung loai:"); print(pc.to_string())
+    pc.to_csv(OUT / f"per_class_{a.tag}.csv")
+
+
+if __name__ == "__main__":
+    main()

rootscope/api.py

@@ -15,53 +15,88 @@ from pathlib import Path
 import pandas as pd
 
 from . import predict as _predict
-from .weights import resolve_cnn_weights, resolve_model_dir
+from .weights import is_v4, resolve_cnn_weights, resolve_model_dir
 
 
 def _load(model_dir=None, cnn_weights=None):
+    """Resolve weights for the selected model version.
+
+    v4 returns ``(model_dir, None, None, None, cnn)``: its classifiers are
+    loaded inside the v4 stage (one bundle per seed, each carrying its own
+    scaler and feature list), so there is nothing to pre-load here.
+    """
     mdir = resolve_model_dir(model_dir)
     cnn = resolve_cnn_weights(cnn_weights)
+    if is_v4():
+        return mdir, None, None, None, cnn
     models_dict, scalers, feature_cols, le = _predict.load_models(str(mdir))
     if not models_dict:
         raise RuntimeError(f"No usable models found in {mdir}.")
     return models_dict, scalers, feature_cols, le, cnn
 
 
+def _gpu_available():
+    try:
+        import torch
+        return bool(torch.cuda.is_available())
+    except Exception:  # noqa: BLE001
+        return False
+
+
+def _run_one(tif, loaded, out_dir, gpu, um_per_px, max_rounds, label_cells):
+    if gpu is None:                      # library default: use a GPU if there is one
+        gpu = _gpu_available()
+    models_dict, scalers, feature_cols, le, cnn = loaded
+    if is_v4():
+        return _predict.predict_single_tif_v4(
+            str(tif), model_dir=str(models_dict), um_per_px=um_per_px, gpu=gpu,
+            out_dir=str(out_dir), max_rounds=max_rounds,
+            cnn_weights=str(cnn) if cnn else None, label_cells=label_cells,
+        )
+    return _predict.predict_single_tif(
+        str(tif), models_dict, scalers, feature_cols, le,
+        um_per_px=um_per_px, gpu=gpu, out_dir=str(out_dir),
+        max_rounds=max_rounds, cnn_weights=str(cnn) if cnn else None,
+        label_cells=label_cells,
+    )
+
+
 def predict_tif(
     tif,
     out_dir="results",
-    gpu=True,
+    gpu=None,
     model_dir=None,
     cnn_weights=None,
-    um_per_px=1.0,
-    max_rounds=10,
+    um_per_px=None,
+    max_rounds=None,
     label_cells=False,
 ):
-    """Segment + predict cell types for a single TIFF. Returns a DataFrame."""
-    models_dict, scalers, feature_cols, le, cnn = _load(model_dir, cnn_weights)
-    df = _predict.predict_single_tif(
-        str(tif), models_dict, scalers, feature_cols, le,
-        um_per_px=um_per_px, gpu=gpu, out_dir=str(out_dir),
-        max_rounds=max_rounds, cnn_weights=str(cnn) if cnn else None,
-        label_cells=label_cells,
-    )
-    return df
+    """Segment + predict cell types for a single TIFF. Returns a DataFrame.
+
+    ``um_per_px=None`` reads the pixel size from the TIFF's OME metadata and
+    fails loudly if it is absent -- a wrong scale distorts every size feature.
+    ``max_rounds=None`` uses the version's own setting (6 for v4, 10 for v2).
+    """
+    loaded = _load(model_dir, cnn_weights)
+    return _run_one(tif, loaded, out_dir, gpu, um_per_px,
+                    _predict.default_rounds(max_rounds), label_cells)
 
 
 def predict_folder(
     tif_dir,
     out_dir="results",
-    gpu=True,
+    gpu=None,
     model_dir=None,
     cnn_weights=None,
-    um_per_px=1.0,
-    max_rounds=10,
+    um_per_px=None,
+    max_rounds=None,
     label_cells=False,
     pattern="*.tif",
 ):
     """Segment + predict for every TIFF in a folder. Returns a combined
     DataFrame and writes ``all_predictions.csv`` into ``out_dir``."""
-    models_dict, scalers, feature_cols, le, cnn = _load(model_dir, cnn_weights)
+    loaded = _load(model_dir, cnn_weights)
+    rounds = _predict.default_rounds(max_rounds)
     tif_paths = sorted(Path(tif_dir).glob(pattern))
     if not tif_paths:
         raise FileNotFoundError(f"No files matching {pattern} in {tif_dir}")
@@ -69,12 +104,7 @@ def predict_folder(
     tables = []
     for tp in tif_paths:
         try:
-            df = _predict.predict_single_tif(
-                str(tp), models_dict, scalers, feature_cols, le,
-                um_per_px=um_per_px, gpu=gpu, out_dir=str(out_dir),
-                max_rounds=max_rounds, cnn_weights=str(cnn) if cnn else None,
-                label_cells=label_cells,
-            )
+            df = _run_one(tp, loaded, out_dir, gpu, um_per_px, rounds, label_cells)
             if df is not None:
                 tables.append(df)
         except Exception as e:  # noqa: BLE001

rootscope/classify_v4.py

@@ -0,0 +1,106 @@
+"""RootScope v4 -- tang phan loai.
+
+Khac v2 o ba diem, va ca ba deu nam o day chu khong o phan trich dac trung:
+
+  1. Hang xom dang MEM: thay vi chot nhan cua hang xom roi dua vao model, moi huong
+     dong gop ca vector xac suat 9 lop (36 cot) + 6 co dong thuan.
+  2. Dac trung LOP: khoang cach tu moi te bao toi cac lop mo (epidermis, exodermis,
+     endodermis, pericycle) ma chinh model da du doan o vong truoc, do theo ban kinh
+     cuc bo va tinh bang duong kinh te bao.
+  3. Radial-order prior: uu tien thu tu xuyen tam dung giai phau, ap sau vong cuoi.
+
+Cac ham tinh dac trung nam o context.py va duoc CAT NGUYEN VAN tu code nghien cuu.
+"""
+from pathlib import Path
+
+import joblib
+import numpy as np
+import pandas as pd
+
+from . import context as ctx
+from . import _eval_honest as _eh
+from .extract_features import build_cell_adjacency
+from .radial_order import correct
+
+CLASSES = ["cortex", "endodermis", "epidermis", "exodermis", "pericycle",
+           "phloem", "root_cap", "stele", "xylem"]
+PRIOR = (0.3, 2.0)        # (s, lambda) cua model san xuat
+
+
+def load_models(model_dir, seeds=None):
+    """Nap cac goi LightGBM v4. Moi goi chua round_models = model cua tung vong."""
+    model_dir = Path(model_dir)
+    out = []
+    for p in sorted(model_dir.glob("lgbm_s*.joblib")):
+        s = int(p.stem.split("_s")[-1])
+        if seeds is None or s in seeds:
+            b = joblib.load(p)
+            if "round_models" not in b:
+                raise ValueError(
+                    f"{p.name} khong co 'round_models'. Goi nay chi luu model vong cuoi "
+                    f"nen khong chay lai duoc vong lap. Dung goi san xuat.")
+            out.append((s, b))
+    if not out:
+        raise FileNotFoundError(f"khong co lgbm_s*.joblib trong {model_dir}")
+    f0 = list(out[0][1]["fcols"])
+    for s, b in out:
+        if list(b["fcols"]) != f0:
+            raise ValueError(f"seed {s} co bo feature khac")
+        if list(b["classes"]) != CLASSES:
+            raise ValueError(f"seed {s} co thu tu lop khac")
+    return out
+
+
+def classify(df_base, masks, source_file="image.tif", model_dir=None, seeds=None,
+             max_rounds=6, prior=PRIOR, verbose=True):
+    """Tra (nhan, xac suat, nhan vong 1) cho tung dong cua df_base."""
+    bundles = load_models(model_dir, seeds)
+    fcols = list(bundles[0][1]["fcols"])
+
+    geom = df_base.copy()
+    geom["source_file"] = source_file
+    geom = geom.reset_index(drop=True)
+
+    # Nap thang adjacency vao cache cua _eval_honest: khoi phai ghi mask ra dia
+    # roi doc lai nhu duong danh gia trong du an.
+    _eh._ADJ[source_file] = build_cell_adjacency(masks)
+    ctx.MASK = "<in-memory>"
+    ctx.CL = CLASSES
+
+    missing = [c for c in fcols if c not in geom.columns
+               and c not in ctx.RINGCOLS and not c.startswith(("nbp_", "agree_"))]
+    if missing:
+        raise ValueError(f"thieu {len(missing)} cot dac trung, vd {missing[:6]}")
+
+    mats = ctx.dir_mats(geom)
+    P1s, PFs = [], []
+    for s, b in bundles:
+        prev = prevP = None
+        for rnd in range(1, max_rounds + 1):
+            cur = ctx._blank(geom) if rnd == 1 else ctx.fill_neighbors(geom, prev, ctx.MASK)
+            cur = pd.concat([cur.reset_index(drop=True),
+                             ctx.soft_feats(prevP, mats, CLASSES),
+                             ctx.ring_feats(geom, prev, "q", prevP)], axis=1)
+            m_r, s_r = b["round_models"][min(rnd, len(b["round_models"])) - 1]
+            P = m_r.predict_proba(s_r.transform(cur[fcols].values.astype(np.float32)))
+            if rnd == 1:
+                P1s.append(P.astype(np.float64))
+            names = np.array(CLASSES)[P.argmax(1)]
+            chg = -1 if prev is None else int((names != prev).sum())
+            if verbose:
+                print(f"    seed {s} round {rnd}: {chg} labels changed", flush=True)
+            prev, prevP = names, P
+            if chg == 0:
+                break
+        PFs.append(P.astype(np.float64))
+
+    P1 = np.mean(P1s, 0)
+    P = np.mean(PFs, 0)
+    if prior and prior[1]:
+        cy = geom.centroid_y.values.astype(float)
+        cx = geom.centroid_x.values.astype(float)
+        cid = geom.cell_id.values
+        look = {int(i): (y, x) for i, y, x in zip(cid, cy, cx)}
+        P = correct(P, np.full(len(geom), source_file), cid, CLASSES,
+                    geom={source_file: (cy, cx, look)}, s=prior[0], lam=prior[1])
+    return np.array(CLASSES)[P.argmax(1)], P, np.array(CLASSES)[P1.argmax(1)]

rootscope/cli.py

@@ -40,10 +40,15 @@ def build_parser() -> argparse.ArgumentParser:
 
     p.add_argument("--out", "--out-dir", dest="out_dir", default="results",
                    help="Output folder for CSVs + overlays (default: results/).")
-    p.add_argument("--um-per-px", type=float, default=1.0,
-                   help="Microns per pixel (default: 1.0).")
-    p.add_argument("--max-rounds", type=int, default=10,
-                   help="Max iterative refinement rounds (default: 10).")
+    p.add_argument("--um-per-px", type=float, default=None,
+                   help="Microns per pixel. Default: read from the TIFF's OME "
+                        "metadata (PhysicalSizeX); the run stops if it is "
+                        "missing, because a wrong scale distorts every "
+                        "size feature.")
+    p.add_argument("--max-rounds", type=int, default=None,
+                   help="Max iterative refinement rounds (default: 6 for v4, "
+                        "10 for v1/v2 -- the settings each model was "
+                        "evaluated with).")
     p.add_argument("--label-cells", action="store_true",
                    help="Write the cell-type name inside each cell on the "
                         "overlay. Off by default (the text overlaps on dense "
@@ -55,6 +60,14 @@ def build_parser() -> argparse.ArgumentParser:
     g.add_argument("--cpu", dest="gpu", action="store_false",
                    help="Force CPU (slow).")
 
+    p.add_argument("--model-version", choices=("v1", "v2", "v4"), default=None,
+                   help="Which published model to use. v4 (default): the "
+                        "manuscript model -- 915 features (864 + 51 rebuilt "
+                        "every refinement round), one LightGBM per round, "
+                        "3 seeds averaged, radial-order prior; 0.881 accuracy "
+                        "on the clean held-out split. v2: 864 features, "
+                        "RF/XGB/LGBM ensemble, 10 rounds. v1: 480 features, "
+                        "DINOv2 ViT-S/14.")
     p.add_argument("--model-dir", default=None,
                    help="Folder with trained model weights (default: "
                         "auto-download/cache).")
@@ -69,7 +82,12 @@ def build_parser() -> argparse.ArgumentParser:
 
 
 def main(argv=None):
+    import os as _os
     args = build_parser().parse_args(argv)
+    # set before anything resolves weights, so the classifiers and the backbone
+    # are taken from the same version's folder
+    if args.model_version:
+        _os.environ["ROOTSCOPE_MODEL_VERSION"] = args.model_version
     use_gpu = _gpu_available() if args.gpu is None else args.gpu
     if args.gpu is None:
         print(f"[rootscope] GPU {'detected' if use_gpu else 'not found (using CPU)'}.")
@@ -83,7 +101,7 @@ def main(argv=None):
                 label_cells=args.label_cells,
             )
             n = 0 if df is None else len(df)
-            print(f"\n[rootscope] Done — {n} cell predictions written to {args.out_dir}/")
+            print(f"\n[rootscope] Done. {n} cell predictions written to {args.out_dir}/")
         else:
             df = api.predict_folder(
                 args.tif_dir, out_dir=args.out_dir, gpu=use_gpu,
@@ -92,7 +110,7 @@ def main(argv=None):
                 label_cells=args.label_cells, pattern=args.pattern,
             )
             n = 0 if df is None else len(df)
-            print(f"\n[rootscope] Done — {n} total predictions written to {args.out_dir}/")
+            print(f"\n[rootscope] Done. {n} total predictions written to {args.out_dir}/")
     except (FileNotFoundError, RuntimeError) as e:
         print(f"\n[rootscope] ERROR: {e}", file=sys.stderr)
         sys.exit(1)

rootscope/context.py

@@ -0,0 +1,119 @@
+"""Dac trung ngu canh cua RootScope v4:
+  - hang xom bon huong, dang MEM (xac suat) thay vi nhan cung
+  - khoang cach toi cac LOP mo ma chinh model da du doan o vong truoc
+
+CAT NGUYEN VAN tu Working/_figgrid/eval_grid.py (2026-09-21). Khong go lai bang tay:
+mot sai lech nho o day cho ket qua sai ma khong bao loi nao.
+"""
+import os
+import numpy as np, pandas as pd, scipy.sparse as sp
+from ._eval_honest import _adjacency_for, fill_neighbors, _blank   # noqa: F401
+
+MASK = None        # thu muc chua <stem>_masks.npy; predict.py dat truoc khi goi
+CL = None          # danh sach ten lop; predict.py dat truoc khi goi
+
+DIRS = ["inner", "outer", "cw", "ccw"]
+
+
+def dir_mats(g):
+    """4 ma tran thua (n x n), chuan hoa theo hang: hang xom theo tung huong cuc."""
+    r, c = {k: [] for k in DIRS}, {k: [] for k in DIRS}
+    for sf, sub in g.groupby("source_file"):
+        adj = _adjacency_for(sf, MASK)
+        row = {int(i): p for i, p in zip(sub.cell_id.values, sub.index.values)}
+        cy0, cx0 = sub.centroid_y.mean(), sub.centroid_x.mean()
+        cen = {int(i): (y, x) for i, y, x in zip(sub.cell_id, sub.centroid_y, sub.centroid_x)}
+        for cid, (cy, cx) in cen.items():
+            md = np.hypot(cy - cy0, cx - cx0); ma = np.arctan2(cy - cy0, cx - cx0)
+            for n in adj.get(cid, ()):
+                if n not in cen:
+                    continue
+                ny, nx = cen[n]
+                dr = np.hypot(ny - cy0, nx - cx0) - md
+                da = (np.arctan2(ny - cy0, nx - cx0) - ma + np.pi) % (2 * np.pi) - np.pi
+                arc = abs(da) * md if md > 1.0 else abs(da)
+                k = ("inner" if dr < 0 else "outer") if abs(dr) >= arc else ("cw" if da < 0 else "ccw")
+                r[k].append(row[cid]); c[k].append(row[n])
+    n = len(g); mats = {}
+    for k in DIRS:
+        A = sp.csr_matrix((np.ones(len(r[k])), (r[k], c[k])), shape=(n, n))
+        deg = np.asarray(A.sum(1)).ravel()
+        mats[k] = (sp.diags(1 / np.maximum(deg, 1)) @ A, deg > 0)
+    return mats
+
+
+def soft_feats(P, mats, classes):
+    n = len(mats["cw"][1]); cols = {}
+    if P is None:
+        names = [f"nbp_{k}_{c}" for k in DIRS for c in classes] + [f"agree_{k}" for k in DIRS] + ["agree_cw_ccw", "agree_in_out"]
+        return pd.DataFrame(-1.0, index=range(n), columns=names)
+    M = {}
+    for k in DIRS:
+        A, has = mats[k]; M[k] = A @ P
+        for j, c in enumerate(classes):
+            cols[f"nbp_{k}_{c}"] = np.where(has, M[k][:, j], -1.0)
+    for k in DIRS:
+        cols[f"agree_{k}"] = np.where(mats[k][1], (P * M[k]).sum(1), -1.0)
+    cols["agree_cw_ccw"] = np.where(mats["cw"][1] & mats["ccw"][1], (M["cw"] * M["ccw"]).sum(1), -1.0)
+    cols["agree_in_out"] = np.where(mats["inner"][1] & mats["outer"][1], (M["inner"] * M["outer"]).sum(1), -1.0)
+    return pd.DataFrame(cols)
+
+
+
+LAND = ["epidermis", "exodermis", "endodermis", "pericycle"]
+RINGCOLS = ["depth_surf", "frac_r_surf"] + [f"d_{c}" for c in LAND] + ["d_vasc_own", "d_endo_own", "ratio_endo"]
+RING = int(os.environ.get("RINGFEAT", "1"))
+PERROUND = int(os.environ.get("PERROUND", "1"))   # inference uses the model trained FOR each round
+SOFTRING = int(os.environ.get("SOFTRING", "0"))   # landmarks weighted by probability instead of argmax
+DAMP = float(os.environ.get("DAMP", "0"))          # P_used = DAMP*P_prev + (1-DAMP)*P_new
+_STATIC = {}
+def _adiff(a, b): return np.abs((a[:, None] - b[None, :] + np.pi) % (2 * np.pi) - np.pi)
+def ring_feats(g, names, key, Pm=None):
+    """Where is each cell relative to landmarks of ITS OWN section?  Static: the tissue surface.  Dynamic (from the
+    previous round's predictions; out-of-fold during training): the predicted epidermis / exodermis / endodermis /
+    pericycle rings and the outer envelope of the vascular cylinder.  Distances are in cell diameters, measured along
+    the local radius, so a non-circular root or an off-centre stele does not matter."""
+    out = np.full((len(g), len(RINGCOLS)), np.nan, dtype=np.float32)
+    if not RING: return pd.DataFrame(out, columns=RINGCOLS)
+    if key not in _STATIC:
+        st = np.full((len(g), 2), np.nan, dtype=np.float32)
+        for sf, sub in g.groupby("source_file"):
+            idx = sub.index.values; cy, cx = sub.centroid_y.mean(), sub.centroid_x.mean()
+            r = np.hypot(sub.centroid_y - cy, sub.centroid_x - cx).values; th = np.arctan2(sub.centroid_y - cy, sub.centroid_x - cx).values
+            near = _adiff(th, th) < np.deg2rad(15); rs = np.where(near, r[None, :], -1).max(1); dm = np.sqrt(np.median(sub.area_px))
+            st[idx, 0] = (rs - r) / dm; st[idx, 1] = r / np.maximum(rs, 1)
+        _STATIC[key] = st
+    out[:, :2] = _STATIC[key]
+    if names is None: return pd.DataFrame(out, columns=RINGCOLS)
+    names = np.asarray(names)
+    for sf, sub in g.groupby("source_file"):
+        idx = sub.index.values; nm = names[idx]; own = np.sqrt(sub.area_px.values)
+        core = np.isin(nm, ["endodermis", "pericycle", "stele", "xylem", "phloem"])
+        cy, cx = (sub.centroid_y[core].mean(), sub.centroid_x[core].mean()) if core.sum() >= 5 else (sub.centroid_y.mean(), sub.centroid_x.mean())
+        r = np.hypot(sub.centroid_y - cy, sub.centroid_x - cx).values; th = np.arctan2(sub.centroid_y - cy, sub.centroid_x - cx).values
+        def ring(sel, how="median", win=40):
+            m = _adiff(th, th[sel]) < np.deg2rad(win); A = np.where(m, r[sel][None, :], np.nan)
+            v = np.nanmedian(A, 1) if how == "median" else np.nanmax(A, 1)
+            fb = np.median(r[sel]) if how == "median" else np.percentile(r[sel], 90)
+            return np.where(m.sum(1) >= 2, v, fb)
+        import warnings
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")
+            for j, c in enumerate(LAND):
+                if SOFTRING and Pm is not None:
+                    w = Pm[idx, CL.index(c)]; sel = w > 0.2
+                    if w[sel].sum() < 4: continue
+                    ww = w[sel] ** 2; mwin = (_adiff(th, th[sel]) < np.deg2rad(40)) * ww[None, :]
+                    tot = mwin.sum(1); glob = np.average(r[sel], weights=ww)
+                    rc = np.where(tot > ww.max() * 1.5, (mwin * r[sel][None, :]).sum(1) / np.maximum(tot, 1e-9), glob)
+                    dm = np.sqrt(np.average(sub.area_px.values[sel], weights=ww))
+                else:
+                    sel = nm == c
+                    if sel.sum() < 6: continue
+                    rc = ring(sel); dm = np.sqrt(np.median(sub.area_px.values[sel]))
+                out[idx, 2 + j] = (r - rc) / dm
+                if c == "endodermis": out[idx, 7] = (r - rc) / own; out[idx, 8] = r / np.maximum(rc, 1)
+            sel = np.isin(nm, ["stele", "xylem", "phloem"])
+            if sel.sum() >= 5: out[idx, 6] = (r - ring(sel, "max", 30)) / own
+    return pd.DataFrame(out, columns=RINGCOLS)
+

rootscope/extract_features.py

@@ -675,7 +675,7 @@ def extract_all_features(masks, img_rgb, um_per_px, layer_lookup, adjacency,
         #    have fewer cells than thick cortex
         cells_in_same_layer = layer_cell_counts.get(my_layer, 0) if my_layer >= 0 else 0
 
-        # 3. Neighbors inward vs outward — ring cells have neighbors
+        # 3. Neighbors inward vs outward, ring cells have neighbors
         #    on both sides; cortex cells mostly have same-layer neighbors
         inner_neighbor_count = 0
         outer_neighbor_count = 0
@@ -725,7 +725,7 @@ def extract_all_features(masks, img_rgb, um_per_px, layer_lookup, adjacency,
         #    Exodermis = YES (neighbors epidermis), cortex/endodermis = NO
         cell_neighbors_boundary = 1 if cid in neighbors_boundary else 0
 
-        # 8. Min layer among neighbors — exodermis neighbors include
+        # 8. Min layer among neighbors, exodermis neighbors include
         #    epidermis (layer 0-1), cortex neighbors are all mid-layers
         min_neighbor_layer = min(neighbor_layers) if neighbor_layers else -1
         max_neighbor_layer = max(neighbor_layers) if neighbor_layers else -1
@@ -851,7 +851,7 @@ def extract_all_features(masks, img_rgb, um_per_px, layer_lookup, adjacency,
         HEXAGON_COMPACTNESS = 0.9069
         hexagonality = round(1.0 - abs(compactness - HEXAGON_COMPACTNESS), 4)
 
-        # 5. Convex hull vertex count — hexagons have ~6 vertices
+        # 5. Convex hull vertex count, hexagons have ~6 vertices
         try:
             from skimage.measure import approximate_polygon
             coords = p.coords  # (N, 2) array of pixel coordinates
@@ -1266,7 +1266,7 @@ def extract_cnn_embedding_features(masks, img_rgb, use_gpu=True, weights_path=No
                                       weights_path=weights_path, model=model)
     except ImportError:
         import warnings
-        warnings.warn("cnn_embeddings module not found — skipping CNN features.")
+        warnings.warn("cnn_embeddings module not found, skipping CNN features.")
         return None
     except Exception as e:
         import warnings

rootscope/predict.py

@@ -576,18 +576,24 @@ def load_image(tif_path):
 #
 # predict_single_tif() below runs these in order. They are also exposed
 # separately so a caller that pays for GPU time by the second (a web app on
-# ZeroGPU, say) can wrap only stage_segment and stage_embed — the two steps
-# that actually touch the GPU — and run the rest on CPU.
+# ZeroGPU, say) can wrap only stage_segment and stage_embed, the two steps
+# that actually touch the GPU, and run the rest on CPU.
 # ═══════════════════════════════════════════════════════════════
 
 def stage_segment(img_rgb, gpu=True, cellpose_model=None):
-    """GPU stage 1 — Cellpose-SAM. Returns the int32 label mask."""
+    """GPU stage 1: Cellpose-SAM. Returns the int32 label mask."""
     return segment_cellpose_sam(img_rgb, use_gpu=gpu, model=cellpose_model)
 
 
-def stage_features(masks, img_rgb, um_per_px=1.0, verbose=True):
-    """CPU stage — tissue mask, layer index, debris removal, handcrafted
-    features.
+def stage_features(masks, img_rgb, um_per_px=1.0, verbose=True,
+                   remove_debris=True):
+    """CPU stage: tissue mask, layer index, (optional) debris removal,
+    handcrafted features.
+
+    ``remove_debris`` drops segmented objects that lie mostly outside the
+    tissue body. v1/v2 always did this. v4 passes False: the manuscript's
+    evaluation kept every Cellpose object, and the v4 package reproduces that
+    evaluation cell for cell only when the same cells are scored.
 
     Debris removal edits ``masks`` in place, so the (possibly modified) mask is
     returned alongside the feature table.
@@ -605,7 +611,7 @@ def stage_features(masks, img_rgb, um_per_px=1.0, verbose=True):
     # Filter out debris cells outside the main tissue body
     n_cells = int(masks.max())
     debris_cells = set()
-    for cid in range(1, n_cells + 1):
+    for cid in (range(1, n_cells + 1) if remove_debris else ()):
         cell_px = masks == cid
         n_total = cell_px.sum()
         if n_total == 0:
@@ -634,7 +640,7 @@ def stage_features(masks, img_rgb, um_per_px=1.0, verbose=True):
 
 def stage_embed(masks, img_rgb, df_base, gpu=True, cnn_weights=None,
                 dinov2_model=None, verbose=True):
-    """GPU stage 2 — fine-tuned DINOv2 per-cell embeddings, merged into
+    """GPU stage 2: fine-tuned DINOv2 per-cell embeddings, merged into
     ``df_base``. Returns the merged table (unchanged if embeddings are
     unavailable)."""
     if verbose:
@@ -653,7 +659,7 @@ def stage_classify(df_base, masks, img_rgb, layer_lookup, adjacency,
                    models_dict, scalers, feature_cols, le,
                    out_dir="predictions", stem="image", source_name=None,
                    um_per_px=1.0, max_rounds=10, label_cells=False):
-    """CPU stage — iterative prediction with each model plus the weighted
+    """CPU stage: iterative prediction with each model plus the weighted
     ensemble, anatomical post-processing, and per-model CSV + overlay PNG.
 
     Returns the concatenated per-cell table across all models."""
@@ -828,6 +834,133 @@ def stage_classify(df_base, masks, img_rgb, layer_lookup, adjacency,
     return pd.concat(all_dfs, ignore_index=True)
 
 
+
+# ═══════════════════════════════════════════════════════════════
+# V4
+#
+# Same stage_segment / stage_features / stage_embed as above; only the
+# classifier differs. The v4 loop lives in classify_v4.py and the context
+# features it rebuilds every round in context.py (copied verbatim from the
+# research code). These wrappers exist so api.py, cli.py and the web app can
+# treat v4 like the other versions.
+# ═══════════════════════════════════════════════════════════════
+
+def default_rounds(user_value=None):
+    """The refinement-round cap each model version was evaluated with."""
+    if user_value is not None:
+        return int(user_value)
+    from .weights import is_v4
+    return 6 if is_v4() else 10
+
+
+def read_um_per_px(tif_path):
+    """Pixel size in microns from the TIFF's OME metadata, or None."""
+    import re
+    try:
+        import tifffile
+        with tifffile.TiffFile(str(tif_path)) as t:
+            ome = t.ome_metadata or ""
+    except Exception:  # noqa: BLE001
+        return None
+    m = re.search(r'PhysicalSizeX="([0-9.eE+-]+)"', ome)
+    if not m:
+        return None
+    v = float(m.group(1))
+    unit = re.search(r'PhysicalSizeXUnit="([^"]+)"', ome)
+    if unit and unit.group(1) not in ("µm", "um", "\u00b5m", "micron", "microns"):
+        print(f"    WARNING: PhysicalSizeXUnit is '{unit.group(1)}', assuming microns.")
+    return v
+
+
+def stage_classify_v4(df_base, masks, img_rgb, model_dir, out_dir="predictions",
+                      stem="image", source_name=None, um_per_px=1.0,
+                      max_rounds=6, seeds=None, label_cells=False, verbose=True):
+    """CPU stage, v4: per-round LightGBM x 3 seeds, soft-neighbor + layer
+    context rebuilt each round, radial-order prior. Writes
+    ``<stem>_LightGBM_predictions.csv`` and ``<stem>_LightGBM_overlay.png``
+    and returns the per-cell table in the same shape as stage_classify."""
+    from sklearn.preprocessing import LabelEncoder
+    from .classify_v4 import CLASSES, classify
+
+    out_dir = Path(out_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    if source_name is None:
+        source_name = f"{stem}.tif"
+
+    print(f"\n    --- LightGBM v4 (per-round models, max {max_rounds} rounds, "
+          f"{'all seeds' if seeds is None else seeds}) ---")
+    names, P, names_r1 = classify(df_base, masks, source_file=source_name,
+                                  model_dir=model_dir, seeds=seeds,
+                                  max_rounds=max_rounds, verbose=verbose)
+
+    df = df_base.copy().reset_index(drop=True)
+    df["predicted_cell_type"] = names
+    df["prediction_confidence"] = np.round(P.max(axis=1), 4)
+    df["predicted_round1"] = names_r1
+    for i, c in enumerate(CLASSES):
+        df[f"p_{c}"] = P[:, i]
+    df["source_file"] = source_name
+    df["um_per_px"] = um_per_px
+    df["model"] = "LightGBM"
+
+    print("    Final predictions:")
+    for ct, cnt in pd.Series(names).value_counts().items():
+        print(f"      {ct}: {cnt}")
+    print(f"    Labels changed by refinement: {int((names != names_r1).sum())} / {len(names)}")
+
+    csv_path = out_dir / f"{stem}_LightGBM_predictions.csv"
+    df.to_csv(csv_path, index=False)
+    print(f"    Saved: {csv_path}")
+
+    le = LabelEncoder().fit(CLASSES)          # only classes_ is used by the overlay
+    cell_predictions = {int(c): str(n) for c, n in zip(df["cell_id"].values, names)}
+    overlay_path = out_dir / f"{stem}_LightGBM_overlay.png"
+    save_prediction_overlay(img_rgb, masks, cell_predictions, overlay_path, le,
+                            label_cells=label_cells)
+    print(f"    Saved: {overlay_path}")
+    return df
+
+
+def predict_single_tif_v4(tif_path, model_dir, um_per_px=None, gpu=True,
+                          out_dir="predictions", max_rounds=6, cnn_weights=None,
+                          label_cells=False, cellpose_model=None,
+                          dinov2_model=None, seeds=None):
+    """Full v4 pipeline: segment -> features (no debris filter) -> embed ->
+    v4 classify -> overlay."""
+    out_dir = Path(out_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    stem = Path(tif_path).stem.replace(".aivia", "")
+    print(f"\n  Processing: {tif_path}")
+
+    if um_per_px is None:
+        um_per_px = read_um_per_px(tif_path)
+        if um_per_px is None:
+            raise ValueError(
+                f"{Path(tif_path).name}: no PhysicalSizeX in the OME metadata. "
+                f"Pass um_per_px (CLI: --um-per-px). RootScope will not assume 1.0: "
+                f"a wrong scale distorts every size feature.")
+        print(f"    Pixel size from OME metadata: {um_per_px} um/px")
+
+    img_rgb = load_image(tif_path)
+    print(f"    Image shape: {img_rgb.shape}")
+    print("    Segmenting (Cellpose-SAM)...")
+    masks = stage_segment(img_rgb, gpu=gpu, cellpose_model=cellpose_model)
+    n_cells = int(masks.max())
+    print(f"    {n_cells} cells found")
+    if n_cells == 0:
+        print("    No cells found, skipping.")
+        return None
+
+    masks, df_base, layer_lookup, adjacency, n_layers = stage_features(
+        masks, img_rgb, um_per_px=um_per_px, remove_debris=False)
+    df_base = stage_embed(masks, img_rgb, df_base, gpu=gpu,
+                          cnn_weights=cnn_weights, dinov2_model=dinov2_model)
+    return stage_classify_v4(
+        df_base, masks, img_rgb, model_dir=model_dir, out_dir=out_dir, stem=stem,
+        source_name=Path(tif_path).name, um_per_px=um_per_px,
+        max_rounds=max_rounds, seeds=seeds, label_cells=label_cells)
+
+
 def predict_single_tif(tif_path, models_dict, scalers, feature_cols, le,
                         um_per_px=1.0, gpu=True, out_dir="predictions",
                         max_rounds=10, cnn_weights=None, label_cells=False,
@@ -921,7 +1054,7 @@ def main():
         print(f"\n  Found {len(tif_paths)} TIF files in {tif_dir}")
 
     print(f"\n[2] Predicting cell types with {len(models_dict)} models "
-          f"({', '.join(models_dict.keys())}) — iterative, max {args.max_rounds} rounds...")
+          f"({', '.join(models_dict.keys())}), iterative, max {args.max_rounds} rounds...")
     all_tables = []
     for tif_path in tif_paths:
         try:

rootscope/radial_order.py

@@ -0,0 +1,108 @@
+"""Soft radial-order constraint  (fix for error family B: one-layer boundary shifts).
+
+POST-HOC.  Takes the model's probability matrix and the cell centroids; it never looks at
+the 768 DINOv2 columns, so it can be re-applied unchanged after the re-fine-tune.
+
+Method
+  1. Per section, a monotone radial coordinate r_rel in [0,1]:  r_rel = r / r_surface(theta),
+     centre = mean centroid of the cells PREDICTED vascular-core, r_surface = furthest cell
+     within +-15 deg of the cell's own angle (all cells of the section, labels not needed).
+  2. Tissues get a radial RANK: epidermis/root_cap 0 < exodermis 1 < cortex 2 < endodermis 3
+     < pericycle 4 < {stele,xylem,phloem} 5.  The trio shares ONE rank - there is no reliable
+     radial order inside it, so the constraint cannot touch stele/xylem/phloem against each other.
+  3. Per section, an ordered-logit model of rank against -r_rel is fitted to the model's OWN
+     posterior (soft targets, no labels, nothing from the ground truth).  Two passes: the second
+     fit uses the once-corrected posterior.
+  4. P_final  ~  P_model * P_order ** lambda,  renormalised.
+
+Hyperparameters (s = 0.3, lambda = 2.0, 2 passes) were chosen on the 100 out-of-fold training
+sections only; the OOF optimum is a broad plateau (s 0.3-1.2, 2-4 passes, lambda 2).
+
+Usage
+  from apply_radial_order import correct
+  P_final = correct(P, source_file, cell_id, classes)         # P: (n_cells, n_classes)
+"""
+import numpy as np
+from scipy.optimize import minimize
+
+CLASSES_DEFAULT = ['cortex','endodermis','epidermis','exodermis','pericycle','phloem','root_cap','stele','xylem']
+RANK = {'epidermis':0,'root_cap':0,'exodermis':1,'cortex':2,'endodermis':3,'pericycle':4,
+        'stele':5,'xylem':5,'phloem':5}
+NRANK = 6
+CORE = ('endodermis','pericycle','stele','xylem','phloem')
+S_SCALE, LAMBDA, PASSES, FLOOR = 0.3, 2.0, 2, 1e-3
+
+
+def radial_coordinate(cy, cx, cy_all, cx_all, pred_names, win_deg=15.0):
+    """r_rel for the cells (cy,cx); the surface envelope uses every cell of the section
+    (cy_all,cx_all), which may be a superset (unlabelled cells included)."""
+    core = np.isin(pred_names, CORE)
+    if core.sum() >= 5:
+        y0, x0 = cy[core].mean(), cx[core].mean()
+    else:
+        y0, x0 = cy_all.mean(), cx_all.mean()
+    r_all = np.hypot(cy_all - y0, cx_all - x0); th_all = np.arctan2(cy_all - y0, cx_all - x0)
+    r = np.hypot(cy - y0, cx - x0); th = np.arctan2(cy - y0, cx - x0)
+    da = np.abs((th[:, None] - th_all[None, :] + np.pi) % (2 * np.pi) - np.pi)
+    rs = np.maximum(np.where(da < np.deg2rad(win_deg), r_all[None, :], -1.0).max(1), np.maximum(r, 1.0))
+    return r / rs
+
+
+def _nll(par, v, Q, s):
+    t = np.concatenate([[par[0]], par[0] + np.cumsum(np.logaddexp(0.0, par[1:]))])
+    z = np.clip((t[None, :] - v[:, None]) / s, -50, 50)
+    cum = np.concatenate([np.zeros((len(v), 1)), 1 / (1 + np.exp(-z)), np.ones((len(v), 1))], 1)
+    return -(Q * np.log(np.clip(np.diff(cum, axis=1), 1e-9, 1.0))).sum()
+
+
+def _order_prior(v, Q, s):
+    cw = np.clip(np.cumsum(Q.sum(0) + 1e-9)[:-1] / (Q.sum() + 1e-9), 0.01, 0.99)
+    t0 = np.maximum.accumulate(np.quantile(v, cw))
+    par0 = np.concatenate([[t0[0]], np.log(np.expm1(np.maximum(np.diff(t0), 1e-3)))])
+    par = minimize(_nll, par0, args=(v, Q, s), method='L-BFGS-B', options=dict(maxiter=500)).x
+    t = np.concatenate([[par[0]], par[0] + np.cumsum(np.logaddexp(0.0, par[1:]))])
+    z = np.clip((t[None, :] - v[:, None]) / s, -50, 50)
+    cum = np.concatenate([np.zeros((len(v), 1)), 1 / (1 + np.exp(-z)), np.ones((len(v), 1))], 1)
+    return np.clip(np.diff(cum, axis=1), 1e-9, 1.0), t
+
+
+def correct_section(P, cy, cx, classes, cy_all=None, cx_all=None,
+                    s=S_SCALE, lam=LAMBDA, passes=PASSES):
+    cy_all = cy if cy_all is None else cy_all; cx_all = cx if cx_all is None else cx_all
+    pred = np.asarray(classes)[P.argmax(1)]
+    v = -radial_coordinate(cy, cx, cy_all, cx_all, pred)
+    Pc = P
+    for k in range(passes):
+        Q = np.zeros((len(P), NRANK))
+        for j, c in enumerate(classes):
+            Q[:, RANK[c]] += Pc[:, j]
+        Q /= np.maximum(Q.sum(1, keepdims=True), 1e-12)
+        Po, _ = _order_prior(v, Q, s)
+        Po = Po + FLOOR; Po /= Po.sum(1, keepdims=True)
+        pri = np.stack([Po[:, RANK[c]] for c in classes], 1)
+        F = P * pri ** lam; Pc = F / F.sum(1, keepdims=True)
+    return Pc
+
+
+def correct(P, source_file, cell_id, classes=None, geom=None, **kw):
+    """P (n,k) posterior; source_file (n,) section of each row.  geom: optional dict
+    section -> (cy_all, cx_all, cy, cx) if the centroids are not in geom.npz."""
+    classes = classes or CLASSES_DEFAULT
+    if geom is None:
+        raise ValueError(
+            "correct() needs cell centroids. Pass geom={source_file: (cy_all, cx_all, "
+            "{cell_id: (cy, cx)})}. The research build read these from a cached "
+            "geom.npz that is not part of this package.")
+        d = None  # unreachable; kept so the original block below still parses
+        gsf = d['source_file'].astype(str); gid = d['cell_id'].astype(np.int64)
+        geom = {}
+        for s_ in np.unique(source_file):
+            m = gsf == s_
+            geom[s_] = (d['cy'][m], d['cx'][m], dict(zip(gid[m], zip(d['cy'][m], d['cx'][m]))))
+    out = P.copy()
+    for s_ in np.unique(source_file):
+        m = source_file == s_
+        cy_all, cx_all, look = geom[s_]
+        yx = np.array([look[int(c)] for c in np.asarray(cell_id)[m]])
+        out[m] = correct_section(P[m], yx[:, 0], yx[:, 1], classes, cy_all, cx_all, **kw)
+    return out

rootscope/weights.py

@@ -37,7 +37,7 @@ from pathlib import Path
 DEFAULT_HF_REPO = "ct-tranchau/Rootscope"
 
 # Two published model versions live in one Hub repo:
-#   v2 -> repo root      96 morphometric + 768 DINOv2 ViT-B/14  = 864 features
+#   v2 -> "v2/" folder   96 morphometric + 768 DINOv2 ViT-B/14  = 864 features
 #   v1 -> "v1/" folder   96 morphometric + 384 DINOv2 ViT-S/14  = 480 features
 # v2 is the default: it is the model evaluated in the manuscript and the one the
 # hosted demo runs. Its backbone is heavier, but on a GPU the embedding stage is
@@ -45,8 +45,15 @@ DEFAULT_HF_REPO = "ct-tranchau/Rootscope"
 # each other, because Cellpose segmentation dominates and is identical.
 # v1 is kept for reproducing earlier results. Select with --model-version or
 # ROOTSCOPE_MODEL_VERSION.
-MODEL_VERSIONS = ("v1", "v2")
-DEFAULT_MODEL_VERSION = "v2"
+# v4 -> "v4/" folder   96 morphometric + 768 DINOv2 ViT-B/14, PLUS 51 columns the
+#                      classifier rebuilds every refinement round (36 soft-neighbor
+#                      probabilities, 6 agreement flags, 9 layer-distance features),
+#                      = 915 features; one LightGBM per round, 3 seeds averaged,
+#                      radial-order prior after the last round. This is the model
+#                      in the manuscript (0.881 accuracy, 0.852 macro-F1 on the
+#                      clean 26-section split). Default since 4.0.0.
+MODEL_VERSIONS = ("v1", "v2", "v4")
+DEFAULT_MODEL_VERSION = "v4"
 
 
 def model_version() -> str:
@@ -59,8 +66,14 @@ def model_version() -> str:
 
 
 def _hf_prefix() -> str:
-    """Subfolder inside the Hub repo holding the selected version's weights."""
-    return "" if model_version() == "v2" else f"{model_version()}/"
+    """Subfolder inside the Hub repo holding the selected version's weights.
+    Every version has a folder named after it (v1/, v2/, v4/); the repo root
+    holds only the model card. v2 lived at the root until 2026-09-21."""
+    return f"{model_version()}/"
+
+
+def is_v4() -> bool:
+    return model_version() == "v4"
 
 # Optional escape hatch: a plain base URL under which each file name below is
 # directly downloadable, e.g. "https://zenodo.org/records/XXXXXXX/files".
@@ -90,6 +103,20 @@ MODEL_FILES_OPTIONAL = [
 ]
 CNN_WEIGHTS_FILE = "backbone.pt"
 
+# v4 artifacts. One LightGBM bundle per training seed; each holds the per-round
+# models, the scaler, the 915 feature names and the class order. Seed 42 alone
+# gives a working predictor; all three reproduce the manuscript's 3-seed average.
+V4_MODEL_FILES_REQUIRED = ["lgbm_s42.joblib"]
+V4_MODEL_FILES_OPTIONAL = ["lgbm_s1.joblib", "lgbm_s7.joblib"]
+
+
+def _required_files():
+    return V4_MODEL_FILES_REQUIRED if is_v4() else MODEL_FILES_REQUIRED
+
+
+def _optional_files():
+    return V4_MODEL_FILES_OPTIONAL if is_v4() else MODEL_FILES_OPTIONAL
+
 
 def _cache_dir() -> Path:
     root = os.environ.get("ROOTSCOPE_CACHE")
@@ -110,7 +137,7 @@ def _models_base_url() -> str:
 
 
 def _has_required_models(d: Path) -> bool:
-    return d.is_dir() and all((d / f).exists() for f in MODEL_FILES_REQUIRED)
+    return d.is_dir() and all((d / f).exists() for f in _required_files())
 
 
 # ── Hugging Face Hub ─────────────────────────────────────────────────────────
@@ -185,7 +212,7 @@ def resolve_model_dir(user_arg: str | None = None) -> Path:
         if not _has_required_models(d):
             raise FileNotFoundError(
                 f"--model-dir '{d}' does not contain the required model files "
-                f"({', '.join(MODEL_FILES_REQUIRED)})."
+                f"for model version {model_version()} ({', '.join(_required_files())})."
             )
         return d
 
@@ -212,7 +239,7 @@ def resolve_model_dir(user_arg: str | None = None) -> Path:
         # Name the files rather than globbing: "*.joblib" also matches
         # "v1/*.joblib", which would drag the other version's weights down too.
         snap = _hf_snapshot(
-            repo, [f"{pref}{f}" for f in MODEL_FILES_REQUIRED + MODEL_FILES_OPTIONAL]
+            repo, [f"{pref}{f}" for f in _required_files() + _optional_files()]
         )
         if snap is not None and _has_required_models(snap / pref if pref else snap):
             return snap / pref if pref else snap
@@ -221,8 +248,8 @@ def resolve_model_dir(user_arg: str | None = None) -> Path:
     base_url = _models_base_url()
     if base_url:
         print(f"[rootscope] Fetching model weights into {cache} (first run only)...")
-        _download_set(base_url, MODEL_FILES_REQUIRED, cache, skip_missing=False)
-        _download_set(base_url, MODEL_FILES_OPTIONAL, cache, skip_missing=True)
+        _download_set(base_url, _required_files(), cache, skip_missing=False)
+        _download_set(base_url, _optional_files(), cache, skip_missing=True)
         return cache
 
     raise RuntimeError(
@@ -233,7 +260,7 @@ def resolve_model_dir(user_arg: str | None = None) -> Path:
         "  • Put the model files in a folder and pass --model-dir <folder>, or\n"
         "  • export ROOTSCOPE_MODEL_DIR=/path/to/models, or\n"
         "  • export ROOTSCOPE_MODELS_URL=<base url> to self-host them.\n"
-        f"Required files: {', '.join(MODEL_FILES_REQUIRED)}"
+        f"Required files ({model_version()}): {', '.join(_required_files())}"
     )