Compare commits

..

No commits in common. "99bea82841eec5f7cdbf1b7310cf8d7167e1a877" and "41c88b2a67cc1ec21bc2b6f9ef2420fc0beecb20" have entirely different histories.

10 changed files with 51 additions and 697 deletions

View file

@ -22,37 +22,6 @@ onto the T1c grid with SMD+linear. After the repair pass:
volumes, 3893 subjects; unlabeled pool 13280
- tumor volume stats: median 3116 mm3, p95 24.8k, max 3.2M, zero_frac 0.12%
## Dataset exclusions (2026-10-04)
Row-level curation lives in `src/exclusions.py` and is applied at the **first
stage** by `scripts/make_tv_slices.py` (right after the rows are fetched from
the sources, before any split): it classifies every labeled volume and writes
its TV-slice QA figure to `data/tv_slices/inclusion/` (kept) or
`data/tv_slices/exclusion/` (excluded, reason in the title). `05_build_splits`
re-applies the same rules **after** split assignment (kept rows retain their
split; the subject-level shuffle boundaries are unchanged):
- total TV > 1e6 mm3: **1** row (ntuh_6019528, label outlines the full skull)
- single target (1 connected component) with TV < 100 mm3: **109** rows
- non-brain images: **33** rows (all ntuh — spine/neck/torso CyberKnife
cases whose series names are indistinguishable from brain T1c; identified
by visual review of the TV-slice figures, now in `data/tv_slices/exclusion/`)
Audit trail: `data/manifests/excluded.jsonl` (143 rows, reason + split);
pre-exclusion manifests in `data/manifests/backup_20261004/`. New dataset:
**4057** labeled rows → train 3447 / val 205 / test 405 (3771 subjects);
unlabeled pool unchanged (13280); volume stats now median 3212 mm3, p95
23.4k, max 163k (m6 meningioma), zero_frac still 0.12% (5 rows, e.g.
ntuh_3436472 ×3).
After the initial 39-row non-brain list, expert inspection overturned 6
(skull-base / posterior-fossa brain targets that look like cervical spine at
the TV-max slice: ntuh_5902230, ntuh_5231609, ntuh_5083686, ntuh_5157194,
ntuh_5303758, ntuh_6373987). The remaining 33 were re-audited at the
mid-cranial slice (widest head-mask axial) — all still non-brain. Review
protocol now recorded in `src/exclusions.py`: classify at full resolution
on the mid-cranial slice, not just the TV-max slice.
## Smoke run (2026-10-04)
Small/fast end-to-end check that all three pipelines run on the rebuilt data
@ -83,99 +52,6 @@ ntuh). Artifacts: `results/benchmark_pipelines_smoke.json`,
ModuleNotFoundError) — installed `pip install monai==1.6.1 --no-deps`
(torch 2.14+cu126 untouched); `import monai` verified.
## Full-dataset 10-epoch run (2026-10-05, updated dataset)
`--train-epochs 10`, default budgets on the current curated splits: **2750**
train rows = 80% of split_train (3438 labeled), **40** eval rows = 10% of
split_test (404), TTA on. The 2026-10-04 interim run (pre-exclusion budget
2856/42: A 0.0008, C 0.2119) is superseded; its JSONs are kept in
`results/backup_20261004/` (smoke B checkpoint in
`nnu/results/backup_20261004_Dataset221_smoke/`).
```
python scripts/benchmark_pipelines.py --pipelines C --gpu 2 --train-epochs 10 \
--out results/benchmark_full_C.json # done
python scripts/benchmark_pipelines.py --pipelines A --gpu 0 --train-epochs 10 \
--a-fg-ce-weight 500 --out results/benchmark_full_A.json # done
# A baseline (uniform crops) on this dataset: dice 0.0026, see table;
# checkpoint runs_bench/A/best_baseline_10ep_uniform_crop.pt
NNU_ROOT=/home/xfr/longitudinal_nnu BENCH_B_DIR=/home/xfr/longitudinal_benchB \
python scripts/benchmark_pipelines.py --pipelines B --gpu 1 --train-epochs 10 \
--out results/benchmark_full_B.json # done (local-disk NNU_ROOT, see below)
```
| pipeline | train s/epoch | wall | tr VRAM | s/vol (TTA) | in VRAM | params | weights | dice @10ep |
|---|---|---|---|---|---|---|---|---|
| C MONAI UNet(16→128) | 256.7 (916 steps) | 43 min | 0.4 GB | 0.37 ± 0.06 | 0.5 GB | 1.19 M | 14.3 MB | **0.2162** |
| A Unet3D(16, depth 4), improved (fg crops + CE w=500) | 232.3 (916 steps) | 41 min | 2.0 GB | 3.60 ± 0.76 | 0.5 GB | 1.46 M | 17.6 MB | **0.0369** |
| A baseline (uniform crops, unweighted loss) | 216.6 (916 steps) | 45 min | 2.0 GB | 3.33 ± 0.86 | 0.5 GB | 1.46 M | 17.6 MB | **0.0026** |
| B nnU-Net 3d_fullres | 331.5 (incl. 275-case val/epoch) | ~2.5 h (raw 6 m, prep 62 m, train 55 m, infer 14 m) | 9.3 GB | 21.84 | 1.8 GB | 30.8 M | 246 MB | **0.1424** (hard 0.1424) |
C: 25/39 timed test vols with dice > 0.001 (max 0.7392); loss 0.403 → 0.182.
10 epochs is far below convergence for these sparse targets (median tumor
3.2 k mm³); A/C numbers at this budget are intermediate. Checkpoints at
`runs_bench/{A,C}/best.pt` (extend with `--resume --train-epochs <total>`).
### A: why dice ≈ 0, and the fix (2026-10-05)
The A baseline predicts all-background (38/39 test rows dice exactly 0.0).
Root-cause work:
1. **Uniform random 96³ crops starve the loss of foreground.** A 96³ patch is
885 k voxels; the median tumor is ~4.6 k mm³ (0.02–0.05% of a head volume).
Simulated uniform crops contain the whole blob in only 8–37% of patches and
a *median* crop often has 0 tumor voxels. Pipeline C's
`RandCropByPosNegLabeld(pos=1, neg=1)` centers half its patches on the
tumor — ~4× more foreground supervision per epoch.
2. **Fg-aware crops alone are not enough: A under-confides.** With the crop
protocol fixed (`src/dataset.py`), A's loss matches C's at every epoch
(0.193 vs 0.182 @10ep) yet p_tumor maxes at 0.43 (C: ~1.0): it localizes
(top-k dice 0.48 on best rows) but never crosses the 0.5 threshold, so
dice@0.5 stays 0.
3. **Fix: class-weighted CE** — tumor-class weight `w` in the CE part of the
per-sample loss (`--a-fg-ce-weight`, `src/losses.py`). 3-epoch probes on
this dataset: w=20 → pmax median 0.33, dice 0.0; w=100 → 0.41, dice
0.0002; w=500 → 0.62, dice 0.0188 (34/39 rows > 0.5). Official 10-epoch
w=500 run: **dice 0.0369** (vs 0.0026 baseline, 14×; 8/39 rows > 0.001,
best 0.555, pmax median 0.55, 29/39 rows > 0.5) — A now segments, but
still trails C (0.2162) at this budget; its loss is still descending at
epoch 10 (0.78 → 0.51), i.e. under-converged vs C (0.41 → 0.18), so the
gap should narrow with more epochs.
Changes: `src/dataset.py` (pos/neg crop sampling `pos_neg=0.5`, per-worker
RNG reseed), `src/losses.py` (`soft_dice_fg`, `fg_ce_weight`), benchmark flags
`--a-fg-dice` / `--a-fg-ce-weight`, per-row `pmax` in A's `dice_per_row`,
full traceback in failure JSONs. A/B probes: `results/bench_A_*probe*.json`.
### B: CIFS failures and local-disk NNU_ROOT (2026-10-05)
The CIFS `/mnt/b4` mount dropped I/O repeatedly on 2026-10-05 (NAS flapping:
`OSError [Errno 112] Host is down`, `RuntimeError: Error while setting the
slice`, one silent `rmtree` failure that made the raw rebuild collide with
stale entries → `FileExistsError`). Four B attempts died (07:52 plan&prep;
11:12 raw EEXIST; ~12:40 / ~13:50 / ~14:00 plan&prep). Fixes:
- `nnu_common.rmtree_verified()` + defensive per-link unlink in
`make_raw_dataset` (EEXIST on rebuild)
- failure JSONs now carry a full `traceback` key
- plan step retries raised 2 → 5 (`N.run(..., retries=5)`)
- **B's nnU-Net tree moved off CIFS**: `nnu_common.nnu_root()` honors
`NNU_ROOT` (set to `/home/xfr/longitudinal_nnu`, local ZFS ~1.8 T free) and
the benchmark's B pred I/O honors `BENCH_B_DIR`
(`/home/xfr/longitudinal_benchB`). CIFS is then touched only for the
one-time raw image reads from `/mnt/pve` and small log/JSON appends.
The pre-`NNU_ROOT` attempt's half-built preprocessed tree under
`nnu/preprocessed/` is stale — the local tree is the one being used.
The local-disk run succeeded on attempt 1 (2026-10-05 14:07→16:38 UTC):
raw build 6.3 min, plan&preprocess 62 min (140 GB preprocessed), 10 epochs
in 55 min (epoch times 101–167 s train + 275-case validation —
`nnUNet_compile is enabled`, vs the Oct-4 eager-mode pace of ~2.4 h/epoch),
inference 21.8 s/vol TTA. Val pseudo dice 0.091 → 0.459, best EMA 0.285;
test dice 0.1424 (29/39 rows > 0.001, max 0.675). Checkpoints in
`/home/xfr/longitudinal_nnu/results/Dataset221_T1C_BENCH/...fold_0/`
(copy to `nnu/results/` for the record if local space is reclaimed).
## Resuming training (`--resume`)
All three pipelines can continue from their latest benchmark checkpoint:
@ -207,9 +83,8 @@ python scripts/benchmark_pipelines.py --pipelines "A B C" --resume \
## Readiness notes for a full run
- Budget knobs: `--train-epochs` (default 100), `--max-rows` (default 80% of
split_train = 2750 for A/C and for B's `--nnu-cases` on the current curated
splits; was 2856 pre-exclusion), `--max-eval-rows` (default 10% of
split_test = 40), `--no-tta`.
split_train ≈ 2856 for A/C and for B's `--nnu-cases`), `--max-eval-rows`
(default 10% of split_test ≈ 42), `--no-tta`.
- `--skip-train` is only usable for **A** (`runs/round0/best.pt`, trained
2026-09-25 on the *previous* split — a few patients moved between
train/test in the rebuild, so its test dice carries a small leakage

View file

@ -40,8 +40,6 @@ data/
proc/ <key>.nii.gz (1mm, cropped, normalized)
proc/<key>_label.nii.gz
procmeta/<key>.json geometry (origin/direction/crop_vox) + normalization
tv_slices/inclusion/ TV-slice QA figures, kept rows (<vol_mm3>_<key>.png)
tv_slices/exclusion/ TV-slice QA figures, excluded rows (reason in title)
vols.json labeled tumor volume stats (mm3 percentiles)
pseudo/roundK/ in-house pseudo-labels (rows.jsonl, masks, summary.json)
pseudo_nnu/roundK/ nnU-Net pseudo-labels (rows.jsonl, masks, summary.json)
@ -71,13 +69,12 @@ DDP, per-sample weighted loss.
| 02 | `02_build_m6_dataset.py` | Build M6-2025 manifests (GTV CT→T1c registration/warp) |
| 03 | `03_scan_lee_t1c.py` | Scan lee for brain T1c series → raw + selected manifests |
| 04 | `04_reconstruct_lee.py` | Reconstruct lee T1c niftis from JPG slices + txt metadata |
| 05 | `05_build_splits.py` | Patient-level train/val/test splits + unlabeled pool + volume stats (re-applies `src/exclusions` post-split) |
| 05 | `05_build_splits.py` | Patient-level train/val/test splits + unlabeled pool + volume stats |
| 06 | `06_pseudo_label.py` | Pseudo-label the pool with a round model (sliding-window + TTA) |
| 07 | `07_train.py` | DDP training entrypoint (labeled + weighted pseudo rows) |
| 08 | `08_eval.py` | Held-out test evaluation (Dice from probability map) |
| 09 | `09_run_iterative.py` | Orchestrates rounds 0…K and produces the summary table/plot |
| — | `preprocess.py` | Crop/normalize source volumes → `data/proc/` |
| — | `make_tv_slices.py` | First-stage curation (right after source fetch): classifies labeled rows via `src/exclusions.py` + TV-slice QA figures → `data/tv_slices/{inclusion,exclusion}/` |
| — | `scan_procs.py` | Flag corrupt/non-finite processed volumes |
| — | `test_dataloader.py` | Smoke-test the training dataloader |

View file

@ -11,11 +11,6 @@ normalized, used by Pipelines A and C) and native fields (img/label, the raw
source NIfTIs — used by Pipeline B, which feeds nnU-Net its own
plan_and_preprocess). Row membership still requires the processed volume to
exist, so the patient-level splits stay identical across pipelines.
Before splitting, every labeled row is screened with the row-level exclusion
rules (src/exclusions.py: total TV > 1e6 mm3, single target < 100 mm3, and
the non-brain ntuh rows in NON_BRAIN_KEYS). Excluded rows are written to
data/manifests/excluded.jsonl with their reason and never enter a split.
"""
import os
import sys
@ -24,16 +19,20 @@ import json
import argparse
import random
import numpy as np
from src.common import ROOT, DATA, d, load_jsonl, save_jsonl, proc_row
from src.exclusions import label_stats, exclusion_reason
from src.common import ROOT, DATA, d, load_jsonl, save_jsonl, read_nii_arr
def _label_stats_row(r):
"""(key, label stats) for the ProcessPoolExecutor in main()."""
try:
return r["key"], label_stats(r["plabel"])
except Exception:
return r["key"], {"total": 0, "n_cc": 0}
def proc_row(r, prefix):
key = r["key"]
p = os.path.join(d("data/proc"), key + ".nii.gz")
if not os.path.exists(p):
return None
row = {"key": key, "subject": f"{prefix}_{r['subject']}", "pimg": p,
"img": r["img"], "label": r.get("label"), "date": r.get("date")}
lab = os.path.join(d("data/proc"), key + "_label.nii.gz")
if r.get("label") and os.path.exists(lab):
row["plabel"] = lab
return row
def main():
@ -54,31 +53,6 @@ def main():
labeled_rows.append(p)
print(f"labeled processed rows: {len(labeled_rows)}")
# row-level curation (src/exclusions.py): label stats + exclusion reason
# per row. Exclusions are applied AFTER the split assignment, so kept rows
# retain the split of the 2026-10-04 rebuild and excluded.jsonl is the
# audit trail (excluded subjects still count in the shuffle boundaries).
from collections import Counter
from concurrent.futures import ProcessPoolExecutor
tv_vols, reason_of = {}, {}
with ProcessPoolExecutor(max_workers=32) as ex:
for k, stats in ex.map(_label_stats_row, labeled_rows, chunksize=16):
tv_vols[k] = stats["total"]
reason_of[k] = exclusion_reason(k, stats)
n_excl = sum(1 for v in reason_of.values() if v)
def _cat(reason):
return (reason.split(" image", 1)[0] + " image"
if "non-brain" in reason else
reason.split(" TV ", 1)[0] + " TV < 100 mm3"
if "single target" in reason else
"total TV > 1e6 mm3")
rc = Counter(_cat(v) for v in reason_of.values() if v)
print(f"excluded rows: {n_excl}")
for reason, n in sorted(rc.items()):
print(f" excluded {n:4d} {reason}")
subjects = sorted({r["subject"] for r in labeled_rows})
random.shuffle(subjects)
n_test = max(1, int(len(subjects) * args.frac_test))
@ -88,24 +62,18 @@ def main():
train_subj = set(subjects[n_test + n_val:])
print(f"subjects: {len(subjects)} train={len(train_subj)} val={len(val_subj)} test={len(test_subj)}")
splits, excluded = {"train": [], "val": [], "test": []}, []
splits = {"train": [], "val": [], "test": []}
for r in labeled_rows:
r["w"] = 1.0
if r["subject"] in test_subj:
sk = "test"
splits["test"].append(r)
elif r["subject"] in val_subj:
sk = "val"
splits["val"].append(r)
else:
sk = "train"
if reason_of[r["key"]]:
excluded.append({**r, "split": sk, "exclude_reason": reason_of[r["key"]]})
else:
splits[sk].append(r)
splits["train"].append(r)
for k, v in splits.items():
save_jsonl(v, os.path.join(DATA, "manifests", f"split_{k}.jsonl"))
print(f" split_{k}: {len(v)} volumes / {len(set(r['subject'] for r in v))} subjects")
save_jsonl(excluded, os.path.join(DATA, "manifests", "excluded.jsonl"))
kept_rows = [r for v in splits.values() for r in v]
labeled_subjects = train_subj | val_subj | test_subj
@ -146,12 +114,15 @@ def main():
print(f"unlabeled pool: {len(pool)} volumes ({sum(1 for r in pool if r['source']=='m6')} m6, "
f"{sum(1 for r in pool if r['source']=='lee')} lee) from {len(set(r['subject'] for r in pool))} subjects")
# tumor volume stats of the kept dataset rows (1mm voxels == mm3; totals from
# the curation pass)
# tumor volume stats (1mm voxels == mm3)
vols = []
for r in kept_rows:
vols.append({"key": r["key"], "subject": r["subject"],
"vol_mm3": tv_vols[r["key"]]})
for r in labeled_rows:
try:
l = read_nii_arr(r["plabel"])
n = int((l > 0.5).sum())
vols.append({"key": r["key"], "subject": r["subject"], "vol_mm3": n})
except Exception:
pass
v = np.array([x["vol_mm3"] for x in vols])
stats = {"n": int(v.size),
"p2": float(np.percentile(v, 2)), "p5": float(np.percentile(v, 5)),

View file

@ -46,7 +46,6 @@ import subprocess
import sys
import threading
import time
import traceback
from datetime import datetime
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
@ -81,11 +80,6 @@ def parse_args():
ap.add_argument("--workers", type=int, default=4, help="dataloader workers (A/C)")
ap.add_argument("--batch", type=int, default=3, help="per-GPU batch (A/C)")
ap.add_argument("--patch", type=int, default=96, help="patch/window (A/C)")
ap.add_argument("--a-fg-dice", action="store_true",
help="A: use tumor-class-only soft dice in the per-sample loss "
"(instead of the background+tumor mean)")
ap.add_argument("--a-fg-ce-weight", type=float, default=1.0,
help="A: tumor class weight in the CE part of the loss (1.0 = unweighted)")
ap.add_argument("--no-tta", action="store_true", help="disable each pipeline's standard TTA")
ap.add_argument("--skip-train", action="store_true",
help="eval existing study checkpoints (runs/round0, runs_monai/round0, nnu/results 210) "
@ -262,8 +256,7 @@ def bench_A(args, train_rows, test_rows):
wts = wts.to(dev, non_blocking=True)
with torch.autocast("cuda", dtype=torch.bfloat16):
logits = model(img)
losses = per_sample_loss(logits.float(), one_hot(lab, 2), fg_only=args.a_fg_dice,
fg_ce_weight=args.a_fg_ce_weight)
losses = per_sample_loss(logits.float(), one_hot(lab, 2))
if torch.isnan(losses).any() or torch.isinf(losses).any():
losses = torch.zeros_like(losses)
loss = (losses * wts).sum() / wts.sum().clamp(min=1e-6)
@ -297,7 +290,7 @@ def bench_A(args, train_rows, test_rows):
args.patch, step, tta)
torch.cuda.synchronize(dev)
reset_peak(dev)
times, dices, pmaxs = [], [], []
times, dices = [], []
with torch.no_grad():
for r in test_rows[args.eval_warmup:]:
vol = read_nii_arr(r["pimg"]).astype("float32")
@ -309,7 +302,6 @@ def bench_A(args, train_rows, test_rows):
# probs are on the padded grid; crop back to the volume/label shape
osh = vol.shape
probs = probs[:osh[0], :osh[1], :osh[2]]
pmaxs.append(probs.max())
lab = read_nii_arr(r["plabel"]).astype("uint8")
if lab.shape != osh:
lab = np.pad(lab, [(0, max(osh[i] - lab.shape[i], 0)) for i in range(3)])[:osh]
@ -319,9 +311,8 @@ def bench_A(args, train_rows, test_rows):
res["infer_vram_mb"] = peak_mb(dev)
res["dice"] = round(float(np.mean(dices)), 4) if dices else 0.0
res["dice_per_row"] = [{"key": test_rows[args.eval_warmup + i]["key"],
"dice": round(dices[i], 4), "s": round(times[i], 2),
"pmax": round(float(pm.max()), 4)}
for i, pm in zip(range(len(dices)), pmaxs)]
"dice": round(dices[i], 4), "s": round(times[i], 2)}
for i in range(len(dices))]
del model
free_gpu()
return res
@ -472,8 +463,7 @@ def bench_B(args, train_rows, test_rows):
res["raw_build_s"] = round(time.time() - t0, 1)
if not args.nnu_skip_prep:
t0 = time.time()
N.run(N.plan_preprocess_cmd(8), os.path.join(d("logs/nnu"), "bench_plan.log"), env=N.nnu_env(),
retries=5) # CIFS /mnt/b4 drops during the ~3h preprocess; extra attempts are cheap
N.run(N.plan_preprocess_cmd(8), os.path.join(d("logs/nnu"), "bench_plan.log"), env=N.nnu_env())
res["plan_preprocess_s"] = round(time.time() - t0, 1)
else:
res["plan_preprocess_s"] = 0.0
@ -528,9 +518,8 @@ def bench_B(args, train_rows, test_rows):
te_keys = [r["key"] for r in te]
te_native = native_rows_for(te_keys)
pred_base = os.environ.get("BENCH_B_DIR", os.path.join(d("runs_bench"), "B"))
in_dir = os.path.join(pred_base, "pred_in")
out_dir = os.path.join(pred_base, "pred_out")
in_dir = d("runs_bench/B/pred_in")
out_dir = d("runs_bench/B/pred_out")
shutil.rmtree(in_dir, ignore_errors=True)
shutil.rmtree(out_dir, ignore_errors=True)
os.makedirs(in_dir)
@ -680,9 +669,8 @@ def main():
r = runners[p](args, train_rows, test_rows)
r["wall_s"] = round(time.time() - t0, 1)
except Exception as e: # noqa
tb = traceback.format_exc()
r = {"pipeline": p, "error": repr(e), "traceback": tb, "wall_s": round(time.time() - t0, 1)}
print(f"[bench] pipeline {p} FAILED:\n{tb}")
r = {"pipeline": p, "error": repr(e), "wall_s": round(time.time() - t0, 1)}
print(f"[bench] pipeline {p} FAILED: {e!r}")
results.append(r)
free_gpu()

View file

@ -1,230 +0,0 @@
"""First-stage curation + TV-slice QA figures for labeled rows.
Runs on the labeled source rows (data/manifests/{ntuh,m6_labeled}.jsonl, 1mm
processed) — the first place rows exist after fetching from the sources,
before any split (05). For each labeled volume:
1. classify with the exclusion rules in src/exclusions.py (total TV > 1e6
mm3, single target < 100 mm3, curated non-brain list),
2. render one PNG with the three orthogonal slices (axial | coronal |
sagittal) that carry the largest TV (target-volume) area, with the TV
contour overlaid; when the label has several connected targets
(26-connectivity), only the largest one is shown,
3. write the figure to data/tv_slices/inclusion/ (kept) or
data/tv_slices/exclusion/ (excluded, reason in the title).
Slices are picked per body axis (S/P/L) from the procmeta direction, so the
panel order follows radiological conventions (same mapping as
scripts/make_screenshots.py); volumes that are not body-aligned fall back to
the image-axis slices with a note in the title. Filenames are
<total_vol_mm3:07d>_<key>.png, so each folder lists small TV -> large TV.
The split files (05) are the dataset of record; this stage only curates and
documents. 05 applies the same src/exclusions rules after split assignment,
so the two always agree on kept/excluded.
Usage:
python scripts/make_tv_slices.py # all source rows, 32 workers
python scripts/make_tv_slices.py --from-splits # rows/tag from split files
python scripts/make_tv_slices.py --keys k1 k2 # specific rows
python scripts/make_tv_slices.py --limit 50 # first N rows
python scripts/make_tv_slices.py --force # overwrite existing PNGs
"""
import os
import sys
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import glob
import json
import argparse
from collections import Counter
import numpy as np
import SimpleITK as sitk
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from scipy import ndimage
from src.common import DATA, d, load_jsonl, proc_row, labeled_source_rows
from src.exclusions import label_stats, exclusion_reason
from scripts.make_screenshots import vol_to_body
def out_name(vol_total, key):
"""Flat-folder filename: zero-padded TOTAL TV volume (mm3, all targets)
first so a plain listing sorts small -> large, then the row key."""
return f"{int(vol_total):07d}_{key}.png"
def largest_cc(mask, structure=None):
"""Largest connected component of a boolean mask (26-conn by default)."""
if structure is None:
structure = np.ones((3, 3, 3), dtype=int)
lab, n = ndimage.label(mask, structure=structure)
if n == 0:
return mask, 0
sizes = ndimage.sum(mask, lab, range(1, n + 1))
return lab == (int(np.argmax(sizes)) + 1), int(n)
def max_area_slices(cc):
"""Per-axis index of the slice containing the largest CC area."""
return tuple(int(np.argmax(cc.sum(axis=tuple(j for j in range(cc.ndim) if j != i))))
for i in range(cc.ndim))
def render(key, tag, mr, cc, n_cc, vol_total, direction, reason, out, dpi=110):
"""1x3 figure: axial | coronal | sagittal, max-CC-area slice each, with the
largest-CC contour in red. Writes to `out`. Returns the path."""
body, ok, _ = vol_to_body(mr, direction)
if ok:
c, _, _ = vol_to_body(cc.astype(np.float32), direction)
si, pi, li = max_area_slices(c)
panels = [
(body[si], c[si], "axial", "S", si),
(body[:, pi, :][::-1], c[:, pi, :][::-1], "coronal", "P", pi),
(body[:, :, li][::-1], c[:, :, li][::-1], "sagittal", "L", li),
]
else: # not body-aligned: image-axis fallback
zi, yi, xi = max_area_slices(cc)
panels = [
(mr[zi], cc[zi], "axial*", "z", zi),
(mr[:, yi, :][::-1], cc[:, yi, :][::-1], "coronal*", "y", yi),
(mr[:, :, xi][::-1], cc[:, :, xi][::-1], "sagittal*", "x", xi),
]
lo, hi = np.percentile(mr[mr > 0], [1, 99.5])
panel_specs = []
for im, c2, title, axn, idx in panels:
nr, nc = im.shape
a = 1.0 # processed volumes are 1mm isotropic
panel_specs.append((im, c2, f"{title} {axn}={idx}", a, nr, nc))
max_in, gap_in, m_in, top_in, bot_in = 4.4, 0.25, 0.3, 1.7, 0.2
w_in, h_in = [], []
for _, _, _, a, nrow, ncol in panel_specs:
pw, ph = ncol, nrow * a
m = max(pw, ph)
w_in.append(max_in * pw / m)
h_in.append(max_in * ph / m)
plot_h = max_in
fig_w = sum(w_in) + 2 * gap_in + 2 * m_in
fig_h = plot_h + top_in + bot_in
fig = plt.figure(figsize=(fig_w, fig_h))
left = m_in / fig_w
for (im, c2, title, a, _, _), w, h in zip(panel_specs, w_in, h_in):
axi = fig.add_axes([left, (bot_in + (plot_h - h) / 2) / fig_h, w / fig_w, h / fig_h])
axi.imshow(im, cmap="gray", vmin=lo, vmax=hi, origin="upper", aspect=a)
if c2.any():
axi.contour(c2, levels=[0.5], colors="r", linewidths=1.2)
axi.set_title(title, color="w", fontsize=11)
axi.set_xticks([])
axi.set_yticks([])
for spine in axi.spines.values():
spine.set_edgecolor("0.35")
left += (w + gap_in) / fig_w
tkey = key.split("_", 1)[1] if "_" in key else key
align = "body-aligned" if ok else "NOT body-aligned (image axes)"
status = f"EXCLUDED: {reason}" if reason else "kept"
fig.suptitle(
f"{tkey} [{tag}] TV total={int(vol_total)} mm3, "
f"largest={int(cc.sum())} mm3 ({n_cc} targets) {status} {align}",
color="w", fontsize=12)
fig.patch.set_facecolor("k")
fig.savefig(out, facecolor="k", dpi=dpi)
plt.close(fig)
return out
def worker(args):
key, tag, pimg, plabel, out_base, force, dpi = args
try:
meta = json.load(open(os.path.join(d("data/procmeta"), key + ".json")))
stats = label_stats(plabel)
except Exception:
return key, "error: unreadable label/procmeta"
reason = exclusion_reason(key, stats)
cat = reason.split(" (", 1)[0] if reason else None
# reconcile a stale figure for this key left in the other folder by an
# earlier run under a different exclusion status (e.g. a key newly added to
# NON_BRAIN_KEYS would otherwise keep its old inclusion/ figure)
folder = "exclusion" if reason else "inclusion"
other = "inclusion" if reason else "exclusion"
for stale in glob.glob(os.path.join(out_base, other, f"*_{key}.png")):
os.remove(stale)
if stats["total"] == 0:
return key, "empty-label", cat
status = "excluded" if reason else "included"
out = os.path.join(out_base, folder, out_name(stats["total"], key))
if os.path.exists(out) and not force:
return key, f"skip-{status}", cat
try:
mr = np.asarray(sitk.GetArrayFromImage(sitk.ReadImage(pimg)), dtype=np.float32)
lab = sitk.GetArrayFromImage(sitk.ReadImage(plabel))
cc, n_cc = largest_cc(lab > 0.5)
try:
render(key, tag, mr, cc, n_cc, int(stats["total"]),
np.array(meta["direction"]), reason, out, dpi)
finally:
plt.close("all")
return key, status, cat
except Exception as e: # noqa
return key, f"error: {e!r}", cat
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--from-splits", action="store_true",
help="rows/tags from split_{train,val,test}.jsonl instead of "
"the source manifests (ntuh + m6_labeled)")
ap.add_argument("--keys", nargs="*", default=[])
ap.add_argument("--limit", type=int, default=0)
ap.add_argument("--workers", type=int, default=32)
ap.add_argument("--force", action="store_true")
ap.add_argument("--out", default=d("data/tv_slices"))
ap.add_argument("--dpi", type=int, default=110)
args = ap.parse_args()
os.makedirs(os.path.join(args.out, "inclusion"), exist_ok=True)
os.makedirs(os.path.join(args.out, "exclusion"), exist_ok=True)
if args.from_splits:
rows = []
for s in ("train", "val", "test"):
for r in load_jsonl(os.path.join(DATA, "manifests", f"split_{s}.jsonl")):
rows.append((r["key"], s, r["pimg"], r.get("plabel")))
else:
rows = [(r["key"],
"ntuh" if r["key"].startswith("ntuh_") else "m6",
r["pimg"], r.get("plabel")) for r in labeled_source_rows()]
if args.keys:
keep = set(args.keys)
rows = [r for r in rows if r[0] in keep]
if args.limit:
rows = rows[: args.limit]
print(f"rows: {len(rows)}")
from concurrent.futures import ProcessPoolExecutor, as_completed
tasks = [(k, tag, pimg, plabel, args.out, args.force, args.dpi)
for k, tag, pimg, plabel in rows]
stats = {}
reasons = Counter()
done = 0
with ProcessPoolExecutor(max_workers=args.workers) as ex:
futs = {ex.submit(worker, t): t for t in tasks}
for fu in as_completed(futs):
k, status, cat = fu.result()
stats[status] = stats.get(status, 0) + 1
if status == "excluded" and cat:
reasons[cat] += 1
done += 1
if done % 100 == 0 or done == len(tasks):
print(f" {done}/{len(tasks)} {stats}", flush=True)
print(f"done: {stats}")
if reasons:
print("excluded by rule:")
for r, n in sorted(reasons.items()):
print(f" {n:4d} {r}")
if __name__ == "__main__":
main()

View file

@ -31,11 +31,7 @@ TRAINERS_DIR = os.path.join(SCRIPTS_NNU, "trainers")
def nnu_root():
# NNU_ROOT override: the CIFS /mnt/b4 mount drops writes transiently
# ("Host is down", "Error while setting the slice"), which kills the ~3h
# plan&preprocess and the multi-day train. Point it at local disk
# (e.g. /home/xfr/longitudinal_nnu) to keep all nnU-Net I/O off CIFS.
return os.environ.get("NNU_ROOT", d(os.path.join("nnu")))
return d(os.path.join("nnu"))
def raw_ds():
@ -160,29 +156,14 @@ def native_label_target(img_path, lab_path, link_path):
sitk.WriteImage(out, link_path, True)
def rmtree_verified(path, tries=10, pause=2.0):
"""Remove a tree and verify it is actually gone. CIFS (/mnt/b4, soft mount)
can transiently fail individual unlinks; a silent ignore_errors rmtree leaves
stale entries behind and the next rebuild's os.symlink dies with EEXIST."""
for i in range(tries):
if not os.path.isdir(path):
return
shutil.rmtree(path, ignore_errors=True)
if not os.path.isdir(path) or len(os.listdir(path)) == 0:
return
time.sleep(pause * (i + 1))
raise RuntimeError(f"could not fully remove {path} "
f"({len(os.listdir(path)) if os.path.isdir(path) else '?'} entries left)")
def make_raw_dataset(rows, channel="0000"):
"""rows: list of {key, img, label} — img/label are native (unprocessed) NIfTI
paths. Rebuilds imagesTr/labelsTr + dataset.json. Images are symlinked
as-is; labels via native_label_target. See section note."""
img_dir = os.path.join(raw_ds(), "imagesTr")
lab_dir = os.path.join(raw_ds(), "labelsTr")
rmtree_verified(img_dir)
rmtree_verified(lab_dir)
shutil.rmtree(img_dir, ignore_errors=True)
shutil.rmtree(lab_dir, ignore_errors=True)
os.makedirs(img_dir, exist_ok=True)
os.makedirs(lab_dir, exist_ok=True)
seen = set()
@ -195,13 +176,8 @@ def make_raw_dataset(rows, channel="0000"):
raise FileNotFoundError(f"image missing for {key}: {img}")
if not os.path.exists(lab):
raise FileNotFoundError(f"label missing for {key}: {lab}")
img_link = os.path.join(img_dir, f"{key}_{channel}.nii.gz")
lab_link = os.path.join(lab_dir, f"{key}.nii.gz")
for lp in (img_link, lab_link): # belt-and-braces against stale CIFS state
if os.path.lexists(lp):
os.unlink(lp)
os.symlink(os.path.abspath(img), img_link)
native_label_target(img, lab, lab_link)
os.symlink(os.path.abspath(img), os.path.join(img_dir, f"{key}_{channel}.nii.gz"))
native_label_target(img, lab, os.path.join(lab_dir, f"{key}.nii.gz"))
ds_json = {"channel_names": {"0": "t1c"},
"labels": {"background": 0, "tumor": 1},
"numTraining": len(seen),

View file

@ -223,45 +223,6 @@ def write_arr(arr, path, itk_img=None):
sitk.WriteImage(itk_img, path, True)
def proc_row(r, prefix):
"""Processed-volume view of a source manifest row, or None when the
processed volume does not exist.
Source rows come from data/manifests/{ntuh,m6_labeled,m6_unlabeled}.jsonl
({key, subject, img, label?, date?} — bare subjects, no source prefix).
Returns {key, subject (prefixed), pimg, img, label, date[, plabel]} where
pimg/plabel point at data/proc/<key>{,_label}.nii.gz (1mm processed).
"""
key = r["key"]
proc_root = d("data/proc")
p = os.path.join(proc_root, key + ".nii.gz")
if not os.path.exists(p):
return None
row = {"key": key, "subject": f"{prefix}_{r['subject']}", "pimg": p,
"img": r["img"], "label": r.get("label"), "date": r.get("date")}
lab = os.path.join(proc_root, key + "_label.nii.gz")
if r.get("label") and os.path.exists(lab):
row["plabel"] = lab
return row
def labeled_source_rows():
"""All labeled rows fetched from the sources (ntuh + m6_labeled), as
processed rows — the first dataset stage, before any split (05)."""
rows = []
for name, prefix in (("ntuh", "ntuh"), ("m6_labeled", "m6")):
p = os.path.join(DATA, "manifests", name + ".jsonl")
if not os.path.exists(p):
continue
for r in load_jsonl(p):
if not r.get("label"):
continue
pr = proc_row(r, prefix)
if pr:
rows.append(pr)
return rows
def dice(a, b, eps=1e-7):
a = a > 0
b = b > 0

View file

@ -21,24 +21,13 @@ class PatchDataset(Dataset):
row: {key, pimg, plabel?, is_neg?, w (loss weight, default 1.0)}
yields (img(1,H,W,D), label(H,W,D) or None, weight)
Training crops for labeled rows: `pos_neg` fraction are foreground-centered
(random tumor voxel as crop center), the rest are foreground-free background
crops (up to 20 random retries, plain uniform if the volume is too full to
find one) — mirrors MONAI RandCropByPosNegLabeld(pos=1, neg=1), which is
Pipeline C's protocol. Uniform-only random crops put a ~20 mm tumor blob at
a random edge position in most 96^3 patches and starve the loss of
foreground signal on sparse targets. Unlabeled rows use plain uniform crops.
"""
def __init__(self, rows, patch=(96, 96, 96), train=True, seed=0, pos_neg=0.5):
def __init__(self, rows, patch=(96, 96, 96), train=True, seed=0):
self.rows = list(rows)
self.patch = tuple(patch)
self.train = train
self.pos_neg = pos_neg
self.seed = seed
self._rng_wid = None
self._rng = None
self.rng = np.random.RandomState(seed)
def __len__(self):
return len(self.rows)
@ -73,7 +62,8 @@ class PatchDataset(Dataset):
raise RuntimeError(f"label file unreadable for {row.get('key')}: {row['plabel']}")
if self.train:
s0 = self._crop_origin(img.shape, label)
s0 = [self.rng.randint(0, img.shape[i] - self.patch[i] + 1) if img.shape[i] >= self.patch[i] else 0
for i in range(3)]
img = img[s0[0]:s0[0] + self.patch[0], s0[1]:s0[1] + self.patch[1], s0[2]:s0[2] + self.patch[2]]
if label is not None:
label = label[s0[0]:s0[0] + self.patch[0], s0[1]:s0[1] + self.patch[1], s0[2]:s0[2] + self.patch[2]]
@ -88,41 +78,6 @@ class PatchDataset(Dataset):
lab = torch.from_numpy(label) if label is not None else None
return t, lab, torch.tensor(w, dtype=torch.float32)
def _worker_rng(self):
# forked dataloader workers inherit the parent's rng state; reseed per
# worker so each one draws an independent crop-position stream
info = torch.utils.data.get_worker_info()
wid = info.id if info is not None else 0
if self._rng_wid != wid:
self._rng = np.random.RandomState(self.seed + 104729 * wid)
self._rng_wid = wid
return self._rng
def _crop_origin(self, shape, label):
"""Crop origin: fg-centered for `pos_neg` of labeled rows, otherwise a
foreground-free background crop (20 retries, then plain uniform)."""
rng = self._worker_rng()
win = self.patch
def uniform():
return [int(rng.randint(0, n - w + 1)) if n >= w else 0 for n, w in zip(shape, win)]
has_fg = label is not None and int(label.sum()) > 0
if has_fg and rng.rand() < self.pos_neg:
idx = np.argwhere(label > 0)
c = idx[int(rng.randint(len(idx)))]
s = [int(c[i] - win[i] // 2) for i in range(3)]
return [min(max(s[i], 0), shape[i] - win[i]) for i in range(3)]
s = uniform()
if has_fg:
sl = tuple(slice(s[i], s[i] + win[i]) for i in range(3))
for _ in range(20):
if int(label[sl].sum()) == 0:
return s
s = uniform()
sl = tuple(slice(s[i], s[i] + win[i]) for i in range(3))
return s
@staticmethod
def _augment(img, label):
if np.random.rand() > 0.5:
@ -158,9 +113,8 @@ def _collate(items):
return imgs, labs, wts
def make_dataloader(rows, patch, batch_size, train, num_workers=8, seed=0, persistent=True,
pos_neg=0.5):
ds = PatchDataset(rows, patch=patch, train=train, seed=seed, pos_neg=pos_neg)
def make_dataloader(rows, patch, batch_size, train, num_workers=8, seed=0, persistent=True):
ds = PatchDataset(rows, patch=patch, train=train, seed=seed)
dl = DataLoader(ds, batch_size=batch_size, shuffle=train, num_workers=num_workers,
collate_fn=_collate, drop_last=train and len(ds) > batch_size, pin_memory=True,
persistent_workers=(train and num_workers > 0 and persistent),

View file

@ -1,117 +0,0 @@
"""Row-level exclusion rules for the labeled dataset.
Used at two places, both via exclusion_reason() (single source of truth):
- scripts/make_tv_slices.py — first-stage curation right after the rows are
fetched from the sources (ntuh + m6_labeled): classifies each labeled
volume and writes its QA figure to data/tv_slices/inclusion/ or
data/tv_slices/exclusion/;
- scripts/05_build_splits.py — applied after the patient-level split
assignment so kept rows retain their split; excluded rows are recorded
in data/manifests/excluded.jsonl.
Rules (2026-10-04 curation):
1. total TV volume > MAX_TOTAL_VOL mm3 (implausible targets, e.g. a
full-skull outline instead of a lesion)
2. a single target (1 connected component, 26-conn) with total volume
< MIN_SINGLE_TARGET_VOL mm3 (sub-target fragments)
3. non-brain images (spine/neck CyberKnife cases in
ntuh; their series names are indistinguishable from brain T1c, so the
rows were identified by visual review of the TV-slice figures,
scripts/make_tv_slices.py — 43 rows, see NON_BRAIN_KEYS)
label_stats() / exclusion_reason() are the single source of truth; 05 calls
them for every labeled row and records the excluded rows (with reason) in
data/manifests/excluded.jsonl.
"""
import numpy as np
from scipy import ndimage
MAX_TOTAL_VOL = 1_000_000 # mm3 (1mm isotropic voxels == mm3)
MIN_SINGLE_TARGET_VOL = 100 # mm3
# Non-brain rows identified by visual review of the TV-slice figures on
# 2026-10-04 (spine / neck / torso T1c with vertebral-slab or neck targets;
# series filenames carry no anatomical marker). All ntuh.
#
# Review protocol: classify at FULL resolution (data/tv_slices/exclusion/),
# using the mid-cranial slice (widest head-mask axial), not just the
# TV-max slice — skull-base / posterior-fossa brain targets look like
# cervical spine at the TV-max level (6 such false positives were found and
# removed on 2026-10-04 after expert inspection: ntuh_5902230, ntuh_5231609,
# ntuh_5083686, ntuh_5157194, ntuh_5303758, ntuh_6373987). TV-size gating
# also misses small vertebral targets (e.g. ntuh_5022161, 16k mm3); the
# mid-cranial slice is the only reliable discriminator, so the full ntuh
# population was audited slice-by-slice on 2026-10-05.
NON_BRAIN_KEYS = frozenset({
# cervical / thoracic spine slabs
"ntuh_5237879_case2009.12.10.13.25.32_20091210120310",
"ntuh_5584533_case2011.03.18.13.24.12_20110318115157",
"ntuh_4500992_case2009.08.13.13.51.21_20090813120210",
"ntuh_2922239_case2010.08.23.15.43.52_20100823123647",
"ntuh_0915393_case2012.12.14.17.58.27_20121214121137",
"ntuh_5085147_case2013.05.31.12.51.10_20130531121922",
"ntuh_2486132_case2013.04.24.14.32.34_20130424122008",
"ntuh_4493086_case2009.08.06.13.37.10_20090806114937",
"ntuh_5185487_case2009.07.27.14.18.34_20090727114447",
"ntuh_2177312_case2012.09.10.15.19.55_20120910111101",
"ntuh_5302085_case2014.01.06.14.04.54_20140106123022",
"ntuh_6069590_case2014.08.28.14.06.40_20140828105732",
"ntuh_4408551_case2012.06.28.16.21.19_20120316123153",
"ntuh_4408551_case2012.06.28.16.21.19_20120628114905",
"ntuh_4408551_case2012.03.16.15.05.30_20120316123153",
"ntuh_2444262_case2009.07.07.13.36.39_20090707113235",
"ntuh_5406976_case2009.02.12.14.05.41_20090212112402",
"ntuh_8009538_case2009.11.04.12.46.23_20091104123842",
"ntuh_3563834_case2010.04.20.13.15.45_20100420113442",
"ntuh_5712878_case2011.12.05.14.11.00_20111205120244",
"ntuh_4272908_case2016.06.24.13.33.59_20160624113223",
"ntuh_4494910_case2008.06.03.13.40.34_20080603114804",
"ntuh_2167997_case2016.05.23.13.44.36_20160523114640",
"ntuh_4572226_case2017.11.06.14.11.35_20171106115826",
"ntuh_3420930_case2012.07.30.15.01.30_20120730111910",
"ntuh_4219462_case2010.05.31.13.50.31_20100601120540",
"ntuh_4572226_case2017.11.06.14.12.33_20171106115826",
"ntuh_5022161_case2011.03.02.13.24.52_20110302113819",
"ntuh_5650997_case2011.06.29.13.21.10_20110629113315",
"ntuh_6264497_case2016.09.21.14.48.29_20160921170338",
"ntuh_6182961_case2016.08.08.13.07.21_20160808123711",
"ntuh_2325341_case2010.12.10.13.47.32_20101210122902",
"ntuh_5552572_case2010.04.13.13.27.12_20100413120757",
"ntuh_0711619_case2011.05.10.13.47.39_20110509161659",
"ntuh_5745927_case2013.09.17.13.27.40_20130917115939",
"ntuh_4615006_case2010.07.29.13.43.30_20100722121347",
"ntuh_5332969_case2008.07.21.13.40.07_20080721140145",
"ntuh_5808597_case2012.07.02.15.22.10_20120702115701",
"ntuh_5886224_case2013.12.26.14.59.16_20131226115857",
# byte-identical duplicate images of reviewed non-brain rows above (same
# T1c file registered under a second case-folder name; matched by MD5 on
# 2026-10-05), so they escape the row-level list and the TV-size gates
"ntuh_2177312_case2012.09.10.16.05.26_20120910111101",
"ntuh_2922239_case2010.08.23.14.46.37_20100823123647",
"ntuh_4408551_case2012.06.28.16.14.15_20120628114905",
"ntuh_5085147_case2013.05.31.12.56.43_20130531121922",
})
_CC_STRUCTURE = np.ones((3, 3, 3), dtype=int)
def label_stats(label_path):
"""{total: voxel count, n_cc: connected-component count (26-conn)} for a
processed label (1mm isotropic, so voxels == mm3)."""
import SimpleITK as sitk
m = sitk.GetArrayFromImage(sitk.ReadImage(label_path)) > 0.5
total = int(m.sum())
n_cc = int(ndimage.label(m, structure=_CC_STRUCTURE)[1])
return {"total": total, "n_cc": n_cc}
def exclusion_reason(key, stats):
"""Return the exclusion reason for a labeled row, or None to keep it."""
if stats["total"] > MAX_TOTAL_VOL:
return f"total TV {stats['total']} mm3 > {MAX_TOTAL_VOL} mm3"
if stats["n_cc"] == 1 and stats["total"] < MIN_SINGLE_TARGET_VOL:
return (f"single target, TV {stats['total']} mm3 "
f"< {MIN_SINGLE_TARGET_VOL} mm3")
if key in NON_BRAIN_KEYS:
return "non-brain image (spine/neck, reviewed 2026-10-04)"
return None

View file

@ -22,32 +22,11 @@ class SoftDiceLoss(nn.Module):
return torch.stack(losses, dim=0).mean(dim=0)
def soft_dice_fg(logits, target, eps=1e-5):
"""Soft dice loss on the tumor class (channel 1) only. Averaging the dice
over background+tumor dilutes the foreground signal ~10^4-fold on sparse
targets; the tumor-only variant keeps the same 0.5/0.5 CE+Dice blend but
points all of the Dice term at the tumor."""
p = F.softmax(logits, dim=1)[:, 1]
t = target[:, 1].float()
dim = tuple(range(1, p.dim()))
inter = (p * t).sum(dim)
denom = p.sum(dim) + t.sum(dim)
return 1 - (2 * inter + eps) / (denom + eps)
def per_sample_loss(logits, target, dice_weight=0.5, fg_only=False, fg_ce_weight=1.0):
"""target: (B,C,...) one-hot -> per-sample loss (B,)
fg_ce_weight: class weight [1, w] on the CE part — multiplies the CE on
tumor voxels. Uniform CE on a ~0.05% target leaves the model
under-confident (p_tumor maxes at ~0.4 at a converged loss); weighting the
tumor class sharpens saturation at the cost of more false-positive pressure,
countered by the dice term."""
cls = target.argmax(dim=1)
w = None if fg_ce_weight == 1.0 else torch.tensor([1.0, float(fg_ce_weight)], device=logits.device)
bce = F.cross_entropy(logits, cls, weight=w, reduction="none")
def per_sample_loss(logits, target, dice_weight=0.5):
"""target: (B,C,...) one-hot -> per-sample loss (B,)"""
bce = F.cross_entropy(logits, target.argmax(dim=1), reduction="none")
bce = bce.mean(dim=tuple(range(1, bce.dim())))
dcl = soft_dice_fg(logits, target) if fg_only else SoftDiceLoss()(logits, target)
dcl = SoftDiceLoss()(logits, target)
return (1 - dice_weight) * bce + dice_weight * dcl