From 93e57ab6a32b42b981871b98dddd8b049a2cba2c Mon Sep 17 00:00:00 2001 From: alexprotom Date: Wed, 26 Aug 2026 22:31:35 +0200 Subject: [PATCH 1/5] technical fixes for medsam2 engine --- src/medsam2/engine.rs | 13 +++++++++++++ src/medsam2/model.rs | 17 +++++++++++++++++ tests/medsam2.rs | 44 +++++++++++++++++++++++++++++++++++++------ 3 files changed, 68 insertions(+), 6 deletions(-) diff --git a/src/medsam2/engine.rs b/src/medsam2/engine.rs index b3ca2c1..566eedb 100644 --- a/src/medsam2/engine.rs +++ b/src/medsam2/engine.rs @@ -177,6 +177,19 @@ impl Engine { } } + /// Slices the image encoder has processed since this engine was loaded. + /// + /// Re-prompting a cached slice must leave this unchanged — that is what + /// makes the interactive loop interactive, and the one honest way to + /// assert it, since how *long* a re-prompt takes depends on the machine. + pub fn encode_count(&self) -> usize { + match &self.inner { + Inner::Cpu(model, _) => model.encode_count(), + #[cfg(feature = "gpu")] + Inner::Gpu(model, _) => model.encode_count(), + } + } + /// Forget the cached slice — call this whenever the prepared stack itself /// changes (a different study, or a different intensity window). pub fn clear_cache(&self) { diff --git a/src/medsam2/model.rs b/src/medsam2/model.rs index 805cf5c..1055754 100644 --- a/src/medsam2/model.rs +++ b/src/medsam2/model.rs @@ -12,6 +12,8 @@ //! memory encoder — is another ~26 G but it is strictly sequential, because //! slice *n* needs slice *n-1*'s memory. +use std::sync::atomic::{AtomicUsize, Ordering}; + use anyhow::Result; use burn::tensor::backend::Backend; use burn::tensor::Tensor; @@ -54,6 +56,14 @@ pub struct Medsam2 { /// The sine encoding of the image tokens, `[1, tokens, 256]`. image_pos: Tensor, device: B::Device, + /// How many times [`Medsam2::encode_slice`] has run. + /// + /// The interactive loop's whole premise is the split described at the top + /// of this module: re-prompting a slice must *not* re-encode it. Counting + /// the encodes is how that stays checkable — in a test, or in a profile — + /// without measuring elapsed time, which says as much about the machine + /// as about the code. + encodes: AtomicUsize, } impl Medsam2 { @@ -81,9 +91,15 @@ impl Medsam2 { obj_ptr_tpos_proj: Lin::load(p, "obj_ptr_tpos_proj", MEM_DIM, D_MODEL, dev)?, image_pos, device: dev.clone(), + encodes: AtomicUsize::new(0), }) } + /// Slices encoded since this network was loaded (see [`Self::encodes`]). + pub fn encode_count(&self) -> usize { + self.encodes.load(Ordering::Relaxed) + } + pub fn device(&self) -> &B::Device { &self.device } @@ -91,6 +107,7 @@ impl Medsam2 { /// The image encoder: trunk, neck, and the two high-resolution /// projections the decoder needs. pub fn encode_slice(&self, image: Tensor) -> SliceFeatures { + self.encodes.fetch_add(1, Ordering::Relaxed); let stages = self.trunk.forward(image); let mut levels = self.neck.forward(&stages); // `scalp`: the lowest-resolution level(s) are computed and dropped. diff --git a/tests/medsam2.rs b/tests/medsam2.rs index ea5c749..0fc6a39 100644 --- a/tests/medsam2.rs +++ b/tests/medsam2.rs @@ -329,22 +329,43 @@ fn a_preview_agrees_with_the_prompted_slice_and_reuses_its_features() { ..Config::default() }; + // What "reuses its features" means is that the *image encoder* — the + // expensive half — runs once and once only. Counting the encodes says + // that exactly; timing the calls would only say how loaded the machine + // is, which on a shared CI runner is not a property of this code. + assert_eq!( + engine.encode_count(), + 0, + "nothing encoded before the first prompt" + ); + let t0 = Instant::now(); let preview = engine .preview(&prepared, 1, &prompt, &cfg) .expect("preview"); let cold = t0.elapsed(); assert_eq!(preview.len(), 40 * 48); + assert_eq!( + engine.encode_count(), + 1, + "the first prompt encodes its slice" + ); // The same prompt through the full path must decide that slice the same - // way — the preview is the propagation's first step, not an approximation. + // way — the preview is the propagation's first step, not an approximation + // — and must take the encoded slice from the cache rather than redo it. let full = engine .propagate(&prepared, 1, &prompt, &cfg, &Quiet) .expect("propagate"); assert_eq!(full.masks[1], preview); assert_eq!(full.slices_visited, 1); + assert_eq!( + engine.encode_count(), + 1, + "propagating reused the cached slice" + ); - // And the second call skipped the encoder entirely. + // And so does a second, different prompt on that same slice. let t1 = Instant::now(); let again = engine .preview( @@ -356,11 +377,22 @@ fn a_preview_agrees_with_the_prompted_slice_and_reuses_its_features() { .expect("preview"); let warm = t1.elapsed(); assert_eq!(again.len(), 40 * 48); - assert!( - warm * 2 < cold, - "a cached slice should re-prompt far faster: {cold:?} then {warm:?}" - ); + assert_eq!(engine.encode_count(), 1, "re-prompting must not re-encode"); eprintln!("preview: {cold:?} cold, {warm:?} warm"); + // Clearing the cache is what makes the next prompt pay for the encoder + // again — otherwise the count above would prove nothing. + engine.clear_cache(); + engine + .preview(&prepared, 1, &prompt, &cfg) + .expect("preview"); + assert_eq!(engine.encode_count(), 2, "a cleared cache re-encodes"); + + // A different slice is a different cache entry. + engine + .preview(&prepared, 0, &prompt, &cfg) + .expect("preview"); + assert_eq!(engine.encode_count(), 3, "another slice encodes on its own"); + engine.clear_cache(); } From 5a6cbb1086e6055d4fd819129de5d2d85a3c5c5a Mon Sep 17 00:00:00 2001 From: alexprotom Date: Thu, 27 Aug 2026 11:18:59 +0200 Subject: [PATCH 2/5] Added body segmentation --- Cargo.toml | 2 +- README.md | 16 +- docs/README.md | 1 + docs/architecture.md | 43 +- docs/body-contour.md | 331 +++++++++++++++ docs/segmentation.md | 5 +- examples/autoseg_probe.rs | 2 +- examples/body_cli.rs | 170 ++++++++ src/app/body_win.rs | 556 +++++++++++++++++++++++++ src/app/chrome.rs | 13 +- src/app/dialogs.rs | 1 + src/app/mod.rs | 18 + src/app/models_win.rs | 2 +- src/app/panels.rs | 8 +- src/app/seg.rs | 12 +- src/app/seg_engines.rs | 21 +- src/autoseg/config.rs | 72 +++- src/autoseg/cpu.rs | 2 +- src/autoseg/gpu.rs | 10 +- src/autoseg/infer.rs | 1 + src/autoseg/mod.rs | 78 +++- src/autoseg/net.rs | 16 +- src/autoseg/preprocess.rs | 20 +- src/autoseg/weights.rs | 33 ++ src/bodymask.rs | 856 ++++++++++++++++++++++++++++++++++++++ src/lib.rs | 2 + src/models.rs | 8 + src/morphology.rs | 799 +++++++++++++++++++++++++++++++++++ src/nn/tensor.rs | 97 +++++ tests/autoseg.rs | 2 + tests/body.rs | 408 ++++++++++++++++++ 31 files changed, 3539 insertions(+), 66 deletions(-) create mode 100644 docs/body-contour.md create mode 100644 examples/body_cli.rs create mode 100644 src/app/body_win.rs create mode 100644 src/bodymask.rs create mode 100644 src/morphology.rs create mode 100644 tests/body.rs diff --git a/Cargo.toml b/Cargo.toml index c63582f..8db9542 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "rust-dicom-station" -version = "0.5.0" +version = "0.1.0" edition = "2021" description = "Fast, robust DICOM / RT DICOM viewer in pure Rust (CT/MR volumes, RTSTRUCT, RTDOSE, RTPLAN) with a three-view MPR layout" license = "MIT" diff --git a/README.md b/README.md index 967205e..479ca1b 100644 --- a/README.md +++ b/README.md @@ -10,7 +10,9 @@ three-view layout, with a second dataset row for comparison, built-in **image registration** (elastix- and plastimatch-style, rigid, deformable and landmark-based, with analytics and a deformation vector field), **structure propagation**, **DRR generation**, **interactive -segmentation**, a live **3D structure view**, **automatic multi-organ +segmentation**, a live **3D structure view**, **automatic body / EXTERNAL +contouring** (the patient outline without the couch, the chair or the +immobilisation, on CT and MR), **automatic multi-organ segmentation** (a pure-Rust re-implementation of TotalSegmentator, 117 structures, CPU or any GPU), **prompt-driven segmentation** (a pure-Rust re-implementation of SegVol — point at anything with a box, a @@ -78,6 +80,17 @@ holds the registration controls and both dataset trees.* view, mask → RTSTRUCT conversion, and **DICOM SEG** import and export (binary and fractional multi-frame masks, read onto their own lattice and resampled onto whichever image series they belong to). +* **Body contouring** - the EXTERNAL structure, found automatically and + without the couch, the chair or the immobilisation inside it. Equipment + is separated from anatomy by two facts no patient has together: it is + thin, and its footprint repeats slice after slice - so an 8 mm opening + and a persistence test along **all three** axes catch a supine couch top + and an upright chair's seat pan alike, while ears, nose and fingers are + given back afterwards. Works on CT by Hounsfield threshold and on MR + after flattening the coil shading; optionally guided by + TotalSegmentator's openly licensed body network, which is what removes a + mask touching the skin with no gap. Lands as an editable mask and as an + RTSTRUCT `EXTERNAL`. * **Auto-segmentation** - TotalSegmentator v2 inference rebuilt natively: official nnU-Net weights downloaded once and converted without Python, hand-written SIMD CPU engine and a wgpu GPU path @@ -169,6 +182,7 @@ ships a real two-phase 4DCT (see | [docs/propagation.md](docs/propagation.md) | Carrying contours and segmentations across a registration | | [docs/drr.md](docs/drr.md) | Digitally reconstructed radiographs: the two projectors and the geometry | | [docs/segmentation.md](docs/segmentation.md) | Brush / eraser / region growing, 3D view, mask → RTSTRUCT | +| [docs/body-contour.md](docs/body-contour.md) | The body / EXTERNAL contour: the classical and model-assisted methods, CT and MR, verification | | [docs/segvol.md](docs/segvol.md) | Prompt-driven segmentation: box / point / text, the SegVol re-implementation | | [docs/medsam2.md](docs/medsam2.md) | Propagating a prompt through a stack: the MedSAM2 re-implementation | | [docs/auto-segmentation.md](docs/auto-segmentation.md) | The pure-Rust TotalSegmentator: models, pipeline, engines, validation, classes, licensing | diff --git a/docs/README.md b/docs/README.md index 8b407d3..78b1bc6 100644 --- a/docs/README.md +++ b/docs/README.md @@ -10,6 +10,7 @@ covers one area in depth. | [rt-objects.md](rt-objects.md) | RT DICOM objects: RTSTRUCT, RTDOSE, RTPLAN, REG spatial registrations, RT treatment records, and how their reference chains are resolved | | [registration.md](registration.md) | Rigid and deformable (B-spline) image registration: algorithms, parameters, the fusion overlay, the transform simulator for registration QA, accuracy verification | | [segmentation.md](segmentation.md) | Interactive segmentation: 2D/3D brush, eraser, geodesic region growing, the live 3D structure view, mask → RTSTRUCT conversion | +| [body-contour.md](body-contour.md) | Automatic body / EXTERNAL contouring: why the couch, the chair and the mask are the hard part, the classical threshold-and-morphology method, the model-assisted method built on TotalSegmentator's body network, CT and MR, verification | | [auto-segmentation.md](auto-segmentation.md) | Automatic multi-organ segmentation — the pure-Rust TotalSegmentator re-implementation: models, usage, the inference pipeline, CPU/GPU engines, validation, the full 117-class table, licensing | | [segvol.md](segvol.md) | Prompt-driven segmentation — the pure-Rust SegVol re-implementation: box / point / text prompts, the two-pass pipeline, weights and licensing, validation status | | [export-and-tools.md](export-and-tools.md) | DICOM export, the interactive anonymizer, the synthetic test-data generator | diff --git a/docs/architecture.md b/docs/architecture.md index 5c78819..af0a366 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -101,6 +101,11 @@ rust-dicom-station │ active registration; the recovered field written back out as one │ ├── Segmentation +│ ├── Body / EXTERNAL contour: thresholding by modality (HU, or MR after bias +│ │ flattening), spacing-aware opening, extruded-equipment removal along all +│ │ three axes, component selection by volume, thin-anatomy recovery, +│ │ slice-wise filling — classically, or guided by TotalSegmentator's body +│ │ network (CT 6 / 1.5 mm, MR) │ ├── Voxel masks: brush / eraser (2D, 3D), geodesic region growing, undo, │ │ slice overlays, hole filling, mask ▶ RTSTRUCT contours, RTSTRUCT ▶ mask, │ │ grouped into segmentation series that live in the study and bind to an @@ -204,7 +209,9 @@ src/ drr_win.rs the DRR window: geometry, projectors, comparison seg.rs interactive segmentation state machine, mask ▶ RTSTRUCT, landing an auto-segmentation result - seg_engines.rs what the three engine windows share: names and glyphs, + body_win.rs the body-contour window: method choice, the modality's + own threshold row, the classical / model-assisted split + seg_engines.rs what the four tool windows share: names and glyphs, device / model-folder / licence / progress rows, result landing, the "still the same dataset" check prompt_seg.rs prompt segmentation window and worker (SegVol) @@ -248,6 +255,15 @@ src/ segmentation.rs voxel masks: brush, geodesic grow, undo, overlays, label map ▶ segmentations, mask ▶ RTSTRUCT contours and RTSTRUCT contours ▶ mask Seg + morphology.rs binary-mask geometry, in millimetres: the exact + anisotropic Euclidean distance transform and the + erode / dilate / open / close it powers, 6-connected + components, slice-wise and 3-D hole filling, the + extruded-equipment test, box-blur smoothing Core + bodymask.rs the body / EXTERNAL contour: foreground by modality (HU, + or bias-flattened MR), equipment removal, component + selection, thin-anatomy recovery, filling — classically + or guided by the body network Seg mesh3d.rs contour / mask ▶ surface meshes (scanline fill, surface nets, Laplacian smoothing) Seg @@ -261,12 +277,15 @@ src/ params.rs shape-checked view of a loaded state dict half.rs binary16 ↔ binary32 conversion tensor.rs Mat [rows, cols] and Act [c, d, h, w]; transposed conv + (a hand-tuned 2× and a general kernel = stride form) linalg.rs gemm-backed linear / matmul, layer norm, softmax, GELU / ReLU / QuickGELU attention.rs multi-head attention, optionally causally masked autoseg/ automatic segmentation (pure-Rust TotalSegmentator) Seg - mod.rs public API: variants, run(), progress phases + mod.rs public API: variants, run(), run_specs() (the engine + minus the question, shared with the body contour), + progress phases classes.rs 117-class table, sub-model maps, organ colors config.rs nnU-Net plans.json parsing weights.rs which models exist, where they are published, the @@ -354,14 +373,15 @@ those: a ROI visibility toggle, for instance, is part of the contour key alone and leaves the dose and fusion textures untouched. Repaints are demand-driven; while background jobs run, the UI polls at 10 Hz. -### The three engine windows +### The segmentation tool windows -Auto-segmentation, prompt segmentation and slice propagation are different -conversations — a batch run with a result-selection dialog, a one-shot -prompt, an interactive box loop — but they are the same kind of tool, and -`app/seg_engines.rs` makes them look and behave alike: +Body contouring, auto-segmentation, prompt segmentation and slice +propagation are different conversations — a parameterised geometric run, a +batch run with a result-selection dialog, a one-shot prompt, an interactive +box loop — but they are the same kind of tool, and `app/seg_engines.rs` +makes them look and behave alike: -* one `ToolInfo` per engine gives the glyph (🤖 🧠 ⏩), the window title +* one `ToolInfo` per tool gives the glyph (👤 🤖 🧠 ⏩), the window title (`🤖 Auto-segmentation — dataset A`, the same pattern as `3D structures — dataset A`), the menu entry (`🤖 Auto-segment dataset A…`) and the small sidebar button (`🤖 Auto…`); @@ -374,8 +394,11 @@ prompt, an interactive box loop — but they are the same kind of tool, and engine, the tool's own inputs, `Name`, a collapsed **Options** header holding the engine's settings plus the shared `Compute: Auto / GPU / CPU` and `Model folder` rows, one small licence line ("… Research / QA use — - not a medical device."), then `▶ Segment` / `▶ Propagate` and `Close`, - and a status line summarising the last result; + not a medical device."), then `▶ Segment` / `▶ Propagate` / `▶ Contour` + and `Close`, and a status line summarising the last result; +* what a tool does not have, it does not show: the body contour's classical + method needs no device, no model and no download, so those rows appear + only when its method is the model-assisted one; * results land the same way — `add_segmentation` with the next palette colour — and a run that finishes after the dataset was replaced is discarded with the same message. diff --git a/docs/body-contour.md b/docs/body-contour.md new file mode 100644 index 0000000..4de4215 --- /dev/null +++ b/docs/body-contour.md @@ -0,0 +1,331 @@ +# Automatic body / EXTERNAL contouring + +The outer patient surface, with the couch, the chair and the immobilisation +left outside it. Two methods in one tool — a deterministic geometric one +that needs nothing, and a model-assisted one that borrows +TotalSegmentator's openly licensed body network for the part geometry +cannot decide. Both work on CT and on MR, supine or upright. + +## Why this contour and not another + +Every downstream calculation starts here. A dose engine needs to know where +the patient begins, because that is where the range budget starts — for +protons a 3 mm error in the entrance surface is a 3 mm error in every +distal edge behind it. A DRR needs to know what is *not* patient, or it +projects the couch straight through the anatomy. A registration wants to +sample inside the body and nowhere else, or it spends its iterations +aligning the table. + +And it is the contour nobody wants to draw. So it has to be right on the +first scan of the day, with no parameters touched. + +## What makes it hard + +Not the skin — that is the largest step in the image and any threshold +finds it. The difficulty is everything else in the field of view: + +* a **couch top** is two carbon skins around a foam core, and the skins are + as dense as bone; +* a **thermoplastic mask**, a **headrest shell**, a **vacuum-bag** fabric + and a **chair backrest** are all thin, dense and *touching the patient*, + so in a threshold mask they are one connected object with them; +* **blankets, cables, positioning pads** drift in and out of the field; +* **the reconstruction circle** leaves a bright rim on some scanners; +* and the patient is not always one object — a leg scan is two, an arm cut + off by the field of view is another. + +Meanwhile the parts of the patient that look most like equipment — an ear, +a nose, a fingertip — are exactly the thin structures that any method +aggressive enough to remove a mask will also remove. + +## Using it + +*Tools ▶ 👤 Body-contour dataset A/B…*, or the **👤 Body…** button in the +sidebar *Segmentations* section. The window follows the shared layout of +the segmentation tools (see +[architecture.md](architecture.md#the-segmentation-tool-windows)). + +* **Method** — *Classical* or *Model-assisted*; everything below adapts. +* **Tissue above** — on CT a Hounsfield threshold (default −300 HU); on MR + a fraction of the bias-corrected 99th percentile, or Otsu. The window + re-seeds this whenever the displayed series changes modality, because a + CT threshold is meaningless on MR and the reverse is just as wrong. +* **Name** and **as EXTERNAL structure** — the mask lands as an ordinary + editable segmentation, and optionally also as an RTSTRUCT ROI of + interpreted type `EXTERNAL`, which is the tag a planning system looks for. + It then renders like any ROI and rides the DICOM export. +* **Options** — the smallest body detail, the equipment test with its shell + thickness and repeat window, the smallest body part, thin-anatomy + recovery, whether the body is reported solid, surface smoothing, and + (model-assisted only) the network margin, compute device and model + folder. + +**▶ Contour** runs on a background thread with the usual progress row and +Cancel. The status line reports the body volume, how many separate bodies +were kept, how much equipment was removed and how much thin anatomy was +given back — which is the number to watch: a run that removes nothing is a +run whose threshold or opening radius is wrong. + +## Method A — classical + +Deterministic, nothing to download, a few seconds on a whole-body CT, and +every step explainable to a physicist doing QA. + +### 1. Foreground + +CT thresholds directly: −300 HU sits in the gap between air and fat. (The +skin edge moves about half a millimetre per 100 HU through the +partial-volume ramp — this is the one number worth agreeing on with your +planning system.) + +MR has no absolute scale: the same tissue is a different number on the next +sequence, and a different number again on the other side of the same slice, +because the receive coils shade the image. So the MR path divides the image +by a heavy blur of itself (σ = 40 mm by default) — a poor man's N4, no +iteration and no histogram model — which flattens the shading while leaving +every edge intact, then thresholds at a low fraction of the 99th percentile. +A low fraction rather than Otsu by default: Otsu splits *bright from dark*, +not tissue from air, so it runs high on fat-suppressed series and bites into +subcutaneous fat. + +### 2. Equipment, by two geometric facts + +Equipment is separated from anatomy without knowing what either looks like, +using two properties no patient has together: + +**It is thin.** An opening removes every shell whose largest inscribed ball +is smaller than its radius, while leaving everything thicker with its +surface *exactly* intact — an opening is the union of every ball that fits +inside the mask, so the ball rolls along the inside of the skin and touches +every point of it. The distance transform behind it is the exact +anisotropic Euclidean one, so the radius means the same along every axis +whatever the slice thickness, and the cost does not depend on it. + +The radius here is **2 mm**, not the 8 mm used later to decide what is big +enough to be a body, and the difference matters more than anything else on +this page. A couch skin is one or two millimetres of carbon and a +thermoplastic mask two or three; the thinnest tissue anyone would miss — +the chest wall over a lung — is five or six. At 2 mm the two are cleanly +separated. At 3 mm a six-millimetre chest wall becomes a candidate too, and +since it repeats slice after slice it is then indistinguishable from a +couch skin: the ribcage goes with the table. (This is not hypothetical. It +is what the first version of this code did to the bundled 4D-Lung study, +and it is why the test suite contains a hollow cylinder.) + +**It is extruded.** A couch top, a backrest, a seat pan, an arm rest: each +is a surface swept along one axis, so its footprint in the orthogonal plane +repeats slice after slice after slice. A pinna is 25 mm long, a nose 30 mm, +a fingertip 15 mm. Requiring a footprint to repeat over 150 mm in 80 % of +that window's slices separates them with room to spare — and the test runs +along **all three** axes, because a supine couch is extruded along z while +an upright chair's seat pan and arm rests are extruded along x. + +Being thin alone is not enough to be discarded, and neither is repeating; +only both together mark a voxel as equipment. + +### 3. A body is a solid object + +A threshold does not see a body. It sees a shell of tissue wrapped round two +lungs, a stomach and a bowel — and left that way, the chest wall over a lung +is a thin sheet that repeats slice after slice, which is to say +indistinguishable from a couch skin by the rules just described. So the +interior is closed *before* any of the size reasoning below: fill what the +slice border cannot reach, and the wall becomes part of a solid object +again. + +Slice by slice, not in three dimensions — the lungs drain to the outside air +through the trachea, so on any scan that includes the neck a 3-D fill leaves +both lungs open, while slice by slice they close. (An open mouth stays a +cavity, which is the usual convention, and so does a lung on the two or +three slices where the airway is actually open to the air.) + +The order is the point: **equipment first, then fill**. Fill first and a +couch top with a closed profile becomes a solid slab before anything has a +chance to recognise it. + +### 4. Which components are a patient + +The opened mask is split into 6-connected components, and every component +of at least 50 cm³ is kept — *not* merely the largest, so a leg scan comes +out as two bodies and a truncated arm as a third. (A shared corner is not +contact: 6-connectivity is what stops a couch rail grazing the skin +diagonally from merging with it.) + +### 5. Giving the thin anatomy back + +The 8 mm opening shaved a rim off the body, and took the ears with it. Two +questions put it back, because there are two different things in there. + +What the opening removed from the body's **own surface** lies, by +construction, within one opening radius of what is left of it — a skin rim, +the edge of a shoulder, the sharp flank of a cross-section. It can run the +whole length of the scan and still be nothing but patient, so its size is +not asked about; it is simply given back. + +What stands **clear** of that is a separate object that happens to touch: +an ear, a nose or a fingertip, which are small, or a pad, a blanket or a +bolus, which are not. There, size is exactly the right question, and +anything more than 100 mm across stays out. Two rounds, because a fingertip +hangs off a finger. + +### 6. Surface + +An optional closing at the end takes the staircase off the contour. + +### Where it fails, stated plainly + +Where a shell touches the skin with **no air gap at all** — the mask on the +forehead and chin, a bare couch skin under the back, a bolus — the shell's +thickness over the contact patch stays inside the body. It is 2–5 mm, over +the contact patches only, and it is not detectable by geometry, because +locally it *is* a slightly thicker patient. In practice cushions mean bare +couch contact is rare, and bolus in the external is the convention anyway. +The model-assisted method is the answer to the rest. + +## Method B — model-assisted + +TotalSegmentator publishes a **body-outline nnU-Net** under the same +Apache-2.0 licence as its "total" task, in three flavours: + +| Model | Dataset | Grid | Download | +|---|---|---|---| +| CT 6 mm | 300 | 6 mm isotropic | 124 MB | +| CT 1.5 mm | 299 | 1.5 mm isotropic | 233 MB | +| MR | 597 | 3.0 × 1.19 × 0.99 mm | 230 MB | + +They run through the *same* engine as the 117-class auto-segmentation — +`autoseg::run_specs`, the same `PlainConvUNet` rebuilt from `plans.json`, +the same sliding window, the same CPU/GPU choice — and their weights live +beside it in `models/totalsegmentator/`. The model manager lists them like +any other. + +The network is **not** used as the answer. It is planned at 6 mm or 1.5 mm; +its boundary is far too coarse to be a skin surface. It is used as a +*classifier*: + +``` +body = threshold(image) ∧ dilate(network_body, 6 mm) +``` + +The network decides **what** is patient — it removes a mask contact patch or +a couch sliver semantically, which no geometry can — and the threshold still +decides **where** the skin is, at full image resolution. The result then goes +through exactly the same components / thin-recovery / fill / closing steps as +the classical method. + +The network's own two classes (`body_trunc`, `body_extremities`) are cleaned +first the way the reference implementation cleans them: the trunk is one +object, so only its largest blob survives; extremities are several, so they +are filtered at 50 000 mm³ — the same constant used upstream. The body is +the union. + +The equipment test still runs. It has little left to do once the network has +answered, but the guide is used *dilated* — 6 mm by default — and a margin +that generous can pull a touching rail back in. Two cheap passes over the +volume are a fair price for not having to think about that. + +### What is new in the engine for this + +Supporting the MR body model meant three additions, all of them additive — +the existing models take exactly the code path they took before, so their +numerics are untouched: + +* `ZScoreNormalization` alongside `CTNormalization`. Where CT normalizes + against dataset constants from `plans.json`, MR normalizes against *this + image*, so its constants are only knowable after resampling. +* **Anisotropic target spacing.** The MR model plans 3.0 × 1.19 × 0.99 mm; + `SarMap` now takes a spacing per axis rather than one number. +* **A general transposed convolution.** The MR decoder upsamples + `[1, 2, 2]` at two of its five stages. The hand-tuned 2× routine is still + what every isotropic model uses; a general `kernel = stride` version + handles the rest, on CPU and through burn on the GPU. + +## Cost + +The classical method is a handful of distance transforms and flood fills. +Measured on the bundled 4D-Lung study — 512 × 512 × 133 at 0.98 × 0.98 × +3 mm, on two throttled cores — it takes **13.8 s** end to end and reports a +23.2 L body with 85 cm³ of couch left out. It allocates about one byte per +voxel per intermediate mask, plus four bytes per *set* voxel for the +component lists. + +The model-assisted method adds one nnU-Net inference: 34 s for the 6 mm +model on the same two cores (50 s in total), minutes with the 1.5 mm or MR +model. Since the network only has to say which side of the skin a voxel is +on, **6 mm is the sensible default** — the resolution comes from the +threshold, not from it. + +## Verification + +First, on real data. The bundled 4D-Lung study carries a real couch rail at +the bottom of the field; the classical method removes it on every slice and +follows the skin to the voxel, including the three separate pieces — two +arms and the neck — that the most superior slices contain. + +The two methods are also each other's check, and on that study they pass +it: run separately, the classical geometry and the 6 mm network agree on +**8 098 425 of 8 098 443 voxels** — eighteen voxels apart, Dice 0.999999. +Neither was tuned against the other; they simply have to be looking at the +same surface. + +Then `tests/body.rs`, which builds a phantom containing every failure mode +deliberately: an elliptical body, a couch skin and a rail under the back, a +*moulded* mask shell that stands 2 mm clear of the skin over most of its +span and presses against it over a patch, ears thin enough for the opening +to shave off, lungs draining through an airway, a cable — with anisotropic +2 × 2 × 5 mm voxels, since that is where a voxel-counting implementation +goes wrong and a millimetre-aware one does not. It asserts: + +* Dice > 0.99 against the body it was built from; +* no couch skin, rail, cable or free-standing mask shell anywhere; +* the shell **is** kept where it presses against the skin, and the total + error beyond the patient stays under 10 cm³ — the documented limitation, + pinned so that a change which quietly makes it worse is caught; +* both ears kept, and a non-zero recovered-anatomy count; +* both lungs inside the body on every slice past the airway; +* a hollow cylinder — a 6 mm wall around a cavity, beside a 2 mm couch skin + — comes out with its wall intact and its couch gone, which is the + regression test for the failure real data taught; +* two separated legs come out as two bodies, not as the larger one; +* an MR version with an exponential receive gradient still comes out whole + at both ends of the field (Dice > 0.97); +* a model folder that cannot exist is an error, not a panic. + +`the_model_assisted_method_runs_the_published_network` is `#[ignore]`d +because it downloads 124 MB; it runs the real Dataset300 weights end to end +through the hybrid. + +`src/morphology.rs`'s own tests check the pieces underneath: the distance +transform against brute force on an anisotropic grid, the opening against a +sheet and a block, 6-connectivity against a shared corner, slice-wise versus +3-D filling, the persistence test against an extruded rail and a bump, and +the blur against a constant and a step. + +## Command-line tool + +``` +cargo run --release --example body_cli -- \ + [--method classical|model] [--model ct6|ct15|mr] \ + [--hu -300] [--mr-fraction 0.12] [--mr-otsu] [--bias-sigma 40] \ + [--open 8] [--thin-shell 2] [--no-devices] [--window 150] [--frac 0.8] \ + [--min-cm3 50] [--no-thin] [--thin-extent 100] [--margin 6] \ + [--thin-shell 2] [--no-fill] [--close 0] \ + [--models DIR] [--device auto|gpu|cpu] [--out mask.bin] +``` + +For batch checks over a folder of scans: `--out` writes a raw `u8` mask on +the original grid, one byte per voxel in `Volume::data` order — the same +convention as the other example tools, so masks can be compared byte for +byte between methods. + +## Licensing and citation + +The classical method has no weights and no third-party code. The +model-assisted method uses TotalSegmentator's `body` and `body_mr` tasks, +which the authors publish under **Apache-2.0** as openly available for any +usage, commercial included. If you use it in academic work, cite +TotalSegmentator and nnU-Net as in +[auto-segmentation.md](auto-segmentation.md#licensing-and-citation). + +As with everything in this viewer: research and QA use — not a medical +device, not for clinical decision-making. diff --git a/docs/segmentation.md b/docs/segmentation.md index 517ce5d..a083463 100644 --- a/docs/segmentation.md +++ b/docs/segmentation.md @@ -3,8 +3,9 @@ MITK-style manual and semi-automatic segmentation, implemented entirely in Rust and CPU-side, plus a Slicer-style 3D surface view that follows every edit in essentially real time. For the neural-network auto-segmentation -see [auto-segmentation.md](auto-segmentation.md) — its results land as the -same editable masks described here. +see [auto-segmentation.md](auto-segmentation.md), and for the patient +outline [body-contour.md](body-contour.md) — both land as the same editable +masks described here. ## Segmentation masks diff --git a/examples/autoseg_probe.rs b/examples/autoseg_probe.rs index 016a643..5dd434c 100644 --- a/examples/autoseg_probe.rs +++ b/examples/autoseg_probe.rs @@ -34,7 +34,7 @@ fn main() -> anyhow::Result<()> { cfg.n_stages() ); - let map = preprocess::SarMap::new(vol, cfg.spacing[0]); + let map = preprocess::SarMap::new(vol, cfg.spacing); let vm = preprocess::resample_to_model(vol, &map); eprintln!("model grid {:?}", map.model_dims); // extract the corner patch (like the first sliding-window tile) diff --git a/examples/body_cli.rs b/examples/body_cli.rs new file mode 100644 index 0000000..0a0c69e --- /dev/null +++ b/examples/body_cli.rs @@ -0,0 +1,170 @@ +//! Headless body / EXTERNAL contouring: point it at a DICOM folder. +//! +//! ```text +//! cargo run --release --example body_cli -- \ +//! [--method classical|model] [--model ct6|ct15|mr] \ +//! [--hu -300] [--mr-fraction 0.12] [--mr-otsu] [--bias-sigma 40] \ +//! [--open 8] [--no-devices] [--window 150] [--frac 0.8] \ +//! [--min-cm3 50] [--no-thin] [--thin-extent 100] [--margin 6] \ +//! [--thin-shell 3] [--no-fill] [--close 0] \ +//! [--models DIR] [--device auto|gpu|cpu] [--out FILE] +//! ``` +//! +//! `--out` writes a raw `u8` mask on the original volume's grid, one byte +//! per voxel in `Volume::data` order — the same convention as the other +//! example tools, so the masks can be compared byte for byte. +//! +//! Batch use, which is what this is for: run it over a folder of upright +//! chair scans and check the reported equipment volume. A run that removes +//! nothing is a run whose threshold or opening radius is wrong. + +use std::path::PathBuf; + +use rust_dicom_station::bodymask::{self, BodyModel, BodyParams, Foreground, Method}; +use rust_dicom_station::loader; +use rust_dicom_station::models::{self, Engine}; +use rust_dicom_station::nn::device::DevicePref; +use rust_dicom_station::progress::Progress; + +mod common; + +fn main() -> anyhow::Result<()> { + let mut args = std::env::args().skip(1); + let mut dicom: Option = None; + let mut models_dir: Option = None; + let mut out: Option = None; + let mut p = BodyParams::default(); + // Set by --hu / --mr-*; otherwise the series' modality decides. + let mut foreground: Option = None; + let mut bias_sigma = 40.0f64; + let mut mr_otsu = false; + let mut mr_fraction: Option = None; + let mut model_forced = false; + while let Some(a) = args.next() { + let mut next = || args.next().expect("missing value"); + match a.as_str() { + "--method" => { + p.method = match next().as_str() { + "classical" => Method::Classical, + "model" => Method::ModelAssisted, + other => panic!("--method classical|model, not {other:?}"), + } + } + "--model" => { + model_forced = true; + p.model = match next().as_str() { + "ct6" => BodyModel::Ct6mm, + "ct15" => BodyModel::Ct15mm, + "mr" => BodyModel::Mr, + other => panic!("--model ct6|ct15|mr, not {other:?}"), + } + } + "--hu" => foreground = Some(Foreground::Hu(next().parse().expect("number"))), + "--mr-fraction" => mr_fraction = Some(next().parse().expect("number")), + "--mr-otsu" => mr_otsu = true, + "--bias-sigma" => bias_sigma = next().parse().expect("number"), + "--open" => p.open_mm = next().parse().expect("number"), + "--no-devices" => p.remove_devices = false, + "--window" => p.persist_window_mm = next().parse().expect("number"), + "--frac" => p.persist_frac = next().parse().expect("number"), + "--min-cm3" => p.min_volume_cm3 = next().parse().expect("number"), + "--no-thin" => p.recover_thin = false, + "--thin-extent" => p.thin_max_extent_mm = next().parse().expect("number"), + "--margin" => p.guide_margin_mm = next().parse().expect("number"), + "--no-fill" => p.fill_interior = false, + "--thin-shell" => p.device_thin_mm = next().parse().expect("number"), + "--close" => p.close_mm = next().parse().expect("number"), + "--models" => models_dir = Some(PathBuf::from(next())), + "--device" => { + p.device = DevicePref::from_key(&next()).expect("auto, gpu or cpu"); + } + "--out" => out = Some(PathBuf::from(next())), + other if dicom.is_none() => dicom = Some(PathBuf::from(other)), + other => panic!("unexpected argument {other:?}"), + } + } + let dicom = dicom.expect("usage: body_cli [options]"); + let models_dir = models_dir + .unwrap_or_else(|| models::engine_dir(&models::default_root(), Engine::TotalSegmentator)); + + let progress = Progress::default(); + let study = loader::load_directory(&dicom, &progress)?; + let modality = study + .series + .get(study.active_series) + .map(|s| s.modality.to_uppercase()) + .unwrap_or_default(); + eprintln!( + "{} {:?} {} × {} × {} at {:.2} × {:.2} × {:.2} mm", + modality, + study.meta.patient_id, + study.volume.dims[0], + study.volume.dims[1], + study.volume.dims[2], + study.volume.spacing[0], + study.volume.spacing[1], + study.volume.spacing[2], + ); + + // Modality-appropriate defaults, overridden by whatever was asked for. + p.foreground = Foreground::for_modality(&modality); + if !model_forced { + p.model = BodyModel::for_modality(&modality); + } + if let Some(f) = foreground { + p.foreground = f; + } else if mr_otsu { + p.foreground = Foreground::MrOtsu { + sigma_mm: bias_sigma, + }; + } else if let Some(fraction) = mr_fraction { + p.foreground = Foreground::MrRelative { + fraction, + sigma_mm: bias_sigma, + }; + } + eprintln!("method {:?}, foreground {:?}", p.method, p.foreground); + + let ap = Progress::default(); + let t = std::time::Instant::now(); + let done = std::sync::atomic::AtomicBool::new(false); + let result = std::thread::scope(|s| { + let (done_ref, ap_ref) = (&done, &ap); + let printer = s.spawn(move || { + let mut last = String::new(); + while !done_ref.load(std::sync::atomic::Ordering::Relaxed) { + let msg = ap_ref.get(); + if msg != last { + eprintln!( + "[{:6.1}s] {:5.1}% {}", + t.elapsed().as_secs_f64(), + ap_ref.frac() * 100.0, + msg + ); + last = msg; + } + std::thread::sleep(std::time::Duration::from_millis(500)); + } + }); + let r = bodymask::contour_body(&study.volume, &p, &models_dir, &ap); + done.store(true, std::sync::atomic::Ordering::Relaxed); + let _ = printer.join(); + r + }); + let result = result?; + eprintln!( + "{:.0} cm3 body in {:.1} s{}; {} piece(s); removed {} voxels of equipment, \ + recovered {} thin voxels", + result.cm3, + result.elapsed_secs, + if result.device.is_empty() { + String::new() + } else { + format!(" on {}", result.device) + }, + result.pieces.len(), + result.removed_voxels, + result.recovered_voxels, + ); + common::finish_mask(&result.mask, &study.volume, out.as_deref()) +} diff --git a/src/app/body_win.rs b/src/app/body_win.rs new file mode 100644 index 0000000..8a40636 --- /dev/null +++ b/src/app/body_win.rs @@ -0,0 +1,556 @@ +//! The body-contour tool window — the fourth of the segmentation tools, and +//! the only one that can answer without a network. +//! +//! It shares its bones with the other three ([`super::seg_engines`]): a +//! description, the tool's own inputs, an `Options` section, a licence line +//! and a button row that becomes a progress row while a run is in flight. +//! What it does differently is that its *method* choice changes what the +//! rest of the window means — the classical method has no device, no model +//! and nothing to download, so those rows appear only when they matter. +//! +//! The window re-seeds its thresholds from the displayed series' modality +//! whenever that changes, because −300 HU is meaningless on MR and a +//! fraction of the 99th percentile is meaningless on CT. A threshold the +//! user has edited by hand is left alone. + +use std::path::PathBuf; +use std::sync::Arc; + +use crate::bodymask::{self, BodyModel, BodyParams, BodyResult, Foreground, Method}; +use crate::models::Engine as ModelsEngine; +use crate::progress::Progress; + +use super::*; + +/// The tool window's state; it stays open across runs. +pub(super) struct BodyDialog { + pub slot: usize, + pub params: BodyParams, + /// The modality the parameters were seeded from, so a series switch can + /// re-seed them — and only then. + pub seeded_for: String, + /// One-line summary of the last finished run. + pub status: Option, +} + +/// Everything a run needs, snapshotted from the window when it starts. +struct BodyRequest { + params: BodyParams, + models_dir: PathBuf, +} + +impl ViewerApp { + /// The modality of the series a slot is showing, upper-cased. + fn slot_modality(&self, slot: usize) -> String { + self.slots[slot] + .study + .as_ref() + .and_then(|st| st.series.get(st.active_series)) + .map(|s| s.modality.to_uppercase()) + .unwrap_or_default() + } + + /// Tools ▶ body contour: open the tool window for `slot`. + pub(super) fn open_body_dialog(&mut self, slot: usize) { + if self.slots[slot].study.is_none() { + return; + } + let modality = self.slot_modality(slot); + match &mut self.body_dialog { + // Re-target an open window unless it is busy with the other slot. + Some(d) if self.body_job.is_none() => d.slot = slot, + Some(_) => {} + None => { + self.body_dialog = Some(BodyDialog { + slot, + params: BodyParams::for_modality(&modality), + seeded_for: modality, + status: None, + }); + } + } + } + + /// Snapshot the parameters and the volume and run on a worker thread. + pub(super) fn start_body(&mut self) { + if self.body_job.is_some() { + return; + } + let Some(d) = &self.body_dialog else { + return; + }; + let Some(study) = self.slots[d.slot].study.as_ref() else { + return; + }; + let volume = study.volume.clone(); + let slot = d.slot; + let mut params = d.params.clone(); + params.name = params.name.trim().to_string(); + if params.name.is_empty() { + params.name = "BODY".into(); + } + let req = BodyRequest { + params, + models_dir: self.engine_models_dir(ModelsEngine::TotalSegmentator), + }; + self.persist_settings(); + let progress = Arc::new(Progress::default()); + progress.set("Preparing…"); + self.body_slot = slot; + self.body_job = Some(Job::spawn(progress, move |p| { + ( + slot, + bodymask::contour_body(&volume, &req.params, &req.models_dir, p), + ) + })); + } + + /// A run finished: verify the slot still shows the same volume, land the + /// mask, and — when asked — file it as an RTSTRUCT `EXTERNAL` too. + pub(super) fn on_body_done(&mut self, slot: usize, result: BodyResult) { + if !self.slot_still_shows(slot, result.volume_dims, &result.frame_of_reference_uid) { + self.error = Some(stale_result(&BODY_CONTOUR)); + return; + } + if result.voxels == 0 { + self.error = Some( + "The body contour came out empty — lower the threshold, or reduce the \ + opening radius." + .into(), + ); + return; + } + let idx = self.add_colored_segmentation( + slot, + result.name.clone(), + // Bone-white: the outline is a reference, not one more coloured + // structure competing with the anatomy inside it. + [230, 230, 220], + result.volume_dims, + &result.mask, + ); + if result.make_external { + self.seg_to_rtstruct(slot, idx, "EXTERNAL"); + } + // The cm³ conversions read `self`, so they happen before the + // dialog is borrowed mutably. + let removed_cm3 = self.voxels_to_cm3(slot, result.removed_voxels); + let recovered_cm3 = self.voxels_to_cm3(slot, result.recovered_voxels); + let pieces = match result.pieces.len() { + 0 | 1 => String::new(), + n => format!(", {n} separate bodies"), + }; + let device = if result.device.is_empty() { + String::new() + } else { + format!(" on {}", result.device) + }; + if let Some(d) = &mut self.body_dialog { + d.status = Some(format!( + "✔ {}: {:.0} cm³{pieces} in {:.1} s{device} — {:.0} cm³ of couch, chair, \ + immobilisation and stray objects left out{}", + result.name, + result.cm3, + result.elapsed_secs, + removed_cm3, + if result.recovered_voxels > 0 { + format!(", {recovered_cm3:.0} cm³ of thin anatomy kept") + } else { + String::new() + } + )); + } + } + + fn voxels_to_cm3(&self, slot: usize, voxels: u64) -> f64 { + let sp = self.slots[slot] + .study + .as_ref() + .map(|s| s.volume.spacing) + .unwrap_or([1.0; 3]); + voxels as f64 * sp[0] * sp[1] * sp[2] / 1000.0 + } + + /// The tool window; while a run is in flight its buttons become the + /// progress row. + pub(super) fn body_window(&mut self, ctx: &egui::Context) { + let Some(slot) = self.body_dialog.as_ref().map(|d| d.slot) else { + return; + }; + if self.slots[slot].study.is_none() { + self.body_dialog = None; + return; + } + // Everything that reads the whole of `self` is settled before the + // dialog is borrowed mutably for the frame. + let modality = self.slot_modality(slot); + let idle = self.body_job.is_none(); + let models_dir = models::engine_dir( + &models::root_from_setting(&self.models_dir), + ModelsEngine::TotalSegmentator, + ); + let Some(d) = &mut self.body_dialog else { + return; + }; + // Re-seed on a series switch: a CT threshold cannot be carried over + // to an MR series, and the reverse is just as wrong. + if modality != d.seeded_for && idle { + let name = d.params.name.clone(); + d.params = BodyParams::for_modality(&modality); + d.params.name = name; + d.seeded_for = modality; + } + let running = self.body_job.as_ref().filter(|_| self.body_slot == slot); + let mut open = true; + let (mut run, mut close, mut browse, mut cancel) = (false, false, false, false); + egui::Window::new(BODY_CONTOUR.title(d.slot)) + .id(egui::Id::new("body_window")) + .collapsible(true) + .resizable(false) + .default_width(430.0) + .open(&mut open) + .show(ctx, |ui| { + ui.label( + "Finds the patient's outer surface and leaves the couch, the chair and \ + the immobilisation outside it — the EXTERNAL structure everything \ + downstream starts from.", + ); + ui.separator(); + ui.label("Method:"); + for m in Method::ALL { + let hint = match m { + Method::Classical => { + "Threshold and morphology only. Nothing to download, a few \ + seconds, and the same answer every time. Equipment is found \ + by being thin and repeating slice after slice." + } + Method::ModelAssisted => { + "A body-outline network decides what is patient, the threshold \ + still places the skin. Slower and needs a one-off download, \ + but it is the only thing that removes a mask or a couch \ + touching the skin with no gap." + } + }; + ui.radio_value(&mut d.params.method, m, m.label()) + .on_hover_text(hint); + } + ui.separator(); + + // ---- what counts as tissue ---------------------------- + foreground_row(ui, &mut d.params.foreground); + + if d.params.method == Method::ModelAssisted { + ui.horizontal(|ui| { + ui.label("Model:"); + egui::ComboBox::from_id_salt("body_model") + .selected_text(d.params.model.label()) + .show_ui(ui, |ui| { + for m in BodyModel::ALL { + ui.selectable_value(&mut d.params.model, m, m.label()); + } + }); + let need = bodymask::download_needed(d.params.model, &models_dir); + ui.weak(if need == 0 { + "cached ✓".to_string() + } else { + format!("{} MB to download once", need / 1_000_000) + }); + }); + if matches!(d.params.model, BodyModel::Mr) != d.params.foreground.is_mr() { + ui.label( + egui::RichText::new( + "The chosen model was trained on the other modality.", + ) + .small() + .color(warn_color(ui.visuals())), + ); + } + } + + ui.horizontal(|ui| { + ui.label("Name:"); + ui.add(egui::TextEdit::singleline(&mut d.params.name).desired_width(140.0)); + ui.checkbox(&mut d.params.make_external, "as EXTERNAL structure") + .on_hover_text( + "Also file the result as an RTSTRUCT ROI of type EXTERNAL — \ + what a planning system looks for to find the patient surface. \ + It rides the DICOM export like any other contour.", + ); + }); + + ui.separator(); + ui.collapsing("Options", |ui| { + ui.horizontal(|ui| { + ui.label("Smallest body detail:"); + ui.add( + egui::Slider::new(&mut d.params.open_mm, 0.0..=20.0) + .suffix(" mm") + .fixed_decimals(1), + ) + .on_hover_text( + "The opening that decides what is solid enough to be a body. \ + Anything thicker than twice this keeps its exact surface; \ + anything thinner is left to the thin-anatomy step below.", + ); + }); + ui.checkbox( + &mut d.params.remove_devices, + "Remove equipment that repeats slice after slice", + ) + .on_hover_text( + "A couch, a backrest, a seat pan and an arm rest are surfaces swept \ + along one axis, so their footprint repeats. An ear or a finger \ + never does.", + ); + ui.add_enabled_ui(d.params.remove_devices, |ui| { + ui.horizontal(|ui| { + ui.label(" shells up to:"); + ui.add( + egui::DragValue::new(&mut d.params.device_thin_mm) + .range(0.5..=10.0) + .speed(0.1) + .suffix(" mm"), + ) + .on_hover_text( + "Half the thickness a shell may have to count as equipment. \ + A couch skin is one or two millimetres of carbon, a mask \ + two or three; the thinnest tissue you would miss is a good \ + deal thicker. Raising this is how you start deleting \ + patients.", + ); + }); + ui.horizontal(|ui| { + ui.label(" repeating over:"); + ui.add( + egui::DragValue::new(&mut d.params.persist_window_mm) + .range(20.0..=600.0) + .suffix(" mm"), + ); + ui.label("in at least"); + ui.add( + egui::DragValue::new(&mut d.params.persist_frac) + .range(0.5..=1.0) + .speed(0.01) + .fixed_decimals(2), + ); + ui.label("of its slices"); + }); + }); + ui.horizontal(|ui| { + ui.label("Smallest body part:"); + ui.add( + egui::DragValue::new(&mut d.params.min_volume_cm3) + .range(1.0..=5000.0) + .suffix(" cm³"), + ) + .on_hover_text( + "Kept as a volume, not as 'the largest object', so a leg scan \ + comes out as two bodies rather than one.", + ); + }); + ui.checkbox( + &mut d.params.recover_thin, + "Keep thin anatomy (skin, ears, fingers)", + ) + .on_hover_text( + "What the opening shaved off the body's own surface is always \ + given back. What stands clear of it — an ear, a nose, a \ + fingertip — is given back if it is small enough; a pad, a \ + blanket or a bolus is not.", + ); + ui.add_enabled_ui(d.params.recover_thin, |ui| { + ui.horizontal(|ui| { + ui.label(" standing clear, up to:"); + ui.add( + egui::DragValue::new(&mut d.params.thin_max_extent_mm) + .range(10.0..=400.0) + .suffix(" mm"), + ); + ui.label("across"); + }); + }); + ui.checkbox( + &mut d.params.fill_interior, + "Solid body (fill lungs and gas)", + ) + .on_hover_text( + "Lungs and bowel gas belong inside the body. Unticked gives \ + the tissue shell the threshold sees instead, cavities open.", + ); + ui.horizontal(|ui| { + ui.label("Surface smoothing:"); + ui.add( + egui::Slider::new(&mut d.params.close_mm, 0.0..=10.0) + .suffix(" mm") + .fixed_decimals(1), + ) + .on_hover_text( + "A closing applied last, to take the staircase off the contour.", + ); + }); + if d.params.method == Method::ModelAssisted { + ui.separator(); + ui.horizontal(|ui| { + ui.label("Network margin:"); + ui.add( + egui::Slider::new(&mut d.params.guide_margin_mm, 0.0..=20.0) + .suffix(" mm") + .fixed_decimals(1), + ) + .on_hover_text( + "How far the network's answer is grown before it is used as \ + a mask. It is planned at 6 mm or 1.5 mm, so it needs room \ + not to clip the skin the threshold found.", + ); + }); + device_row(ui, &mut d.params.device); + browse = models_dir_row( + ui, + &mut self.models_dir, + ModelsEngine::TotalSegmentator, + ); + } + }); + ui.separator(); + match d.params.method { + Method::Classical => licence_line( + ui, + "No weights and no network: thresholding and morphology, computed here.", + false, + ), + Method::ModelAssisted => { + let (note, warn) = weights_licence(ModelsEngine::TotalSegmentator); + licence_line(ui, note, warn) + } + } + ui.separator(); + match running { + Some(job) => cancel = progress_row(ui, &job.progress), + None => { + ui.horizontal(|ui| { + if ui + .button("▶ Contour") + .on_hover_text("Find the patient surface in the displayed series") + .clicked() + { + run = true; + } + if ui.button("Close").clicked() { + close = true; + } + }); + } + } + if let Some(status) = &d.status { + ui.separator(); + ui.weak(status); + } + }); + if browse { + if let Some(dir) = Self::pick_folder("Model folder") { + self.models_dir = dir.display().to_string(); + } + } + if cancel { + if let Some(job) = &self.body_job { + job.progress.cancel(); + } + } + if run { + self.start_body(); + } + if !open || close { + // The run, if any, carries on; the sidebar still shows it. + self.body_dialog = None; + self.persist_settings(); + } + } +} + +/// The threshold row — a different question on CT and on MR, so a different +/// row rather than one control that means two things. +fn foreground_row(ui: &mut egui::Ui, fg: &mut Foreground) { + match fg { + Foreground::Hu(t) => { + ui.horizontal(|ui| { + ui.label("Tissue above:"); + ui.add(egui::Slider::new(t, -900.0..=200.0).suffix(" HU")) + .on_hover_text( + "−300 HU sits in the gap between air and fat. The skin edge moves \ + about half a millimetre per 100 HU through the partial-volume ramp, \ + so this is the one number worth agreeing on with your planning \ + system.", + ); + }); + } + _ => { + let mut otsu = matches!(fg, Foreground::MrOtsu { .. }); + let (mut fraction, mut sigma) = match *fg { + Foreground::MrRelative { fraction, sigma_mm } => (fraction, sigma_mm), + Foreground::MrOtsu { sigma_mm } => (0.12, sigma_mm), + Foreground::Hu(_) => (0.12, 40.0), + }; + ui.horizontal(|ui| { + ui.label("Tissue above:"); + ui.radio_value(&mut otsu, false, "a fraction of the signal") + .on_hover_text( + "A low fraction of the 99th percentile. Robust, because the \ + boundary being looked for is the largest step in the image.", + ); + ui.radio_value(&mut otsu, true, "Otsu").on_hover_text( + "No constant to pick, but Otsu splits bright from dark rather than \ + tissue from air, so it runs high on fat-suppressed series.", + ); + }); + ui.horizontal(|ui| { + ui.add_enabled( + !otsu, + egui::Slider::new(&mut fraction, 0.02..=0.6) + .prefix("× p99 ") + .fixed_decimals(2), + ); + ui.label("bias blur:"); + ui.add( + egui::DragValue::new(&mut sigma) + .range(0.0..=200.0) + .suffix(" mm"), + ) + .on_hover_text( + "The receive coils shade the image, so one threshold cannot hold across \ + it. Dividing by the image blurred far beyond any anatomy flattens the \ + shading and leaves every edge intact.", + ); + }); + *fg = if otsu { + Foreground::MrOtsu { sigma_mm: sigma } + } else { + Foreground::MrRelative { + fraction, + sigma_mm: sigma, + } + }; + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::nn::device::DevicePref; + + #[test] + fn the_tool_names_itself_like_the_others() { + assert_eq!(BODY_CONTOUR.title(0), "👤 Body contour — dataset A"); + assert_eq!(BODY_CONTOUR.menu_entry(1), "👤 Body-contour dataset B…"); + assert_eq!(BODY_CONTOUR.short_button(), "👤 Body…"); + } + + #[test] + fn a_device_preference_is_only_meaningful_with_a_network() { + // The classical method never resolves a device; the default stays + // Auto so that switching methods needs no second decision. + let p = BodyParams::default(); + assert_eq!(p.method, Method::Classical); + assert_eq!(p.device, DevicePref::Auto); + } +} diff --git a/src/app/chrome.rs b/src/app/chrome.rs index 3df533f..19db3dc 100644 --- a/src/app/chrome.rs +++ b/src/app/chrome.rs @@ -180,8 +180,14 @@ impl ViewerApp { }); ui.menu_button("Tools", |ui| { // The three segmentation engines, one block per dataset: - // the same three entries, in the same order, for A and B. - let tools: [(&ToolInfo, &str); 3] = [ + // the same four entries, in the same order, for A and B. + let tools: [(&ToolInfo, &str); 4] = [ + ( + &BODY_CONTOUR, + "Outline the patient and leave the couch, the chair and the \ + immobilisation outside — the EXTERNAL structure. Works on CT \ + and MR, with or without a network.", + ), ( &AUTOSEG, "Automatic multi-organ segmentation of the displayed CT \ @@ -222,6 +228,9 @@ impl ViewerApp { } } match open_tool { + Some((slot, t)) if t.glyph == BODY_CONTOUR.glyph => { + self.open_body_dialog(slot) + } Some((slot, t)) if t.glyph == AUTOSEG.glyph => { self.open_autoseg_dialog(slot) } diff --git a/src/app/dialogs.rs b/src/app/dialogs.rs index e5c9465..1f89b2a 100644 --- a/src/app/dialogs.rs +++ b/src/app/dialogs.rs @@ -48,6 +48,7 @@ impl ViewerApp { self.rename_window(ctx); self.autoseg_run_window(ctx); self.segvol_window(ctx); + self.body_window(ctx); self.medsam2_window(ctx); self.autoseg_result_window(ctx); if let Some(msg) = self.notice.clone() { diff --git a/src/app/mod.rs b/src/app/mod.rs index 5c49448..4936997 100644 --- a/src/app/mod.rs +++ b/src/app/mod.rs @@ -14,6 +14,7 @@ use rayon::prelude::*; use crate::anonymize; use crate::autoseg; +use crate::bodymask; use crate::dicom_export; use crate::extras; use crate::gen_test_data::{self, GenParams}; @@ -32,6 +33,7 @@ use crate::settings::{self, Settings}; use crate::simulate::{self, SimParams}; use crate::volume::{ViewPlane, Volume}; +mod body_win; mod box_seg; mod chrome; mod d3; @@ -857,6 +859,13 @@ pub struct ViewerApp { /// Finished result awaiting organ selection. autoseg_pending: Option, + // Body / External contouring (see `bodymask`) — the one tool that can + // answer with no network at all. + body_job: Option>, + body_slot: usize, + /// The tool window, when open; it stays open across runs. + body_dialog: Option, + // Prompt-driven segmentation (SegVol re-implementation, see `segvol`). segvol_job: Option>, segvol_slot: usize, @@ -1005,6 +1014,10 @@ impl ViewerApp { autoseg_slot: 0, autoseg_dialog: None, autoseg_pending: None, + body_job: None, + body_slot: 0, + body_dialog: None, + segvol_job: None, segvol_slot: 0, segvol_dialog: None, @@ -1199,6 +1212,11 @@ impl eframe::App for ViewerApp { { self.on_segvol_done(slot, result); } + if let Some((slot, result)) = + poll_tool_job(&mut self.body_job, &ctx, BODY_CONTOUR.name, &mut self.error) + { + self.on_body_done(slot, result); + } // Poll background registration. if let Some((fixed_slot, out)) = diff --git a/src/app/models_win.rs b/src/app/models_win.rs index d00a577..65846dc 100644 --- a/src/app/models_win.rs +++ b/src/app/models_win.rs @@ -148,7 +148,7 @@ impl ViewerApp { .open(&mut open) .show(ctx, |ui| { ui.label( - "Every model the three segmentation engines can fetch. Weights are \ + "Every model the segmentation tools can fetch. Weights are \ downloaded once, converted to a cache beside them, and never touched \ again — this window is where that inventory is managed.", ); diff --git a/src/app/panels.rs b/src/app/panels.rs index cc1d13b..623461d 100644 --- a/src/app/panels.rs +++ b/src/app/panels.rs @@ -866,6 +866,11 @@ impl ViewerApp { new_series = true; } for (tool, hint) in [ + ( + &BODY_CONTOUR, + "Outline the patient without the couch, the chair or the \ + immobilisation (EXTERNAL)", + ), ( &AUTOSEG, "Automatic multi-organ segmentation (TotalSegmentator, \ @@ -1062,6 +1067,7 @@ impl ViewerApp { self.create_seg(slot); } match open_tool.map(|t| t.glyph) { + Some(g) if g == BODY_CONTOUR.glyph => self.open_body_dialog(slot), Some(g) if g == AUTOSEG.glyph => self.open_autoseg_dialog(slot), Some(g) if g == PROMPT_SEG.glyph => self.open_segvol_dialog(slot), Some(_) => self.open_medsam2_panel(slot), @@ -1092,7 +1098,7 @@ impl ViewerApp { } } if let Some(i) = to_struct { - self.seg_to_rtstruct(slot, i); + self.seg_to_rtstruct(slot, i, "ORGAN"); } if set_act.is_some() { self.set_action = set_act; diff --git a/src/app/seg.rs b/src/app/seg.rs index 50419d0..684f4db 100644 --- a/src/app/seg.rs +++ b/src/app/seg.rs @@ -299,7 +299,12 @@ impl ViewerApp { /// Convert a segmentation into RTSTRUCT contours: appends a ROI to the /// slot's active structure set (creating an in-memory set if the study /// has none), so it displays like any ROI and rides the DICOM export. - pub(super) fn seg_to_rtstruct(&mut self, slot: usize, seg_idx: usize) { + /// + /// `roi_type` is the RT ROI Interpreted Type the ROI is filed under. + /// Everything painted by hand is an `ORGAN`; the body contour is the + /// one thing that has to be an `EXTERNAL`, because that is the tag a + /// planning system looks for to find the patient surface. + pub(super) fn seg_to_rtstruct(&mut self, slot: usize, seg_idx: usize, roi_type: &str) { // The contours are built first, under a shared borrow, because the // segments now live inside the very study the ROI is appended to. let s = &self.slots[slot]; @@ -315,7 +320,8 @@ impl ViewerApp { .get(s.active_structs) .map(|ss| ss.rois.iter().map(|r| r.number).max().unwrap_or(0) + 1) .unwrap_or(1); - let roi = segmentation::mask_to_roi(seg, &vol.grid(), number); + let mut roi = segmentation::mask_to_roi(seg, &vol.grid(), number); + roi.roi_type = roi_type.to_string(); let StudySlot { study, @@ -390,7 +396,7 @@ impl ViewerApp { s.active_seg = first_new; if p.also_rs { for i in first_new..first_new + added { - self.seg_to_rtstruct(slot, i); + self.seg_to_rtstruct(slot, i, "ORGAN"); } } } diff --git a/src/app/seg_engines.rs b/src/app/seg_engines.rs index e360877..ee167fa 100644 --- a/src/app/seg_engines.rs +++ b/src/app/seg_engines.rs @@ -44,6 +44,14 @@ pub(super) const SLICE_PROP: ToolInfo = ToolInfo { name: "Slice propagation", verb: "Propagate through", }; +/// The fourth tool. Its glyph is a person because that is what it outlines, +/// and because it is one of the few figures egui's bundled emoji font +/// actually carries. +pub(super) const BODY_CONTOUR: ToolInfo = ToolInfo { + glyph: "👤", + name: "Body contour", + verb: "Body-contour", +}; impl ToolInfo { /// `🤖 Auto-segmentation — dataset A`, the window title. @@ -163,6 +171,9 @@ impl ViewerApp { { return Some((&SLICE_PROP, &job.progress)); } + if let Some(job) = self.body_job.as_ref().filter(|_| self.body_slot == slot) { + return Some((&BODY_CONTOUR, &job.progress)); + } None } } @@ -290,10 +301,16 @@ mod tests { assert_eq!(AUTOSEG.short_button(), "🤖 Auto…"); assert_eq!(PROMPT_SEG.short_button(), "🧠 Prompt…"); assert_eq!(SLICE_PROP.short_button(), "⏩ Propagate…"); - let mut glyphs = vec![AUTOSEG.glyph, PROMPT_SEG.glyph, SLICE_PROP.glyph]; + assert_eq!(BODY_CONTOUR.menu_entry(0), "👤 Body-contour dataset A…"); + let mut glyphs = vec![ + AUTOSEG.glyph, + PROMPT_SEG.glyph, + SLICE_PROP.glyph, + BODY_CONTOUR.glyph, + ]; glyphs.sort(); glyphs.dedup(); - assert_eq!(glyphs.len(), 3, "every tool has its own glyph"); + assert_eq!(glyphs.len(), 4, "every tool has its own glyph"); } #[test] diff --git a/src/autoseg/config.rs b/src/autoseg/config.rs index bf0a174..28273aa 100644 --- a/src/autoseg/config.rs +++ b/src/autoseg/config.rs @@ -9,8 +9,25 @@ use anyhow::{bail, Context, Result}; use serde_json::Value; +/// How the model wants its input scaled — nnU-Net's `normalization_schemes`. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Norm { + /// `CTNormalization`: clip to the training set's [p0.5, p99.5] window, + /// then z-score with the dataset fingerprint's mean and standard + /// deviation. Every constant comes from `plans.json`, so the same CT + /// always normalizes the same way. + Ct, + /// `ZScoreNormalization`: subtract *this image's* mean and divide by + /// *this image's* standard deviation. MR has no absolute scale, so the + /// MR models use it — and it is why an MR run fills its normalization + /// constants in only after resampling ([`ModelConfig::apply_image_norm`]). + ZScore, +} + #[derive(Clone, Debug)] pub struct ModelConfig { + /// Which scheme [`Self::clip_lo`] … [`Self::std`] are to be read under. + pub norm: Norm, /// Sliding-window patch size per spatial axis. pub patch_size: [usize; 3], /// Target voxel spacing (mm) per spatial axis (isotropic for these models). @@ -25,7 +42,8 @@ pub struct ModelConfig { pub n_conv_per_stage: Vec, /// Convs per decoder stage. pub n_conv_per_stage_decoder: Vec, - /// CT normalization: clip bounds (HU) then z-score. + /// Clip bounds then z-score. For [`Norm::ZScore`] the bounds are + /// infinite and the mean/std belong to the image, not the dataset. pub clip_lo: f32, pub clip_hi: f32, pub mean: f32, @@ -37,6 +55,30 @@ impl ModelConfig { self.features.len() } + /// Fill the normalization constants from the resampled image itself, as + /// `ZScoreNormalization` requires. A no-op for CT models, whose + /// constants are fixed by the training set. + pub fn apply_image_norm(&mut self, voxels: &[f32]) { + if self.norm != Norm::ZScore || voxels.is_empty() { + return; + } + let n = voxels.len() as f64; + let mean = voxels.iter().map(|&v| v as f64).sum::() / n; + let var = voxels + .iter() + .map(|&v| { + let d = v as f64 - mean; + d * d + }) + .sum::() + / n; + // nnU-Net: `image -= mean; image /= (std + 1e-8)`. + self.clip_lo = f32::NEG_INFINITY; + self.clip_hi = f32::INFINITY; + self.mean = mean as f32; + self.std = var.sqrt() as f32 + 1e-8; + } + pub fn from_plans_json(text: &str) -> Result { let root: Value = serde_json::from_str(text).context("parse plans.json")?; let tf = root @@ -60,9 +102,11 @@ impl ModelConfig { .pointer("/normalization_schemes/0") .and_then(|v| v.as_str()) .unwrap_or(""); - if norm != "CTNormalization" { - bail!("plans.json: unsupported normalization scheme {norm:?}"); - } + let norm = match norm { + "CTNormalization" => Norm::Ct, + "ZScoreNormalization" => Norm::ZScore, + other => bail!("plans.json: unsupported normalization scheme {other:?}"), + }; let usize3 = |v: &Value, what: &str| -> Result<[usize; 3]> { let a: Vec = v .as_array() @@ -129,16 +173,28 @@ impl ModelConfig { if n_conv_per_stage.len() != n_stages || n_conv_per_stage_decoder.len() != n_stages - 1 { bail!("plans.json: conv-per-stage lengths do not match stage count"); } - let fg = root - .pointer("/foreground_intensity_properties_per_channel/0") - .context("plans.json: intensity properties missing")?; + // A z-score model never reads these: its constants come from the + // image in `apply_image_norm`. The values left here are the + // identity, so a model that somehow skips that step is merely + // un-normalized rather than scaled by nonsense. + let fg = root.pointer("/foreground_intensity_properties_per_channel/0"); let f = |key: &str| -> Result { - fg.get(key) + if norm == Norm::ZScore { + return Ok(match key { + "percentile_00_5" => f32::NEG_INFINITY, + "percentile_99_5" => f32::INFINITY, + "std" => 1.0, + _ => 0.0, + }); + } + fg.context("plans.json: intensity properties missing")? + .get(key) .and_then(|v| v.as_f64()) .map(|v| v as f32) .with_context(|| format!("plans.json: intensity {key}")) }; Ok(ModelConfig { + norm, patch_size, spacing, features, diff --git a/src/autoseg/cpu.rs b/src/autoseg/cpu.rs index 679d50f..ef0437b 100644 --- a/src/autoseg/cpu.rs +++ b/src/autoseg/cpu.rs @@ -15,7 +15,7 @@ use rayon::prelude::*; // SegVol mask decoder, so they live in `nn`; re-exported here because this // is where the nnU-Net code has always reached for them. use crate::nn::tensor::SendPtr; -pub use crate::nn::tensor::{conv_transpose3d_2x, Act}; +pub use crate::nn::tensor::{conv_transpose3d_2x, conv_transpose3d_stride, Act}; #[inline] fn conv_out(len: usize, k: usize, s: usize) -> usize { diff --git a/src/autoseg/gpu.rs b/src/autoseg/gpu.rs index 55c7e6c..6bc3f48 100644 --- a/src/autoseg/gpu.rs +++ b/src/autoseg/gpu.rs @@ -33,6 +33,7 @@ struct GBlock { struct GTransp { w: Tensor, b: Tensor, + stride: [usize; 3], } /// The network with weights resident on the GPU. @@ -91,8 +92,13 @@ impl GpuNet { .transp .iter() .map(|t| GTransp { - w: upload5(d, &t.w, [t.cin, t.cout, 2, 2, 2]), + w: upload5( + d, + &t.w, + [t.cin, t.cout, t.stride[0], t.stride[1], t.stride[2]], + ), b: upload1(d, &t.b), + stride: t.stride, }) .collect(); let head_w = upload5(d, &unet.head.w, [unet.head.classes, unet.head.cin, 1, 1, 1]); @@ -130,7 +136,7 @@ impl GpuNet { cur, tc.w.clone(), Some(tc.b.clone()), - ConvTransposeOptions::new([2, 2, 2], [0, 0, 0], [0, 0, 0], [1, 1, 1], 1), + ConvTransposeOptions::new(tc.stride, [0, 0, 0], [0, 0, 0], [1, 1, 1], 1), ); let skip = skips.pop().unwrap(); cur = Tensor::cat(vec![cur, skip], 1); diff --git a/src/autoseg/infer.rs b/src/autoseg/infer.rs index fdb5e67..9d1cc32 100644 --- a/src/autoseg/infer.rs +++ b/src/autoseg/infer.rs @@ -274,6 +274,7 @@ mod tests { } } let cfg = ModelConfig { + norm: crate::autoseg::config::Norm::Ct, patch_size: [8, 8, 8], spacing: [3.0, 3.0, 3.0], features: vec![], diff --git a/src/autoseg/mod.rs b/src/autoseg/mod.rs index a168e53..af9fd4a 100644 --- a/src/autoseg/mod.rs +++ b/src/autoseg/mod.rs @@ -121,6 +121,7 @@ struct Hooks<'a> { progress: &'a Progress, /// (model index, model count) for progress text. model: (usize, usize), + /// What the run is called in the progress line. label: &'static str, } @@ -145,21 +146,25 @@ impl infer::InferHooks for Hooks<'_> { } } -/// Run auto-segmentation on a CT volume. Blocking — call from a worker -/// thread; observe/cancel through `progress`. +/// Run one or more nnU-Net models over a volume and return their merged +/// labels **on the volume's own grid**, plus the description of the device +/// the work ran on. /// -/// `parts` selects sub-models for [`Variant::HighRes15mm`] -/// (organs, vertebrae, cardiac, muscles, ribs) and is ignored otherwise. -pub fn run( +/// This is the whole engine minus the question being asked: the 117-class +/// "total" task ([`run`]) and the two-class body-outline task +/// ([`crate::bodymask`]) differ in which checkpoints they load and what +/// they do with the answer, not in how a checkpoint is fetched, converted, +/// resampled onto, tiled over or mapped back from. +/// +/// `label` names the run in progress messages. +pub fn run_specs( volume: &Volume, - variant: Variant, + specs: &[ModelSpec], + label: &'static str, device: DevicePref, - parts: [bool; 5], models_dir: &Path, progress: &Progress, -) -> Result { - let t_start = std::time::Instant::now(); - let specs = variant.specs(parts); +) -> Result<(Vec, String)> { if specs.is_empty() { bail!("no sub-models selected"); } @@ -186,7 +191,6 @@ pub fn run( bail!("sub-models disagree on spacing"); } } - let target = spacing[0]; // ---- engine ---------------------------------------------------------- progress.set("Choosing the compute device…"); @@ -199,8 +203,18 @@ pub fn run( // ---- preprocess ------------------------------------------------------ progress.set_phase(0.15, 0.05); - progress.report(0.0, &format!("Resampling volume to {target} mm…")); - let map = preprocess::SarMap::new(volume, target); + progress.report( + 0.0, + &format!( + "Resampling volume to {} mm…", + if spacing[0] == spacing[1] && spacing[1] == spacing[2] { + format!("{}", spacing[0]) + } else { + format!("{} × {} × {}", spacing[0], spacing[1], spacing[2]) + } + ), + ); + let map = preprocess::SarMap::new(volume, spacing); let vol_model = preprocess::resample_to_model(volume, &map); if progress.cancelled() { bail!(CANCELLED); @@ -211,7 +225,11 @@ pub fn run( let infer_span = 0.75 / n_models as f32; for (mi, model) in models.iter().enumerate() { progress.set_phase(0.2 + mi as f32 * infer_span, infer_span); - let unet = net::UNet::build(model.config.clone(), &model.tensors) + // A z-score model normalizes against this image, so its constants + // are only knowable now, with the resampled volume in hand. + let mut cfg = model.config.clone(); + cfg.apply_image_norm(&vol_model); + let unet = net::UNet::build(cfg, &model.tensors) .with_context(|| format!("assemble network ({})", model.spec.label))?; let classes = unet.num_classes(); let forward: ForwardFn = match &gpu { @@ -242,7 +260,7 @@ pub fn run( forward, progress, model: (mi, n_models), - label: variant.label(), + label, }; let local = infer::predict( &vol_model, @@ -262,10 +280,38 @@ pub fn run( } } - // ---- back-map to the CT grid + statistics ---------------------------- + // ---- back-map to the CT grid ---------------------------------------- progress.set_phase(0.95, 0.05); progress.report(0.0, "Mapping labels back to the CT grid…"); let labels = preprocess::labels_to_volume_grid(&global, &map, volume); + Ok((labels, device_desc)) +} + +/// Run auto-segmentation on a CT volume. Blocking — call from a worker +/// thread; observe/cancel through `progress`. +/// +/// `parts` selects sub-models for [`Variant::HighRes15mm`] +/// (organs, vertebrae, cardiac, muscles, ribs) and is ignored otherwise. +pub fn run( + volume: &Volume, + variant: Variant, + device: DevicePref, + parts: [bool; 5], + models_dir: &Path, + progress: &Progress, +) -> Result { + let t_start = std::time::Instant::now(); + let specs = variant.specs(parts); + let (labels, device_desc) = run_specs( + volume, + &specs, + variant.label(), + device, + models_dir, + progress, + )?; + + // ---- statistics ------------------------------------------------------ let mut counts = [0u64; 256]; for l in &labels { counts[*l as usize] += 1; diff --git a/src/autoseg/net.rs b/src/autoseg/net.rs index d0b7b38..6e16b16 100644 --- a/src/autoseg/net.rs +++ b/src/autoseg/net.rs @@ -17,7 +17,7 @@ use anyhow::{bail, Context, Result}; use std::collections::HashMap; use super::config::ModelConfig; -use super::cpu::{concat, conv3d, conv_transpose3d_2x, instance_norm_lrelu, Act}; +use super::cpu::{concat, conv3d, conv_transpose3d_stride, instance_norm_lrelu, Act}; use crate::nn::cache::WTensor; /// One Conv3d → InstanceNorm → LeakyReLU block. @@ -37,6 +37,10 @@ pub struct TranspConv { pub b: Vec, pub cin: usize, pub cout: usize, + /// Kernel = stride of this upsampling step — the encoder stride it + /// undoes. `[2, 2, 2]` for every isotropic model; the MR models plan + /// `[1, 2, 2]` where the through-plane spacing is already coarse. + pub stride: [usize; 3], } pub struct SegHead { @@ -89,13 +93,16 @@ impl UNet { for t in 0..n - 1 { let c_below = cfg.features[n - 1 - t]; let c_skip = cfg.features[n - 2 - t]; + // Decoder step `t` undoes the stride of encoder stage n-1-t. + let stride = cfg.strides[n - 1 - t]; let tw = take(tensors, &format!("decoder.transpconvs.{t}.weight"))?; let tb = take(tensors, &format!("decoder.transpconvs.{t}.bias"))?; - if tw.shape != [c_below, c_skip, 2, 2, 2] { + let want = [c_below, c_skip, stride[0], stride[1], stride[2]]; + if tw.shape != want { bail!( "decoder.transpconvs.{t}.weight has shape {:?}, expected {:?}", tw.shape, - [c_below, c_skip, 2, 2, 2] + want ); } transp.push(TranspConv { @@ -103,6 +110,7 @@ impl UNet { b: tb.data.clone(), cin: c_below, cout: c_skip, + stride, }); let mut blocks = Vec::new(); let stage_kernel = cfg.kernels[n - 2 - t]; @@ -170,7 +178,7 @@ impl UNet { } let mut cur = skips.pop().unwrap(); for (t, tc) in self.transp.iter().enumerate() { - let up = conv_transpose3d_2x(&cur, &tc.w, &tc.b, tc.cout); + let up = conv_transpose3d_stride(&cur, &tc.w, &tc.b, tc.cout, tc.stride); let skip = skips.pop().unwrap(); cur = concat(&up, &skip); for blk in &self.dec[t] { diff --git a/src/autoseg/preprocess.rs b/src/autoseg/preprocess.rs index d4d55cd..c4940cf 100644 --- a/src/autoseg/preprocess.rs +++ b/src/autoseg/preprocess.rs @@ -35,20 +35,18 @@ pub struct SarMap { } impl SarMap { - pub fn new(vol: &Volume, target_spacing: f64) -> SarMap { + /// `target` is the model's voxel spacing in [S, A, R] order. It is an + /// array rather than a scalar because not every nnU-Net model is + /// isotropic — the MR body model plans 3.0 × 1.19 × 0.99 mm. + pub fn new(vol: &Volume, target: [f64; 3]) -> SarMap { let (perm, flip) = vol.canonical_axes(); let dims = [vol.dims[0], vol.dims[1], vol.dims[2]]; let spac = [vol.spacing[0], vol.spacing[1], vol.spacing[2]]; let orig_dims = [dims[perm[0]], dims[perm[1]], dims[perm[2]]]; let orig_spacing = [spac[perm[0]], spac[perm[1]], spac[perm[2]]]; - let model_dims = [ - ((orig_dims[0] as f64 * orig_spacing[0] / target_spacing).round_ties_even() as usize) - .max(1), - ((orig_dims[1] as f64 * orig_spacing[1] / target_spacing).round_ties_even() as usize) - .max(1), - ((orig_dims[2] as f64 * orig_spacing[2] / target_spacing).round_ties_even() as usize) - .max(1), - ]; + let model_dims = std::array::from_fn(|a| { + ((orig_dims[a] as f64 * orig_spacing[a] / target[a]).round_ties_even() as usize).max(1) + }); SarMap { perm, flip, @@ -231,7 +229,7 @@ mod tests { fn standard_axial_mapping() { // Standard axial LPS volume: i→+x(L), j→+y(P), k→+z(S). let vol = axial_volume([512, 512, 133], [0.9766, 0.9766, 3.0]); - let map = SarMap::new(&vol, 3.0); + let map = SarMap::new(&vol, [3.0; 3]); // S axis = volume z (no flip), A axis = volume y (flip: +y is P), // R axis = volume x (flip: +x is L). assert_eq!(map.perm, [2, 1, 0]); @@ -251,7 +249,7 @@ mod tests { } } } - let map = SarMap::new(&vol, 2.0); // same spacing → pure reorientation + let map = SarMap::new(&vol, [2.0; 3]); // same spacing → pure reorientation assert_eq!(map.model_dims, [10, 20, 20]); let res = resample_to_model(&vol, &map); // voxel count preserved under pure flips/permutation diff --git a/src/autoseg/weights.rs b/src/autoseg/weights.rs index fdfaeca..f9263b4 100644 --- a/src/autoseg/weights.rs +++ b/src/autoseg/weights.rs @@ -52,6 +52,38 @@ pub const SPEC_6MM: ModelSpec = ModelSpec { global_offset: 0, }; +/// The body-outline task's models. Same nnU-Net architecture, same open +/// Apache-2.0 licence, a different question: two classes, trunk and +/// extremities, whose union is the patient. They are what the body-contour +/// tool's model-assisted method uses to tell patient from equipment — see +/// [`crate::bodymask`]. +pub const SPEC_BODY_15MM: ModelSpec = ModelSpec { + key: "body_1_5mm", + label: "body 1.5 mm", + url: "https://github.com/wasserth/TotalSegmentator/releases/download/v2.0.0-weights/Dataset299_body_1559subj.zip", + zip_bytes: 233_211_222, + global_offset: 0, +}; + +pub const SPEC_BODY_6MM: ModelSpec = ModelSpec { + key: "body_6mm", + label: "body 6 mm", + url: "https://github.com/wasserth/TotalSegmentator/releases/download/v2.0.0-weights/Dataset300_body_6mm_1559subj.zip", + zip_bytes: 124_286_256, + global_offset: 0, +}; + +/// The MR counterpart (TotalSegmentator v2.5 weights). Plans a z-score +/// normalization and an anisotropic 3.0 × 1.19 × 0.99 mm grid, which is why +/// the engine carries both. +pub const SPEC_BODY_MR: ModelSpec = ModelSpec { + key: "body_mr", + label: "body MR", + url: "https://github.com/wasserth/TotalSegmentator/releases/download/v2.5.0-weights/Dataset597_mri_body_139subj.zip", + zip_bytes: 229_810_326, + global_offset: 0, +}; + pub const SPECS_15MM: [ModelSpec; 5] = [ ModelSpec { key: "total_part1_organs", @@ -99,6 +131,7 @@ pub fn all_specs() -> Vec { let mut v = vec![SPEC_3MM]; v.extend(SPECS_15MM); v.push(SPEC_6MM); + v.extend([SPEC_BODY_15MM, SPEC_BODY_6MM, SPEC_BODY_MR]); v } diff --git a/src/bodymask.rs b/src/bodymask.rs new file mode 100644 index 0000000..7c23d34 --- /dev/null +++ b/src/bodymask.rs @@ -0,0 +1,856 @@ +//! Automatic **BODY / External** contouring — the outer patient surface, +//! with the couch, the chair and the immobilisation left outside it. +//! +//! Every downstream calculation starts here. A dose engine needs to know +//! where the patient begins, because that is where the range budget starts; +//! a DRR needs to know what is not patient, or it projects the couch through +//! the anatomy; a registration wants to sample inside the body and nowhere +//! else. So this is the one contour that has to be right, on the first scan +//! of the day, without anybody drawing it. +//! +//! Two methods, sharing everything after the first step: +//! +//! * [`Method::Classical`] — thresholding and morphology, deterministic, +//! instantaneous, nothing to download. Equipment is separated from anatomy +//! by two geometric facts and no semantics: a device shell is *thin* +//! (a couch top is two carbon skins around foam that is already below the +//! threshold; a thermoplastic mask is 2–3 mm), and it is *extruded* — +//! the same footprint repeats slice after slice, which no part of a +//! patient does. See [`morphology::axis_persistence`]. +//! +//! * [`Method::ModelAssisted`] — TotalSegmentator's openly licensed +//! body-outline nnU-Net (Apache-2.0, the same engine as +//! [`crate::autoseg`]) decides *what* is patient; the threshold still +//! decides *where* the skin is. The network's own output is far too +//! coarse to be a skin surface — it is planned at 6 mm or 1.5 mm — so it +//! is used dilated, as a mask on the thresholded image, never as the +//! answer. This is what removes a device in gap-free contact with the +//! skin, which no amount of geometry can. +//! +//! Both end with the same post-processing: keep the components big enough to +//! be a patient, give back the thin anatomy the opening took (ears, nose, +//! fingers), fill the interior slice by slice, and optionally close the +//! staircase off the surface. +//! +//! Known limitation, stated rather than hidden: where a shell touches the +//! skin with no air gap at all, the classical method keeps the shell's +//! thickness over the contact patch — 2–5 mm, over the patches only. It is +//! not detectable by geometry, because locally it *is* a slightly thicker +//! patient. The model-assisted method is the answer to that. + +use anyhow::{bail, Result}; +use rayon::prelude::*; +use std::path::Path; + +use crate::autoseg::weights::{ModelSpec, SPEC_BODY_15MM, SPEC_BODY_6MM, SPEC_BODY_MR}; +use crate::morphology as morph; +use crate::nn::device::DevicePref; +use crate::progress::{Progress, ProgressSink, CANCELLED}; +use crate::volume::Volume; + +/// How the patient is told apart from everything else in the image. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Method { + /// Threshold + morphology. No download, no network, a few seconds. + Classical, + /// A body-outline network decides what is patient; the threshold still + /// places the skin. + ModelAssisted, +} + +impl Method { + pub const ALL: [Method; 2] = [Method::Classical, Method::ModelAssisted]; + pub fn label(&self) -> &'static str { + match self { + Method::Classical => "Classical (threshold + morphology)", + Method::ModelAssisted => "Model-assisted (TotalSegmentator body)", + } + } +} + +/// Which body-outline network the model-assisted method runs. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum BodyModel { + /// Dataset300, 6 mm — 124 MB, seconds even on a CPU. Plenty, because + /// the network only has to say *which side of the skin* a voxel is on. + Ct6mm, + /// Dataset299, 1.5 mm — 233 MB, minutes on a CPU. + Ct15mm, + /// Dataset597, the MR body model — 230 MB. + Mr, +} + +impl BodyModel { + pub const ALL: [BodyModel; 3] = [BodyModel::Ct6mm, BodyModel::Ct15mm, BodyModel::Mr]; + pub fn label(&self) -> &'static str { + match self { + BodyModel::Ct6mm => "CT 6 mm (fast)", + BodyModel::Ct15mm => "CT 1.5 mm", + BodyModel::Mr => "MR", + } + } + pub fn spec(&self) -> ModelSpec { + match self { + BodyModel::Ct6mm => SPEC_BODY_6MM, + BodyModel::Ct15mm => SPEC_BODY_15MM, + BodyModel::Mr => SPEC_BODY_MR, + } + } + /// The model to reach for on a series of this modality. + pub fn for_modality(modality: &str) -> BodyModel { + if modality.eq_ignore_ascii_case("MR") { + BodyModel::Mr + } else { + BodyModel::Ct6mm + } + } +} + +/// How the foreground — "anything at all, patient or not" — is found. +/// +/// CT has an absolute scale, so a fixed HU threshold is exactly right. MR +/// has none: the same tissue is a different number on the next sequence, +/// and the receive coils make it a different number on the other side of +/// the same slice. So the MR path divides out a smooth estimate of the coil +/// sensitivity first, then thresholds relative to what is left. +#[derive(Clone, Copy, Debug, PartialEq)] +pub enum Foreground { + /// Hounsfield units. −300 sits in the fat/air gap; the skin edge moves + /// about half a millimetre per 100 HU through the partial-volume ramp. + Hu(f32), + /// Fraction of the bias-corrected 99th percentile, with the bias field + /// estimated by a `sigma_mm` blur. + MrRelative { fraction: f32, sigma_mm: f64 }, + /// Otsu's threshold on the bias-corrected image — no constant to pick, + /// but it splits bright from dark rather than tissue from air, so it + /// runs high on fat-suppressed series. + MrOtsu { sigma_mm: f64 }, +} + +impl Foreground { + /// The sensible default for a series of this modality. + pub fn for_modality(modality: &str) -> Foreground { + if modality.eq_ignore_ascii_case("MR") { + Foreground::MrRelative { + fraction: 0.12, + sigma_mm: 40.0, + } + } else { + Foreground::Hu(-300.0) + } + } + pub fn is_mr(&self) -> bool { + !matches!(self, Foreground::Hu(_)) + } +} + +/// Everything one run of the body contour is told. +#[derive(Clone, Debug, PartialEq)] +pub struct BodyParams { + pub method: Method, + pub model: BodyModel, + pub device: DevicePref, + pub foreground: Foreground, + /// Radius of the opening that decides what is big enough to *be* a + /// body. Everything thinner than twice this is set aside for the + /// thin-anatomy step to judge. + pub open_mm: f64, + /// A shell whose largest inscribed ball is smaller than this is a + /// candidate for equipment — so shells up to about twice it. + /// + /// Deliberately far smaller than [`Self::open_mm`], and not a knob to + /// turn up. A couch skin is one or two millimetres of carbon and a + /// thermoplastic mask two or three; the thinnest tissue anyone would + /// miss — the chest wall over a lung — is five or six. At 2 mm the two + /// are cleanly separated. At 3 mm a six-millimetre chest wall is + /// itself a candidate, and since it repeats slice after slice it is + /// then indistinguishable from a couch: the whole ribcage goes. + pub device_thin_mm: f64, + /// Run the extruded-equipment test. Off by default in the + /// model-assisted method, where the network has already answered. + pub remove_devices: bool, + /// How far a device footprint has to repeat to count as extruded. + pub persist_window_mm: f64, + /// And in what fraction of that window's slices. + pub persist_frac: f64, + /// Components smaller than this are noise, a cable or a pillow — not a + /// patient. Kept as a *volume* so two legs both survive. + pub min_volume_cm3: f64, + /// Give back the thin pieces the opening removed. + pub recover_thin: bool, + /// How big a piece standing clear of the body's own surface may be and + /// still count as anatomy — an ear, a nose, a fingertip. A pad, a + /// blanket or a bolus is larger, and stays out. + pub thin_max_extent_mm: f64, + /// How far the network's coarse answer is grown before it is used as a + /// mask, in the model-assisted method. + pub guide_margin_mm: f64, + /// Report the body as the solid object it is, rather than as the shell + /// of tissue the threshold sees. Off gives the tissue mask instead, + /// with the lungs and the bowel gas left open. + pub fill_interior: bool, + /// Closing radius applied last, to take the staircase off the surface. + pub close_mm: f64, + pub name: String, + /// Also append an RTSTRUCT ROI of type EXTERNAL. + pub make_external: bool, +} + +impl Default for BodyParams { + fn default() -> Self { + BodyParams { + method: Method::Classical, + model: BodyModel::Ct6mm, + device: DevicePref::Auto, + foreground: Foreground::Hu(-300.0), + open_mm: 8.0, + device_thin_mm: 2.0, + remove_devices: true, + persist_window_mm: 150.0, + persist_frac: 0.8, + min_volume_cm3: 50.0, + recover_thin: true, + thin_max_extent_mm: 100.0, + guide_margin_mm: 6.0, + fill_interior: true, + close_mm: 0.0, + name: "BODY".to_string(), + make_external: true, + } + } +} + +impl BodyParams { + /// The defaults that suit a series of this modality — the CT thresholds + /// are meaningless on MR and vice versa, so the tool window re-seeds + /// itself whenever the displayed series changes. + pub fn for_modality(modality: &str) -> BodyParams { + BodyParams { + foreground: Foreground::for_modality(modality), + model: BodyModel::for_modality(modality), + ..BodyParams::default() + } + } +} + +/// One piece of the finished contour, for the results line. +#[derive(Clone, Debug)] +pub struct Piece { + pub voxels: u64, + pub cm3: f64, +} + +/// What a finished run hands back. +pub struct BodyResult { + /// 0/1 per voxel, in [`Volume::data`] index order. Omitted from the + /// `Debug` output, which is otherwise 35 MB of ones and zeros. + pub mask: Vec, + pub dims: [usize; 3], + pub voxels: u64, + pub cm3: f64, + /// The separate bodies kept — two legs are two pieces, and saying so is + /// more use than silently keeping the larger one. + pub pieces: Vec, + /// Voxels above the threshold that were judged not to be patient — the + /// equipment the extrusion test caught, plus every component too small + /// or too detached to be a body. + pub removed_voxels: u64, + /// Thin voxels handed back to the body after the opening. + pub recovered_voxels: u64, + pub method: Method, + /// Which device the network ran on; empty for the classical method. + pub device: String, + pub elapsed_secs: f64, + pub name: String, + pub make_external: bool, + /// Identity of the volume this was computed on. + pub frame_of_reference_uid: String, + pub volume_dims: [usize; 3], +} + +/// The mask is 35 MB of ones and zeros on a normal CT, so it is named +/// rather than printed — everything else is what one wants to see when a +/// test or a batch run reports a surprise. +impl std::fmt::Debug for BodyResult { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("BodyResult") + .field("mask", &format_args!("<{} voxels>", self.mask.len())) + .field("dims", &self.dims) + .field("voxels", &self.voxels) + .field("cm3", &self.cm3) + .field("pieces", &self.pieces) + .field("removed_voxels", &self.removed_voxels) + .field("recovered_voxels", &self.recovered_voxels) + .field("method", &self.method) + .field("device", &self.device) + .field("elapsed_secs", &self.elapsed_secs) + .field("name", &self.name) + .finish_non_exhaustive() + } +} + +/// Total download still needed for the model-assisted method, in bytes. +pub fn download_needed(model: BodyModel, models_dir: &Path) -> u64 { + let spec = model.spec(); + if crate::autoseg::weights::is_cached(&spec, models_dir) { + 0 + } else { + spec.zip_bytes + } +} + +// --------------------------------------------------------------------------- +// The pipeline +// --------------------------------------------------------------------------- + +/// Contour the patient. Blocking — call from a worker thread and watch it +/// through `progress`. +pub fn contour_body( + volume: &Volume, + params: &BodyParams, + models_dir: &Path, + progress: &Progress, +) -> Result { + let t0 = std::time::Instant::now(); + let dims = volume.dims; + let spacing = volume.spacing; + let voxel_cm3 = spacing[0] * spacing[1] * spacing[2] / 1000.0; + + // ---- 1. foreground --------------------------------------------------- + // Six equal steps follow, over whatever is left of the bar: all of it + // for the classical method, the last 30 % when a network has had the + // first 70 %. + const STEPS: f32 = 6.0; + let (fg_base, fg_span) = match params.method { + Method::Classical => (0.0, 1.0 / STEPS), + Method::ModelAssisted => (0.70, 0.30 / STEPS), + }; + let mut device = String::new(); + let mut guide: Option> = None; + if params.method == Method::ModelAssisted { + progress.set_phase(0.0, 0.70); + let spec = params.model.spec(); + let (labels, dev) = crate::autoseg::run_specs( + volume, + &[spec], + "body outline", + params.device, + models_dir, + progress, + )?; + device = dev; + guide = Some(body_from_labels(&labels, dims, spacing, progress)?); + } + progress.set_phase(fg_base, fg_span); + progress.report(0.0, "Separating tissue from air…"); + let mut fg = foreground(volume, params.foreground); + if progress.cancelled() { + bail!(CANCELLED); + } + + // The network's answer is coarse by construction — it is planned at + // 6 mm or 1.5 mm — so it is grown by a margin and used as a *mask* on + // the thresholded image. What survives has the network's semantics and + // the image's resolution. + if let Some(g) = &guide { + let grown = morph::dilate_mm(g, dims, spacing, params.guide_margin_mm); + fg.par_iter_mut().zip(grown.par_iter()).for_each(|(m, &g)| { + if g == 0 { + *m = 0; + } + }); + } + let mut mask = fg.clone(); + + // ---- 2. extruded equipment ------------------------------------------ + // Judged on the *unfilled* foreground, while a shell is still a shell, + // and at a radius small enough that only equipment qualifies. + progress.set_phase(fg_base + fg_span, fg_span); + if params.remove_devices { + progress.report(0.0, "Looking for couch, chair and immobilisation…"); + let opened = morph::open_mm(&mask, dims, spacing, params.device_thin_mm); + let mut thin: Vec = mask + .par_iter() + .zip(opened.par_iter()) + .map(|(&m, &o)| u8::from(m != 0 && o == 0)) + .collect(); + let mut device_mask = vec![0u8; mask.len()]; + for axis in 0..3 { + if progress.cancelled() { + bail!(CANCELLED); + } + let p = morph::axis_persistence( + &thin, + dims, + spacing, + axis, + params.persist_window_mm, + params.persist_frac, + ); + device_mask + .par_iter_mut() + .zip(p.par_iter()) + .for_each(|(d, &v)| *d |= v); + progress.report((axis + 1) as f32 / 3.0, ""); + } + // A device is thin *and* extruded; whatever is only one of the two + // stays for the component test to judge. + thin.par_iter_mut() + .zip(device_mask.par_iter()) + .for_each(|(t, &d)| *t &= d); + mask.par_iter_mut() + .zip(thin.par_iter()) + .for_each(|(m, &d)| { + if d != 0 { + *m = 0; + } + }); + } + if progress.cancelled() { + bail!(CANCELLED); + } + + // ---- 3. a body is a solid object ------------------------------------- + // Everything below reasons about the *body*, and a threshold does not + // see one: it sees a shell of tissue wrapped round two lungs. Left that + // way, the chest wall over a lung is a five-millimetre sheet that + // repeats slice after slice — which is to say, indistinguishable from a + // couch skin, and duly deleted. Filling the interior first is what + // makes the wall part of a solid object again. It happens after the + // equipment step, so that a couch top with a closed profile is not + // turned into a solid slab before it can be recognised. + // + // Axial means "perpendicular to the patient's superior axis", which is + // the volume axis the direction cosines put closest to it, not whichever + // axis happens to be third in the array. + let axial = volume.canonical_axes().0[0]; + progress.set_phase(fg_base + 2.0 * fg_span, fg_span); + progress.report(0.0, "Closing the interior…"); + let mut solid = mask.clone(); + morph::fill_holes_2d(&mut solid, dims, axial); + if progress.cancelled() { + bail!(CANCELLED); + } + + // ---- 4. which components are a patient ------------------------------- + progress.set_phase(fg_base + 3.0 * fg_span, fg_span); + progress.report(0.0, "Finding the body…"); + let core = morph::open_mm(&solid, dims, spacing, params.open_mm); + let comps = morph::components(&core, dims); + if comps.is_empty() { + bail!( + "nothing above the threshold looks like a body — lower the threshold \ + or reduce the opening radius" + ); + } + let min_voxels = (params.min_volume_cm3 / voxel_cm3).max(1.0) as usize; + // Everything big enough, not merely the largest: a leg scan is two + // bodies, and an arm cut off by the field of view is a third. + let mut kept: Vec<&morph::Component> = comps.iter().filter(|c| c.len() >= min_voxels).collect(); + if kept.is_empty() { + kept.push(&comps[0]); + } + let mut body = vec![0u8; mask.len()]; + for c in &kept { + for &v in &c.voxels { + body[v as usize] = 1; + } + } + if progress.cancelled() { + bail!(CANCELLED); + } + + // ---- 5. give back the thin anatomy ----------------------------------- + progress.set_phase(fg_base + 4.0 * fg_span, fg_span); + let mut recovered = 0u64; + if params.recover_thin { + progress.report(0.0, "Recovering ears, nose and fingers…"); + // Two questions, because there are two kinds of thing here. + // + // What the opening shaved off the body's *own* surface lies, by + // construction, within one opening radius of what is left — a skin + // rim, the edge of a shoulder, the sharp flank of a cross-section. + // It can run the whole length of the scan and still be nothing but + // patient, so its size says nothing and is not asked. + // + // What stands clear of that is a separate object that happens to + // touch: an ear, a nose and a fingertip, which are small, or a pad, + // a blanket and a bolus, which are not. There, size is exactly the + // question. Two rounds, because a fingertip hangs off a finger. + let reach = params.open_mm + spacing.iter().cloned().fold(0.0, f64::max); + for _ in 0..2 { + let residue: Vec = solid + .par_iter() + .zip(body.par_iter()) + .map(|(&m, &b)| u8::from(m != 0 && b == 0)) + .collect(); + let near = morph::dilate_mm(&body, dims, spacing, reach); + let mut grew = false; + for c in morph::components(&residue, dims) { + let shaved_off_the_body = c.voxels.iter().all(|&v| near[v as usize] != 0); + if !shaved_off_the_body && c.extent_mm(spacing) > params.thin_max_extent_mm { + continue; + } + if !morph::touches(&c, &body, dims) { + continue; + } + for &v in &c.voxels { + body[v as usize] = 1; + } + recovered += c.len() as u64; + grew = true; + } + if !grew { + break; + } + } + } + if progress.cancelled() { + bail!(CANCELLED); + } + + // ---- 6. finish ------------------------------------------------------- + progress.set_phase(fg_base + 5.0 * fg_span, fg_span); + progress.report(0.0, "Finishing the surface…"); + // The opening can leave a dent the fill has to close again, and a + // cavity that is open on every slice can still be enclosed in space. + morph::fill_holes_2d(&mut body, dims, axial); + morph::fill_holes_3d(&mut body, dims); + if params.close_mm > 0.0 { + body = morph::close_mm(&body, dims, spacing, params.close_mm); + } + // Everything above the threshold that is not patient: the equipment the + // extrusion test caught, plus every component too small or too detached + // to be a body. Counted here, against the original foreground, because + // by now the body also contains an interior that was never above the + // threshold at all. + let rejected: u64 = fg + .par_iter() + .zip(body.par_iter()) + .map(|(&f, &b)| u64::from(f != 0 && b == 0)) + .sum(); + if !params.fill_interior { + body.par_iter_mut() + .zip(fg.par_iter()) + .for_each(|(b, &f)| *b &= (f != 0) as u8); + } + + // ---- 7. statistics --------------------------------------------------- + let voxels: u64 = body.par_iter().map(|&v| u64::from(v != 0)).sum(); + let pieces: Vec = morph::components(&body, dims) + .into_iter() + .filter(|c| c.len() >= min_voxels) + .map(|c| Piece { + voxels: c.len() as u64, + cm3: c.cm3(spacing), + }) + .collect(); + progress.report(1.0, "Body contour finished"); + Ok(BodyResult { + mask: body, + dims, + voxels, + cm3: voxels as f64 * voxel_cm3, + pieces, + removed_voxels: rejected, + recovered_voxels: recovered, + method: params.method, + device, + elapsed_secs: t0.elapsed().as_secs_f64(), + name: params.name.clone(), + make_external: params.make_external, + frame_of_reference_uid: volume.frame_of_reference_uid.clone(), + volume_dims: dims, + }) +} + +// --------------------------------------------------------------------------- +// Steps +// --------------------------------------------------------------------------- + +/// Everything that is not air, by the rule the modality allows. +pub fn foreground(volume: &Volume, how: Foreground) -> Vec { + match how { + Foreground::Hu(t) => { + let t = t as i16; + volume.data.par_iter().map(|&v| u8::from(v >= t)).collect() + } + Foreground::MrRelative { fraction, sigma_mm } => { + let flat = flatten_bias(volume, sigma_mm); + let hi = high_percentile(&flat, 0.99); + let t = hi * fraction; + flat.par_iter().map(|&v| u8::from(v >= t)).collect() + } + Foreground::MrOtsu { sigma_mm } => { + let flat = flatten_bias(volume, sigma_mm); + let t = otsu(&flat); + flat.par_iter().map(|&v| u8::from(v >= t)).collect() + } + } +} + +/// Divide out a smooth estimate of the receive-coil sensitivity, so that one +/// threshold holds across the whole image. +/// +/// The estimate is a **normalized convolution**: the image blurred far +/// beyond any anatomy (40 mm by default), but weighted so that only voxels +/// plausibly inside *something* contribute, and divided by the same blur of +/// the weights. A plain blur would not do — near the skin, and anywhere the +/// body is thin, it is dominated by the surrounding air and reports a bias +/// that is really just "how much background is nearby", which flattens the +/// anatomy instead of the shading. Weighting fixes exactly that, for the +/// cost of one more pass. +/// +/// It is a poor man's N4 — no iteration, no histogram model — which is all a +/// *body outline* needs, because the boundary it is looking for is the +/// largest step in the image. The output is rescaled to the mean signal +/// inside the object, so the numbers stay readable and a threshold given as +/// a fraction of the 99th percentile means the same thing before and after. +/// +/// `sigma_mm` has to sit between two scales: **well above** any anatomy, or +/// the estimate follows the tissue and flattens the very contrast the +/// threshold needs, and **well below** the body, or it degenerates into a +/// global mean and corrects nothing. 40 mm on a clinical field of view is +/// comfortably inside that window. +pub fn flatten_bias(volume: &Volume, sigma_mm: f64) -> Vec { + let raw: Vec = volume.data.par_iter().map(|&v| v as f32).collect(); + if sigma_mm <= 0.0 { + return raw; + } + // "Plausibly inside something" — deliberately generous, since this only + // has to keep the estimate off the air, not find the body. + let floor = 0.05 * high_percentile(&raw, 0.99); + let weight: Vec = raw.par_iter().map(|&v| f32::from(v >= floor)).collect(); + let signal: Vec = raw + .par_iter() + .zip(weight.par_iter()) + .map(|(&v, &w)| v * w) + .collect(); + let num = morph::blur_mm(&signal, volume.dims, volume.spacing, sigma_mm); + let den = morph::blur_mm(&weight, volume.dims, volume.spacing, sigma_mm); + // The level everything is rescaled to, and the fallback wherever the + // weight is too thin for a local estimate to mean anything. + let (sum, count) = signal + .par_iter() + .zip(weight.par_iter()) + .map(|(&v, &w)| (v as f64, w as f64)) + .reduce(|| (0.0, 0.0), |a, b| (a.0 + b.0, a.1 + b.1)); + if count == 0.0 { + return raw; + } + let mean_in = (sum / count) as f32; + let guard = (mean_in * 0.05).max(1e-3); + raw.par_iter() + .zip(num.par_iter()) + .zip(den.par_iter()) + .map(|((&v, &n), &d)| { + // A weighted mean is meaningful wherever *any* weight reached + // this voxel; only genuine 0/0, far from anything, falls back. + let bias = if d > 1e-6 { n / d } else { mean_in }; + v * mean_in / bias.max(guard) + }) + .collect() +} + +/// The value below which `q` of the (positive) samples fall — a robust +/// stand-in for the maximum, immune to a single hot voxel. +fn high_percentile(v: &[f32], q: f64) -> f32 { + let mut pos: Vec = v.par_iter().copied().filter(|x| *x > 0.0).collect(); + if pos.is_empty() { + return 0.0; + } + let k = ((pos.len() as f64 - 1.0) * q).round() as usize; + let (_, nth, _) = pos.select_nth_unstable_by(k, |a, b| a.total_cmp(b)); + *nth +} + +/// Otsu's threshold over a 256-bin histogram of `[0, p99.9]`. +fn otsu(v: &[f32]) -> f32 { + let hi = high_percentile(v, 0.999); + if hi <= 0.0 { + return 0.0; + } + let bins = 256usize; + let scale = bins as f32 / hi; + let mut hist = vec![0u64; bins]; + for &x in v { + if x <= 0.0 { + hist[0] += 1; + } else { + hist[((x * scale) as usize).min(bins - 1)] += 1; + } + } + let total: u64 = hist.iter().sum(); + let sum: f64 = hist + .iter() + .enumerate() + .map(|(i, &c)| i as f64 * c as f64) + .sum(); + let (mut w0, mut s0, mut best, mut best_t) = (0u64, 0f64, -1f64, 0usize); + for (t, &c) in hist.iter().enumerate() { + w0 += c; + if w0 == 0 { + continue; + } + let w1 = total - w0; + if w1 == 0 { + break; + } + s0 += t as f64 * c as f64; + let m0 = s0 / w0 as f64; + let m1 = (sum - s0) / w1 as f64; + let between = w0 as f64 * w1 as f64 * (m0 - m1) * (m0 - m1); + if between > best { + best = between; + best_t = t; + } + } + // The *upper* edge of the winning bin, not its centre: on a histogram + // with two well-separated peaks every split between them scores the + // same, argmax takes the first, and its centre would sit inside the + // background peak rather than above it. + (best_t as f32 + 1.0) / scale +} + +/// TotalSegmentator's own post-processing of the body task, then the union. +/// +/// The network answers in two classes — trunk and extremities — and the +/// reference implementation cleans each differently: the trunk is one +/// object, so only its largest blob is kept, while extremities are several +/// and are merely filtered for size (50 000 mm³, the same constant +/// upstream uses). The body is what is left of both. +fn body_from_labels( + labels: &[u8], + dims: [usize; 3], + spacing: [f64; 3], + progress: &Progress, +) -> Result> { + progress.report(0.95, "Cleaning up the network's answer…"); + let voxel_mm3 = spacing[0] * spacing[1] * spacing[2]; + let trunk: Vec = labels.par_iter().map(|&l| u8::from(l == 1)).collect(); + let limbs: Vec = labels.par_iter().map(|&l| u8::from(l == 2)).collect(); + let mut out = vec![0u8; labels.len()]; + if let Some(c) = morph::components(&trunk, dims).first() { + for &v in &c.voxels { + out[v as usize] = 1; + } + } + let min_voxels = (50_000.0 / voxel_mm3).max(1.0) as usize; + for c in morph::components(&limbs, dims) { + if c.len() < min_voxels { + continue; + } + for &v in &c.voxels { + out[v as usize] = 1; + } + } + if out.iter().all(|&v| v == 0) { + bail!("the body-outline network found no patient in this volume"); + } + Ok(out) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::geometry::Vec3; + + fn vol(dims: [usize; 3], spacing: [f64; 3], data: Vec) -> Volume { + Volume { + data, + dims, + spacing, + origin: Vec3::new(0.0, 0.0, 0.0), + row_dir: Vec3::new(1.0, 0.0, 0.0), + col_dir: Vec3::new(0.0, 1.0, 0.0), + normal: Vec3::new(0.0, 0.0, 1.0), + frame_of_reference_uid: "1.2.3".into(), + min_value: -1000, + max_value: 1000, + } + } + + #[test] + fn a_hounsfield_threshold_is_exactly_a_threshold() { + let v = vol([4, 1, 1], [1.0; 3], vec![-1000, -301, -300, 40]); + let f = foreground(&v, Foreground::Hu(-300.0)); + assert_eq!(f, vec![0, 0, 1, 1]); + } + + #[test] + fn flattening_rescues_a_threshold_that_the_coil_shading_had_broken() { + // The case that matters on MR, and the one a fixed threshold cannot + // survive: an exponential coil falloff steep enough that background + // *near* the coil is brighter than body *far* from it. No single + // threshold on the raw image can separate them — that is asserted + // below, not assumed. + let dims = [256, 64, 8]; + let mut data = vec![0i16; dims[0] * dims[1] * dims[2]]; + let at = |i: usize, j: usize, k: usize| k * dims[0] * dims[1] + j * dims[0] + i; + for i in 0..dims[0] { + let gain = (-4.0 * i as f32 / dims[0] as f32).exp(); + for j in 0..dims[1] { + let inside = (16..240).contains(&i) && (8..56).contains(&j); + let v = if inside { 1000.0 } else { 60.0 } * gain; + for k in 0..dims[2] { + data[at(i, j, k)] = v.round() as i16; + } + } + } + let v = vol(dims, [1.0; 3], data); + let (body_near, body_far) = (at(30, 32, 4), at(225, 32, 4)); + let (air_near, air_far) = (at(30, 2, 4), at(225, 2, 4)); + assert!( + v.data[body_far] < v.data[air_near], + "the phantom has to be unthresholdable to be worth the test: \ + body {} vs air {}", + v.data[body_far], + v.data[air_near] + ); + + let flat = flatten_bias(&v, 20.0); + assert!( + flat[body_far] > flat[air_near], + "flattening did not reorder body and air: {} vs {}", + flat[body_far], + flat[air_near] + ); + let raw_ratio = v.data[body_near] as f32 / v.data[body_far] as f32; + let flat_ratio = flat[body_near] / flat[body_far]; + assert!( + flat_ratio < raw_ratio / 4.0, + "residual shading {flat_ratio:.1} was barely better than {raw_ratio:.1}" + ); + + // The rule the tool actually applies, on the same phantom. + let mask = foreground( + &v, + Foreground::MrRelative { + fraction: 0.12, + sigma_mm: 20.0, + }, + ); + assert_eq!(mask[body_near], 1, "the bright end of the body"); + assert_eq!(mask[body_far], 1, "the dim end of the body"); + assert_eq!(mask[air_near], 0, "bright air near the coil"); + assert_eq!(mask[air_far], 0, "dim air"); + } + + #[test] + fn otsu_splits_a_two_peaked_histogram_between_the_peaks() { + let mut v = vec![10.0f32; 500]; + v.extend(std::iter::repeat_n(900.0f32, 500)); + let t = otsu(&v); + assert!((10.0..900.0).contains(&t), "threshold {t}"); + } + + #[test] + fn the_defaults_follow_the_modality() { + let ct = BodyParams::for_modality("CT"); + assert_eq!(ct.foreground, Foreground::Hu(-300.0)); + assert_eq!(ct.model, BodyModel::Ct6mm); + let mr = BodyParams::for_modality("mr"); + assert!(mr.foreground.is_mr()); + assert_eq!(mr.model, BodyModel::Mr); + } +} diff --git a/src/lib.rs b/src/lib.rs index da4756e..b6b4f27 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -4,6 +4,7 @@ pub mod anonymize; pub mod app; pub mod autoseg; +pub mod bodymask; pub mod dicom_export; pub mod dicomseg; pub mod drr; @@ -14,6 +15,7 @@ pub mod loader; pub mod medsam2; pub mod mesh3d; pub mod models; +pub mod morphology; pub mod nn; pub mod progress; pub mod propagate; diff --git a/src/models.rs b/src/models.rs index 4d44400..cb20ae0 100644 --- a/src/models.rs +++ b/src/models.rs @@ -185,6 +185,14 @@ pub fn inventory() -> Vec { "All 117 structures at 3 mm — the fast default." } else if spec.key == autoseg::weights::SPEC_6MM.key { "Coarse preview quality — the quickest look." + } else if spec.key == autoseg::weights::SPEC_BODY_6MM.key { + "Patient outline, 6 mm — what the body-contour tool's model-assisted \ + method uses. Plenty, because it only decides which side of the skin \ + a voxel is on." + } else if spec.key == autoseg::weights::SPEC_BODY_15MM.key { + "Patient outline at full resolution; slower, for the same decision." + } else if spec.key == autoseg::weights::SPEC_BODY_MR.key { + "Patient outline on MR — the body-contour tool's model for MR series." } else { "Full-resolution sub-model; the five together are the reference quality." }; diff --git a/src/morphology.rs b/src/morphology.rs new file mode 100644 index 0000000..2e13d1a --- /dev/null +++ b/src/morphology.rs @@ -0,0 +1,799 @@ +//! Binary-mask morphology on voxel grids — the geometry every mask-shaped +//! feature needs and nobody should write twice. +//! +//! Masks are `[u8]` in [`crate::volume::Volume`] index order +//! (`k * nx * ny + j * nx + i`), non-zero meaning *set*. Every operation that +//! has a physical size takes it in **millimetres** and reads the voxel +//! spacing, so a 5 mm opening is 5 mm along each axis whatever the slice +//! thickness — which is the whole point on clinical CT, where in-plane and +//! through-plane spacing routinely differ by a factor of five. +//! +//! The distance transform is the exact anisotropic Euclidean one +//! (Felzenszwalb & Huttenlocher's lower envelope of parabolas, *Theory of +//! Computing* 2012): three separable O(n) passes, no approximation by +//! chamfer weights, no dependence of the cost on the radius. Erosion and +//! dilation are then thresholds on it, so an opening costs four passes +//! whether the radius is 1 mm or 50 mm. +//! +//! Outside the volume counts as **set** for the distance-to-background +//! transform. That convention is what keeps anatomy truncated by the scan +//! FoV — the arms at the edge of the field, the body at the first and last +//! slice — from being eroded away at the cut: nothing is inferred about +//! what was never imaged. + +use rayon::prelude::*; + +/// Squared distance (mm²) from every voxel to the nearest **unset** voxel; +/// zero on unset voxels. Voxels outside the volume count as set. +pub fn dist2_to_background(mask: &[u8], dims: [usize; 3], spacing: [f64; 3]) -> Vec { + let n = dims[0] * dims[1] * dims[2]; + debug_assert_eq!(mask.len(), n); + // Seed: 0 where unset, +inf where set. + let mut f: Vec = mask + .par_iter() + .map(|&v| if v != 0 { f32::INFINITY } else { 0.0 }) + .collect(); + edt_in_place(&mut f, dims, spacing); + f +} + +/// Squared distance (mm²) from every voxel to the nearest **set** voxel; +/// zero on set voxels. +pub fn dist2_to_foreground(mask: &[u8], dims: [usize; 3], spacing: [f64; 3]) -> Vec { + let n = dims[0] * dims[1] * dims[2]; + debug_assert_eq!(mask.len(), n); + let mut f: Vec = mask + .par_iter() + .map(|&v| if v == 0 { f32::INFINITY } else { 0.0 }) + .collect(); + edt_in_place(&mut f, dims, spacing); + f +} + +/// Three separable passes of the 1-D squared-distance transform. +fn edt_in_place(f: &mut [f32], dims: [usize; 3], spacing: [f64; 3]) { + for (axis, step) in spacing.iter().enumerate() { + pass_along(f, dims, axis, *step as f32); + } +} + +/// Stride and length of the lines running along `axis`, and the start index +/// of each line. +fn lines_along(dims: [usize; 3], axis: usize) -> (usize, usize, Vec) { + let [nx, ny, nz] = dims; + match axis { + 0 => (nx, 1, (0..ny * nz).map(|l| l * nx).collect()), + 1 => ( + ny, + nx, + (0..nx * nz) + .map(|l| (l / nx) * nx * ny + (l % nx)) + .collect(), + ), + _ => (nz, nx * ny, (0..nx * ny).collect()), + } +} + +/// The 1-D lower envelope of parabolas along one axis, in place. +fn pass_along(f: &mut [f32], dims: [usize; 3], axis: usize, step: f32) { + let (n, stride, starts) = lines_along(dims, axis); + if n == 0 { + return; + } + let sq = step * step; + // Lines along an axis are disjoint, so each can be lifted out, solved + // and written back independently. + let out: Vec<(usize, Vec)> = starts + .par_iter() + .map(|&base| { + let mut d = vec![0f32; n]; + let mut v = vec![0usize; n]; // parabola centres + let mut z = vec![0f32; n + 1]; // envelope breakpoints + for (q, slot) in d.iter_mut().enumerate() { + *slot = f[base + q * stride]; + } + let mut k = 0usize; + v[0] = 0; + z[0] = f32::NEG_INFINITY; + z[1] = f32::INFINITY; + for q in 1..n { + if d[q].is_infinite() { + // A parabola of infinite height never joins the envelope. + continue; + } + loop { + let p = v[k]; + // Intersection of the parabolas rooted at p and q. + let s = if d[p].is_infinite() { + f32::NEG_INFINITY + } else { + (d[q] + sq * (q * q) as f32 - d[p] - sq * (p * p) as f32) + / (2.0 * sq * (q as f32 - p as f32)) + }; + if s <= z[k] && k > 0 { + k -= 1; + } else { + k += 1; + v[k] = q; + z[k] = s; + z[k + 1] = f32::INFINITY; + break; + } + } + } + // Walk the envelope left to right. + let mut line = vec![0f32; n]; + let mut k = 0usize; + for (q, slot) in line.iter_mut().enumerate() { + while z[k + 1] < q as f32 { + k += 1; + } + let p = v[k]; + *slot = if d[p].is_infinite() { + f32::INFINITY + } else { + let dq = (q as f32 - p as f32) * step; + d[p] + dq * dq + }; + } + (base, line) + }) + .collect(); + for (base, line) in out { + for (q, val) in line.into_iter().enumerate() { + f[base + q * stride] = val; + } + } +} + +/// Erosion by a ball of `radius_mm`: the voxels further than the radius from +/// any background voxel. Voxels outside the volume are not background, so a +/// mask that runs into the volume boundary is not eroded there. +pub fn erode_mm(mask: &[u8], dims: [usize; 3], spacing: [f64; 3], radius_mm: f64) -> Vec { + if radius_mm <= 0.0 { + return mask.to_vec(); + } + let r2 = (radius_mm * radius_mm) as f32; + let d = dist2_to_background(mask, dims, spacing); + d.par_iter().map(|&v| u8::from(v > r2)).collect() +} + +/// Dilation by a ball of `radius_mm`. +pub fn dilate_mm(mask: &[u8], dims: [usize; 3], spacing: [f64; 3], radius_mm: f64) -> Vec { + if radius_mm <= 0.0 { + return mask.to_vec(); + } + let r2 = (radius_mm * radius_mm) as f32; + let d = dist2_to_foreground(mask, dims, spacing); + d.par_iter().map(|&v| u8::from(v <= r2)).collect() +} + +/// Opening — erosion then dilation. Equivalently: the union of every ball of +/// `radius_mm` that fits entirely inside the mask, so everything thinner +/// than twice the radius disappears and everything thicker keeps its exact +/// surface. +pub fn open_mm(mask: &[u8], dims: [usize; 3], spacing: [f64; 3], radius_mm: f64) -> Vec { + if radius_mm <= 0.0 { + return mask.to_vec(); + } + let eroded = erode_mm(mask, dims, spacing, radius_mm); + dilate_mm(&eroded, dims, spacing, radius_mm) +} + +/// Closing — dilation then erosion. Bridges gaps narrower than twice the +/// radius; used to take the staircase off a contour. +pub fn close_mm(mask: &[u8], dims: [usize; 3], spacing: [f64; 3], radius_mm: f64) -> Vec { + if radius_mm <= 0.0 { + return mask.to_vec(); + } + let dilated = dilate_mm(mask, dims, spacing, radius_mm); + erode_mm(&dilated, dims, spacing, radius_mm) +} + +// --------------------------------------------------------------------------- +// Connected components +// --------------------------------------------------------------------------- + +/// One 6-connected component: the voxels it owns and its bounding box. +/// +/// The voxel list rather than a per-voxel label volume is deliberate: a +/// label volume costs four bytes for every voxel of the *image* (313 MB on a +/// 512 × 512 × 300 CT), the lists together cost four bytes per *set* voxel, +/// which for a body mask is an order of magnitude less. +#[derive(Clone, Debug)] +pub struct Component { + pub voxels: Vec, + /// Inclusive voxel bounding box. + pub lo: [usize; 3], + pub hi: [usize; 3], +} + +impl Component { + pub fn len(&self) -> usize { + self.voxels.len() + } + pub fn is_empty(&self) -> bool { + self.voxels.is_empty() + } + /// Volume in cm³ for the given voxel spacing. + pub fn cm3(&self, spacing: [f64; 3]) -> f64 { + self.voxels.len() as f64 * spacing[0] * spacing[1] * spacing[2] / 1000.0 + } + /// Longest side of the bounding box, in millimetres. + pub fn extent_mm(&self, spacing: [f64; 3]) -> f64 { + (0..3) + .map(|a| (self.hi[a] - self.lo[a] + 1) as f64 * spacing[a]) + .fold(0.0, f64::max) + } +} + +/// Every 6-connected component of the mask, largest first. +pub fn components(mask: &[u8], dims: [usize; 3]) -> Vec { + let [nx, ny, nz] = dims; + let sl = nx * ny; + let n = sl * nz; + debug_assert_eq!(mask.len(), n); + let mut seen = vec![false; n]; + let mut out: Vec = Vec::new(); + let mut stack: Vec = Vec::new(); + for start in 0..n { + if mask[start] == 0 || seen[start] { + continue; + } + seen[start] = true; + stack.push(start as u32); + let mut voxels: Vec = Vec::new(); + let mut lo = [usize::MAX; 3]; + let mut hi = [0usize; 3]; + while let Some(p) = stack.pop() { + let idx = p as usize; + let (i, j, k) = (idx % nx, (idx % sl) / nx, idx / sl); + for (a, c) in [i, j, k].into_iter().enumerate() { + lo[a] = lo[a].min(c); + hi[a] = hi[a].max(c); + } + voxels.push(p); + // 6-neighbourhood. Diagonal contact is not contact: two organs + // that share only a corner are two components, and so are a + // couch rail and the skin it grazes. + if i > 0 { + push(&mut stack, &mut seen, mask, idx - 1); + } + if i + 1 < nx { + push(&mut stack, &mut seen, mask, idx + 1); + } + if j > 0 { + push(&mut stack, &mut seen, mask, idx - nx); + } + if j + 1 < ny { + push(&mut stack, &mut seen, mask, idx + nx); + } + if k > 0 { + push(&mut stack, &mut seen, mask, idx - sl); + } + if k + 1 < nz { + push(&mut stack, &mut seen, mask, idx + sl); + } + } + out.push(Component { voxels, lo, hi }); + } + out.sort_by_key(|c| std::cmp::Reverse(c.voxels.len())); + out +} + +#[inline] +fn push(stack: &mut Vec, seen: &mut [bool], mask: &[u8], idx: usize) { + if mask[idx] != 0 && !seen[idx] { + seen[idx] = true; + stack.push(idx as u32); + } +} + +/// True when any voxel of `comp` has a 6-neighbour set in `other`. +pub fn touches(comp: &Component, other: &[u8], dims: [usize; 3]) -> bool { + let [nx, ny, nz] = dims; + let sl = nx * ny; + comp.voxels.par_iter().any(|&p| { + let idx = p as usize; + let (i, j, k) = (idx % nx, (idx % sl) / nx, idx / sl); + (i > 0 && other[idx - 1] != 0) + || (i + 1 < nx && other[idx + 1] != 0) + || (j > 0 && other[idx - nx] != 0) + || (j + 1 < ny && other[idx + nx] != 0) + || (k > 0 && other[idx - sl] != 0) + || (k + 1 < nz && other[idx + sl] != 0) + }) +} + +// --------------------------------------------------------------------------- +// Hole filling +// --------------------------------------------------------------------------- + +/// Fill every background region of a slice that the slice border cannot +/// reach, for each slice perpendicular to `axis`. +/// +/// Two-dimensional on purpose. A lung is connected to the outside air +/// through the trachea, so a three-dimensional fill leaves both lungs open +/// on any scan that includes the neck; slice by slice they close. +pub fn fill_holes_2d(mask: &mut [u8], dims: [usize; 3], axis: usize) { + let [nx, ny, nz] = dims; + let sl = nx * ny; + // In-plane extents and the two in-plane strides, per slicing axis. + let (n_slices, slice_stride, (w, h), (su, sv)) = match axis { + 0 => (nx, 1usize, (ny, nz), (nx, sl)), + 1 => (ny, nx, (nx, nz), (1usize, sl)), + _ => (nz, sl, (nx, ny), (1usize, nx)), + }; + if n_slices == 0 || w == 0 || h == 0 { + return; + } + let filled: Vec<(usize, Vec)> = (0..n_slices) + .into_par_iter() + .map(|s| { + let base = s * slice_stride; + let mut reach = vec![false; w * h]; + let mut stack: Vec = Vec::new(); + let at = |u: usize, v: usize| base + u * su + v * sv; + // Seed from the whole slice border. + let seed = |u: usize, v: usize, reach: &mut Vec, stack: &mut Vec| { + let p = v * w + u; + if !reach[p] && mask[at(u, v)] == 0 { + reach[p] = true; + stack.push(p); + } + }; + for u in 0..w { + seed(u, 0, &mut reach, &mut stack); + seed(u, h - 1, &mut reach, &mut stack); + } + for v in 0..h { + seed(0, v, &mut reach, &mut stack); + seed(w - 1, v, &mut reach, &mut stack); + } + while let Some(p) = stack.pop() { + let (u, v) = (p % w, p / w); + let visit = |qu: usize, qv: usize, reach: &mut Vec, st: &mut Vec| { + let q = qv * w + qu; + if !reach[q] && mask[at(qu, qv)] == 0 { + reach[q] = true; + st.push(q); + } + }; + if u > 0 { + visit(u - 1, v, &mut reach, &mut stack); + } + if u + 1 < w { + visit(u + 1, v, &mut reach, &mut stack); + } + if v > 0 { + visit(u, v - 1, &mut reach, &mut stack); + } + if v + 1 < h { + visit(u, v + 1, &mut reach, &mut stack); + } + } + let mut add: Vec = Vec::new(); + for v in 0..h { + for u in 0..w { + let idx = at(u, v); + if mask[idx] == 0 && !reach[v * w + u] { + add.push(idx as u32); + } + } + } + (s, add) + }) + .collect(); + for (_, add) in filled { + for idx in add { + mask[idx as usize] = 1; + } + } +} + +/// Fill every background region the volume border cannot reach. +pub fn fill_holes_3d(mask: &mut [u8], dims: [usize; 3]) { + let [nx, ny, nz] = dims; + let sl = nx * ny; + let n = sl * nz; + let mut reach = vec![false; n]; + let mut stack: Vec = Vec::new(); + let seed = |idx: usize, reach: &mut Vec, stack: &mut Vec| { + if mask[idx] == 0 && !reach[idx] { + reach[idx] = true; + stack.push(idx as u32); + } + }; + for k in 0..nz { + for j in 0..ny { + for i in 0..nx { + if i == 0 || j == 0 || k == 0 || i + 1 == nx || j + 1 == ny || k + 1 == nz { + seed(k * sl + j * nx + i, &mut reach, &mut stack); + } + } + } + } + while let Some(p) = stack.pop() { + let idx = p as usize; + let (i, j, k) = (idx % nx, (idx % sl) / nx, idx / sl); + let visit = |q: usize, reach: &mut Vec, st: &mut Vec| { + if mask[q] == 0 && !reach[q] { + reach[q] = true; + st.push(q as u32); + } + }; + if i > 0 { + visit(idx - 1, &mut reach, &mut stack); + } + if i + 1 < nx { + visit(idx + 1, &mut reach, &mut stack); + } + if j > 0 { + visit(idx - nx, &mut reach, &mut stack); + } + if j + 1 < ny { + visit(idx + nx, &mut reach, &mut stack); + } + if k > 0 { + visit(idx - sl, &mut reach, &mut stack); + } + if k + 1 < nz { + visit(idx + sl, &mut reach, &mut stack); + } + } + mask.par_iter_mut() + .zip(reach.par_iter()) + .for_each(|(m, r)| { + if *m == 0 && !*r { + *m = 1; + } + }); +} + +// --------------------------------------------------------------------------- +// Axis persistence +// --------------------------------------------------------------------------- + +/// Mark the voxels whose *column* along `axis` is occupied in at least +/// `frac` of the slices of a `window_mm` neighbourhood. +/// +/// This is the shape prior that separates equipment from anatomy without +/// knowing what either looks like. A couch top, a chair backrest, a seat +/// pan, an arm rest, a bright reconstruction-circle rim: each is a surface +/// swept along one axis, so its footprint in the orthogonal plane repeats +/// slice after slice after slice. A pinna, a nose, a finger — the thin +/// anatomy a plain opening also removes — never repeats over anything like +/// the same distance, so a window of 150 mm at 80 % separates them cleanly. +/// +/// The scan has to contain slices where the columns in question are free of +/// the patient. A couch strip directly beneath the body over the *entire* +/// scan length is not distinguishable this way, and is left to the +/// model-assisted method. +pub fn axis_persistence( + mask: &[u8], + dims: [usize; 3], + spacing: [f64; 3], + axis: usize, + window_mm: f64, + frac: f64, +) -> Vec { + let [nx, ny, nz] = dims; + let sl = nx * ny; + let n = sl * nz; + debug_assert_eq!(mask.len(), n); + // Slice count and stride along `axis`, and the in-plane geometry. + let (n_slices, slice_stride, w, h, su, sv) = match axis { + 0 => (nx, 1usize, ny, nz, nx, sl), + 1 => (ny, nx, nx, nz, 1usize, sl), + _ => (nz, sl, nx, ny, 1usize, nx), + }; + let plane = w * h; + let mut out = vec![0u8; n]; + if n_slices == 0 || plane == 0 { + return out; + } + let win = ((window_mm / spacing[axis]).round() as usize).clamp(1, n_slices); + let need = ((win as f64) * frac).ceil().max(1.0) as u32; + let at = |s: usize, p: usize| s * slice_stride + (p % w) * su + (p / w) * sv; + + // The window [lo, hi) advances monotonically with `s`, so each slice is + // added once and removed once: the whole scan costs two passes. + let mut count = vec![0u32; plane]; + let (mut cur_lo, mut cur_hi) = (0usize, 0usize); + for s in 0..n_slices { + let lo = s.saturating_sub(win / 2).min(n_slices - win); + let hi = lo + win; + while cur_hi < hi { + for (p, c) in count.iter_mut().enumerate() { + *c += u32::from(mask[at(cur_hi, p)] != 0); + } + cur_hi += 1; + } + while cur_lo < lo { + for (p, c) in count.iter_mut().enumerate() { + *c -= u32::from(mask[at(cur_lo, p)] != 0); + } + cur_lo += 1; + } + for (p, &c) in count.iter().enumerate() { + if c >= need { + let idx = at(s, p); + if mask[idx] != 0 { + out[idx] = 1; + } + } + } + } + out +} + +// --------------------------------------------------------------------------- +// Smoothing +// --------------------------------------------------------------------------- + +/// Three successive box blurs of `sigma_mm` — a close approximation of a +/// Gaussian (the central limit theorem does the work) at O(voxels) per axis +/// whatever the width, which is what makes a 40 mm blur of a whole MR study +/// affordable. Used to estimate the receive-coil bias field. +pub fn blur_mm(src: &[f32], dims: [usize; 3], spacing: [f64; 3], sigma_mm: f64) -> Vec { + let mut buf = src.to_vec(); + if sigma_mm <= 0.0 { + return buf; + } + for (axis, step) in spacing.iter().enumerate() { + // A box of width w has variance (w² − 1)/12; three of them give 3×. + let var = sigma_mm * sigma_mm / (step * step) / 3.0; + let w = ((12.0 * var + 1.0).sqrt().round() as usize).max(1) | 1; + if w <= 1 { + continue; + } + for _ in 0..3 { + box_pass(&mut buf, dims, axis, w); + } + } + buf +} + +/// One box blur of odd width `w` along `axis`, edges clamped. +fn box_pass(buf: &mut [f32], dims: [usize; 3], axis: usize, w: usize) { + let (n, stride, starts) = lines_along(dims, axis); + if n == 0 || w <= 1 { + return; + } + let r = (w / 2) as isize; + let last = n as isize - 1; + let out: Vec<(usize, Vec)> = starts + .par_iter() + .map(|&base| { + let line: Vec = (0..n).map(|q| buf[base + q * stride]).collect(); + let cl = |t: isize| line[t.clamp(0, last) as usize]; + let mut acc: f32 = (-r..=r).map(cl).sum(); + let mut res = vec![0f32; n]; + for (q, slot) in res.iter_mut().enumerate() { + *slot = acc / w as f32; + acc += cl(q as isize + r + 1) - cl(q as isize - r); + } + (base, res) + }) + .collect(); + for (base, line) in out { + for (q, v) in line.into_iter().enumerate() { + buf[base + q * stride] = v; + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Brute-force squared distance to the nearest background voxel, with + /// the same "outside is not background" convention. + fn brute(mask: &[u8], dims: [usize; 3], sp: [f64; 3]) -> Vec { + let [nx, ny, nz] = dims; + let mut out = vec![f32::INFINITY; nx * ny * nz]; + for k in 0..nz { + for j in 0..ny { + for i in 0..nx { + let idx = k * nx * ny + j * nx + i; + if mask[idx] == 0 { + out[idx] = 0.0; + continue; + } + let mut best = f32::INFINITY; + for kk in 0..nz { + for jj in 0..ny { + for ii in 0..nx { + if mask[kk * nx * ny + jj * nx + ii] != 0 { + continue; + } + let d = (((i as f64 - ii as f64) * sp[0]).powi(2) + + ((j as f64 - jj as f64) * sp[1]).powi(2) + + ((k as f64 - kk as f64) * sp[2]).powi(2)) + as f32; + best = best.min(d); + } + } + } + out[idx] = best; + } + } + } + out + } + + #[test] + fn the_distance_transform_is_exact_and_anisotropic() { + let dims = [9, 7, 5]; + let sp = [0.8, 1.3, 3.0]; + // A deterministic pseudo-random mask. + let mut state = 12345u32; + let mask: Vec = (0..dims[0] * dims[1] * dims[2]) + .map(|_| { + state = state.wrapping_mul(1664525).wrapping_add(1013904223); + u8::from(!(state >> 28).is_multiple_of(4)) + }) + .collect(); + let got = dist2_to_background(&mask, dims, sp); + let want = brute(&mask, dims, sp); + for (g, w) in got.iter().zip(want.iter()) { + assert!( + (g - w).abs() < 1e-3 || (g.is_infinite() && w.is_infinite()), + "{g} vs {w}" + ); + } + } + + #[test] + fn a_fully_set_volume_has_no_background_to_measure() { + let dims = [4, 4, 4]; + let mask = vec![1u8; 64]; + // Outside is not background, so nothing is within any finite + // distance — the convention that stops truncated anatomy eroding. + assert!(dist2_to_background(&mask, dims, [1.0; 3]) + .iter() + .all(|d| d.is_infinite())); + assert_eq!(erode_mm(&mask, dims, [1.0; 3], 100.0), vec![1u8; 64]); + } + + #[test] + fn opening_removes_sheets_thinner_than_the_ball_and_keeps_the_rest() { + // A 40 × 40 × 40 mm block and, 10 mm away, a 2 mm sheet. + let dims = [40, 40, 20]; + let sp = [1.0, 1.0, 2.0]; + let mut mask = vec![0u8; dims[0] * dims[1] * dims[2]]; + let at = |i: usize, j: usize, k: usize| k * dims[0] * dims[1] + j * dims[0] + i; + for k in 2..18 { + for j in 4..36 { + for i in 4..36 { + mask[at(i, j, k)] = 1; + } + } + } + // A one-voxel (2 mm) sheet along the whole i/k extent. + for k in 0..dims[2] { + for i in 0..dims[0] { + mask[at(i, 1, k)] = 1; + } + } + let opened = open_mm(&mask, dims, sp, 5.0); + for k in 0..dims[2] { + for i in 0..dims[0] { + assert_eq!(opened[at(i, 1, k)], 0, "the sheet survived at i={i} k={k}"); + } + } + // The block's interior is untouched, and so is its surface. + assert_eq!(opened[at(20, 20, 10)], 1); + assert_eq!(opened[at(20, 4, 10)], 1, "the block's own face is kept"); + } + + #[test] + fn components_are_six_connected_and_sorted_by_size() { + let dims = [6, 6, 1]; + let mut mask = vec![0u8; 36]; + // Two 2×2 squares touching only at a corner. + for (i, j) in [(0, 0), (1, 0), (0, 1), (1, 1)] { + mask[j * 6 + i] = 1; + } + for (i, j) in [(2, 2), (3, 2), (2, 3), (3, 3), (4, 3)] { + mask[j * 6 + i] = 1; + } + let c = components(&mask, dims); + assert_eq!(c.len(), 2, "a shared corner is not contact"); + assert_eq!(c[0].len(), 5, "largest first"); + assert_eq!(c[1].len(), 4); + assert_eq!(c[0].lo, [2, 2, 0]); + assert_eq!(c[0].hi, [4, 3, 0]); + } + + #[test] + fn slicewise_filling_closes_a_lung_that_a_three_d_fill_leaves_open() { + // A block with a cavity that connects to the outside on one slice + // only — a lung and its trachea. + let dims = [12, 12, 6]; + let at = |i: usize, j: usize, k: usize| k * 144 + j * 12 + i; + let mut mask = vec![0u8; 12 * 12 * 6]; + for k in 0..6 { + for j in 1..11 { + for i in 1..11 { + mask[at(i, j, k)] = 1; + } + } + } + // Cavity in the middle of every slice… + for k in 0..6 { + for j in 4..8 { + for i in 4..8 { + mask[at(i, j, k)] = 0; + } + } + } + // …with a channel to the outside on slice 0 only. + for j in 0..4 { + mask[at(5, j, 0)] = 0; + } + let mut three_d = mask.clone(); + fill_holes_3d(&mut three_d, dims); + assert_eq!( + three_d[at(6, 6, 3)], + 0, + "the whole cavity drains through the one channel" + ); + let mut two_d = mask.clone(); + fill_holes_2d(&mut two_d, dims, 2); + for k in 1..6 { + assert_eq!(two_d[at(6, 6, k)], 1, "slice {k} closed"); + } + assert_eq!( + two_d[at(6, 6, 0)], + 0, + "the slice with the channel stays open" + ); + } + + #[test] + fn persistence_marks_an_extruded_rail_and_spares_a_bump() { + let dims = [20, 20, 40]; + let sp = [1.0, 1.0, 2.0]; // 80 mm along k + let at = |i: usize, j: usize, k: usize| k * 400 + j * 20 + i; + let mut mask = vec![0u8; 20 * 20 * 40]; + // A rail running the whole length in k. + for k in 0..40 { + for i in 5..15 { + mask[at(i, 2, k)] = 1; + } + } + // A bump 8 mm long (4 slices). + for k in 18..22 { + for i in 8..12 { + mask[at(i, 15, k)] = 1; + } + } + let p = axis_persistence(&mask, dims, sp, 2, 60.0, 0.8); + assert_eq!(p[at(9, 2, 20)], 1, "the rail is persistent"); + assert_eq!(p[at(9, 15, 20)], 0, "the bump is not"); + // Nothing outside the mask is ever marked. + assert!(p.iter().zip(mask.iter()).all(|(a, m)| *a == 0 || *m != 0)); + } + + #[test] + fn blurring_preserves_a_constant_and_spreads_a_step() { + let dims = [32, 4, 4]; + let n = dims[0] * dims[1] * dims[2]; + let flat = vec![7.0f32; n]; + let out = blur_mm(&flat, dims, [1.0; 3], 4.0); + for v in &out { + assert!((v - 7.0).abs() < 1e-3, "{v}"); + } + let step: Vec = (0..n) + .map(|idx| if idx % dims[0] < 16 { 0.0 } else { 100.0 }) + .collect(); + let sm = blur_mm(&step, dims, [1.0; 3], 4.0); + let mid = sm[15]; + assert!((0.0..100.0).contains(&mid), "the step is smoothed: {mid}"); + assert!( + sm[0] < 5.0 && sm[dims[0] - 1] > 95.0, + "the plateaus survive" + ); + } +} diff --git a/src/nn/tensor.rs b/src/nn/tensor.rs index 5e1e9fe..d4fb64d 100644 --- a/src/nn/tensor.rs +++ b/src/nn/tensor.rs @@ -170,6 +170,103 @@ impl SendPtr { /// Transposed 3D convolution with kernel = stride = 2: every input voxel /// projects to a disjoint 2x2x2 output block. `weight`: `[cin, cout, 2, 2, 2]` /// — PyTorch's `ConvTranspose3d` layout. +/// Transposed 3-D convolution with `kernel = stride`, for any stride. +/// +/// The 2× case has its own hand-tuned routine below; this is the general +/// one, used by the models whose decoder upsamples anisotropically (the MR +/// body model halves two axes at a time and leaves the third alone). Since +/// kernel equals stride the output tiles are disjoint — no overlap-add, +/// just a GEMM that expands every input voxel into its `s0·s1·s2` outputs +/// and a scatter. +pub fn conv_transpose3d_stride( + x: &Act, + weight: &[f32], + bias: &[f32], + cout: usize, + stride: [usize; 3], +) -> Act { + let [s0, s1, s2] = stride; + if stride == [2, 2, 2] { + return conv_transpose3d_2x(x, weight, bias, cout); + } + let (cin, d, h, w) = (x.c, x.d, x.h, x.w); + let ks = s0 * s1 * s2; + debug_assert_eq!(weight.len(), cin * cout * ks); + let (od, oh, ow) = (d * s0, h * s1, w * s2); + // Repack the weight as [cout*ks, cin] for a row-major GEMM. + let mut wt = vec![0f32; cout * ks * cin]; + for ci in 0..cin { + for co in 0..cout { + for t in 0..ks { + wt[(co * ks + t) * cin + ci] = weight[(ci * cout + co) * ks + t]; + } + } + } + let hw = h * w; + let ohw = oh * ow; + let mut out = Act::zeros(cout, od, oh, ow); + let out_ptr = SendPtr(out.data.as_mut_ptr()); + let od_stride = od * ohw; + // Input slice z writes output slices z*s0 .. z*s0+s0, disjoint across z. + (0..d).into_par_iter().for_each(|z| { + let mut xin = vec![0f32; cin * hw]; + for c in 0..cin { + let src = &x.data[c * d * hw + z * hw..c * d * hw + (z + 1) * hw]; + xin[c * hw..(c + 1) * hw].copy_from_slice(src); + } + let mut tmp = vec![0f32; cout * ks * hw]; + unsafe { + gemm::gemm( + cout * ks, + hw, + cin, + tmp.as_mut_ptr(), + 1, + hw as isize, + false, + wt.as_ptr(), + 1, + cin as isize, + xin.as_ptr(), + 1, + hw as isize, + 0.0f32, + 1.0f32, + false, + false, + false, + gemm::Parallelism::None, + ); + } + for co in 0..cout { + let bv = bias[co]; + for dz in 0..s0 { + let obase = co * od_stride + (z * s0 + dz) * ohw; + for dy in 0..s1 { + for dx in 0..s2 { + let t = (dz * s1 + dy) * s2 + dx; + let src = &tmp[(co * ks + t) * hw..(co * ks + t + 1) * hw]; + for y in 0..h { + let orow = obase + (y * s1 + dy) * ow + dx; + let dst = unsafe { + std::slice::from_raw_parts_mut( + out_ptr.get().add(orow), + (w - 1) * s2 + 1, + ) + }; + let srow = &src[y * w..(y + 1) * w]; + for (xi, sv) in srow.iter().enumerate() { + dst[xi * s2] = sv + bv; + } + } + } + } + } + } + }); + out +} + pub fn conv_transpose3d_2x(x: &Act, weight: &[f32], bias: &[f32], cout: usize) -> Act { let (cin, d, h, w) = (x.c, x.d, x.h, x.w); debug_assert_eq!(weight.len(), cin * cout * 8); diff --git a/tests/autoseg.rs b/tests/autoseg.rs index 4e20aef..a7a9bc0 100644 --- a/tests/autoseg.rs +++ b/tests/autoseg.rs @@ -46,6 +46,7 @@ fn tiny_unet_assembles_and_runs() { strides: vec![[1, 1, 1], [2, 2, 2], [2, 2, 2]], n_conv_per_stage: vec![2, 2, 2], n_conv_per_stage_decoder: vec![2, 2], + norm: rust_dicom_station::autoseg::config::Norm::Ct, clip_lo: -100.0, clip_hi: 100.0, mean: 0.0, @@ -141,6 +142,7 @@ fn shape_mismatch_is_rejected() { strides: vec![[1, 1, 1], [2, 2, 2]], n_conv_per_stage: vec![2, 2], n_conv_per_stage_decoder: vec![2], + norm: rust_dicom_station::autoseg::config::Norm::Ct, clip_lo: 0.0, clip_hi: 1.0, mean: 0.0, diff --git a/tests/body.rs b/tests/body.rs new file mode 100644 index 0000000..b866054 --- /dev/null +++ b/tests/body.rs @@ -0,0 +1,408 @@ +//! The body contour, tested against a phantom whose ground truth is known +//! by construction. +//! +//! The point of a synthetic phantom is that every failure mode of the real +//! problem can be built into it deliberately and then asserted on: a couch +//! skin and a rail under the back, a moulded mask shell hugging the face +//! with a real air gap *and* a real contact patch, ears that a plain opening +//! would shave off, lungs draining through an airway, a cable. Real scans +//! differ in a hundred ways, but a method that gets these wrong is broken, +//! and one that gets them right is worth taking to real data. +//! +//! The voxels are deliberately anisotropic (2 × 2 × 5 mm), because that is +//! where a voxel-counting implementation goes wrong and a millimetre-aware +//! one does not. + +use rust_dicom_station::bodymask::{contour_body, BodyParams, Foreground, Method}; +use rust_dicom_station::geometry::Vec3; +use rust_dicom_station::progress::Progress; +use rust_dicom_station::volume::Volume; + +const NX: usize = 100; +const NY: usize = 70; +const NZ: usize = 40; +const SP: [f64; 3] = [2.0, 2.0, 5.0]; +/// The body: an elliptical cylinder 160 × 104 mm, centred in the field. +const CX: f64 = 100.0; +const CY: f64 = 62.0; +const RX: f64 = 80.0; +const RY: f64 = 52.0; +/// Where the moulded shell is pressed against the skin instead of standing +/// off it — the one thing geometry cannot undo. +const CONTACT: std::ops::Range = 46..55; + +fn idx(i: usize, j: usize, k: usize) -> usize { + k * NX * NY + j * NX + i +} + +/// The first row index at or below the body's upper surface at column `i`. +fn body_top(i: usize) -> Option { + let x = i as f64 * SP[0]; + let t = 1.0 - ((x - CX) / RX).powi(2); + if t < 0.0 { + return None; + } + Some(((CY - RY * t.sqrt()) / SP[1]).ceil() as usize) +} + +/// The row the mask shell occupies at column `i`: two voxels clear of the +/// skin, except over the contact patch, where it touches. +fn shell_row(i: usize) -> Option { + let top = body_top(i)?; + let off = if CONTACT.contains(&i) { 1 } else { 2 }; + top.checked_sub(off) +} + +/// The phantom, plus the body mask it was built from. +struct Phantom { + volume: Volume, + truth: Vec, +} + +fn phantom() -> Phantom { + let mut data = vec![-1000i16; NX * NY * NZ]; + let mut truth = vec![0u8; NX * NY * NZ]; + for k in 0..NZ { + for j in 0..NY { + for i in 0..NX { + let (x, y) = (i as f64 * SP[0], j as f64 * SP[1]); + if ((x - CX) / RX).powi(2) + ((y - CY) / RY).powi(2) <= 1.0 { + data[idx(i, j, k)] = 40; + truth[idx(i, j, k)] = 1; + } + } + } + } + // Ears: 6 mm plates just outside each flank, over 25 mm of the stack. + for k in 8..13 { + for j in 26..36 { + // Overlapping the flank by two voxels, so they are attached to + // the body rather than to the knife edge of a perfect ellipse. + for i in [7, 8, 9, 10, 11, 88, 89, 90, 91, 92] { + data[idx(i, j, k)] = 40; + truth[idx(i, j, k)] = 1; + } + } + } + // Lungs, a comfortable wall away from the skin. + for k in 0..NZ { + for j in 22..43 { + for i in (30..46).chain(55..71) { + data[idx(i, j, k)] = -850; + } + } + } + // Airways reaching both lungs — on the first two slices only, so that + // exactly those two slices have a cavity open to the outside. + for k in 0..2 { + for j in 0..23 { + for i in (40..44).chain(57..61) { + data[idx(i, j, k)] = -850; + truth[idx(i, j, k)] = 0; + } + } + } + // The couch: a carbon skin under the back with a 2 mm air gap, and a + // rail 6 mm below it. The foam between them is below the threshold. + for k in 0..NZ { + for i in 10..90 { + data[idx(i, 59, k)] = 300; + data[idx(i, 62, k)] = 300; + } + } + // The moulded mask shell. + for k in 0..NZ { + for i in 15..86 { + if let Some(j) = shell_row(i) { + data[idx(i, j, k)] = 120; + } + } + } + // A cable: thin, free-standing, running the whole length. + for k in 0..NZ { + data[idx(95, 65, k)] = 200; + } + Phantom { + volume: volume(data), + truth, + } +} + +fn volume(data: Vec) -> Volume { + Volume { + data, + dims: [NX, NY, NZ], + spacing: SP, + origin: Vec3::new(0.0, 0.0, 0.0), + row_dir: Vec3::new(1.0, 0.0, 0.0), + col_dir: Vec3::new(0.0, 1.0, 0.0), + normal: Vec3::new(0.0, 0.0, 1.0), + frame_of_reference_uid: "1.2.826.0.1.3680043.8.498.phantom".into(), + min_value: -1000, + max_value: 300, + } +} + +fn params() -> BodyParams { + BodyParams { + method: Method::Classical, + foreground: Foreground::Hu(-300.0), + open_mm: 8.0, + // The phantom is 200 mm long; a 100 mm window still tells the + // extruded equipment from the 25 mm ears with room to spare. + persist_window_mm: 100.0, + min_volume_cm3: 20.0, + ..BodyParams::default() + } +} + +/// No model folder is ever reached by the classical method. +fn nowhere() -> &'static std::path::Path { + std::path::Path::new("/nonexistent") +} + +fn dice(a: &[u8], b: &[u8]) -> f64 { + let inter = a + .iter() + .zip(b.iter()) + .filter(|(x, y)| **x != 0 && **y != 0) + .count(); + let na = a.iter().filter(|v| **v != 0).count(); + let nb = b.iter().filter(|v| **v != 0).count(); + 2.0 * inter as f64 / (na + nb).max(1) as f64 +} + +const VOXEL_CM3: f64 = SP[0] * SP[1] * SP[2] / 1000.0; + +#[test] +fn the_classical_method_finds_the_patient_and_leaves_the_equipment_out() { + let p = phantom(); + let r = contour_body(&p.volume, ¶ms(), nowhere(), &Progress::default()).expect("a body"); + + // The residual is the two slices whose airway is open to the outside, + // which is the behaviour the slice-wise fill is chosen for. + let d = dice(&r.mask, &p.truth); + assert!(d > 0.99, "Dice against the built body is {d:.4}"); + assert_eq!(r.pieces.len(), 1, "one body"); + assert!(r.removed_voxels > 0, "equipment was reported as removed"); + + for k in [0, NZ / 2, NZ - 1] { + for i in 15..85 { + assert_eq!(r.mask[idx(i, 59, k)], 0, "couch skin kept at i={i} k={k}"); + assert_eq!(r.mask[idx(i, 62, k)], 0, "couch rail kept at i={i} k={k}"); + } + assert_eq!(r.mask[idx(95, 65, k)], 0, "cable kept at k={k}"); + } + + // The mask shell goes wherever it stands clear of the skin. + for k in [0, NZ / 2, NZ - 1] { + for i in 15..86 { + if CONTACT.contains(&i) { + continue; + } + if let Some(j) = shell_row(i) { + assert_eq!(r.mask[idx(i, j, k)], 0, "mask shell kept at i={i} k={k}"); + } + } + } + + // The ears — thin, and exactly what a plain opening shaves off. + assert_eq!(r.mask[idx(8, 30, 10)], 1, "left ear"); + assert_eq!(r.mask[idx(91, 30, 10)], 1, "right ear"); + assert!(r.recovered_voxels > 0, "thin anatomy was recovered"); + + // The lungs are inside the body on every slice past the airway. + for k in 2..NZ { + assert_eq!(r.mask[idx(35, 30, k)], 1, "left lung filled on k={k}"); + assert_eq!(r.mask[idx(65, 30, k)], 1, "right lung filled on k={k}"); + } +} + +#[test] +fn a_shell_pressed_against_the_skin_is_kept_and_the_cost_is_bounded() { + // The documented limitation, pinned as a test rather than left to be + // rediscovered: where a shell touches with no air gap, it is locally + // indistinguishable from a slightly thicker patient, so it stays. What + // must not happen is the error growing beyond the contact patch. + let p = phantom(); + let r = contour_body(&p.volume, ¶ms(), nowhere(), &Progress::default()).expect("a body"); + let kept = CONTACT + .filter_map(shell_row) + .flat_map(|j| (0..NZ).map(move |k| (j, k))) + .zip(CONTACT.cycle()) + .filter(|((j, k), i)| r.mask[idx(*i, *j, *k)] != 0) + .count(); + assert!(kept > 0, "the contact patch is expected to survive"); + let extra = r + .mask + .iter() + .zip(p.truth.iter()) + .filter(|(m, t)| **m != 0 && **t == 0) + .count() as f64 + * VOXEL_CM3; + assert!( + extra < 10.0, + "the error beyond the patient is {extra:.1} cm³ — the contact patch \ + alone should be about 7" + ); +} + +#[test] +fn the_chest_wall_over_a_lung_is_not_mistaken_for_a_couch_skin() { + // The failure real data taught, reduced to its essentials: a hollow + // cylinder — a 6 mm wall around a big cavity — beside a 2 mm couch + // skin. A threshold does not see a body here, it sees a thin shell + // that repeats slice after slice, which is the exact signature the + // equipment test looks for. Both are thin, both are extruded; only one + // of them encloses the patient. + let (nx, ny, nz) = (70usize, 70usize, 40usize); + let at = |i: usize, j: usize, k: usize| k * nx * ny + j * nx + i; + let (cx, cy) = (70.0f64, 70.0f64); + let mut data = vec![-1000i16; nx * ny * nz]; + for k in 0..nz { + for j in 0..ny { + for i in 0..nx { + let r = ((i as f64 * SP[0] - cx).powi(2) + (j as f64 * SP[1] - cy).powi(2)).sqrt(); + if r <= 50.0 { + data[at(i, j, k)] = 40; + } + if r <= 44.0 { + data[at(i, j, k)] = -850; + } + } + } + for i in 5..65 { + data[at(i, 63, k)] = 300; + } + } + let mut vol = volume(vec![0i16; nx * ny * nz]); + vol.data = data; + vol.dims = [nx, ny, nz]; + let r = contour_body(&vol, ¶ms(), nowhere(), &Progress::default()).expect("a body"); + for k in [0, nz / 2, nz - 1] { + for deg in (0..360).step_by(30) { + let a = (deg as f64).to_radians(); + let i = ((cx + 47.0 * a.cos()) / SP[0]).round() as usize; + let j = ((cy + 47.0 * a.sin()) / SP[1]).round() as usize; + assert_eq!(r.mask[at(i, j, k)], 1, "wall deleted at {deg}° on k={k}"); + } + // The cavity is inside the body… + assert_eq!(r.mask[at(35, 35, k)], 1, "cavity not filled on k={k}"); + // …and the couch is not. + for i in 10..60 { + assert_eq!(r.mask[at(i, 63, k)], 0, "couch skin kept at i={i} k={k}"); + } + } +} + +#[test] +fn two_legs_are_two_bodies_not_the_larger_one() { + let (nx, ny, nz) = (60usize, 40usize, 20usize); + let at = |i: usize, j: usize, k: usize| k * nx * ny + j * nx + i; + let mut data = vec![-1000i16; nx * ny * nz]; + for k in 0..nz { + for j in 8..32 { + for i in (4..24).chain(34..50) { + data[at(i, j, k)] = 40; + } + } + } + let mut vol = volume(vec![0i16; nx * ny * nz]); + vol.data = data; + vol.dims = [nx, ny, nz]; + let mut p = params(); + p.remove_devices = false; + let r = contour_body(&vol, &p, nowhere(), &Progress::default()).expect("two bodies"); + assert_eq!(r.pieces.len(), 2, "both legs kept"); + assert_eq!(r.mask[at(14, 20, 10)], 1); + assert_eq!(r.mask[at(42, 20, 10)], 1); +} + +#[test] +fn an_mr_phantom_with_a_coil_gradient_still_comes_out_whole() { + // The same body, but with MR intensities and a receive profile that + // falls off across the field — the case a fixed threshold cannot serve. + let p = phantom(); + let mut data = vec![0i16; NX * NY * NZ]; + for k in 0..NZ { + for j in 0..NY { + for i in 0..NX { + let gain = (-1.6 * i as f32 / NX as f32).exp(); + let raw = if p.truth[idx(i, j, k)] == 0 { + 5.0 + } else if p.volume.data[idx(i, j, k)] < -300 { + 20.0 // lung: dark on MR too + } else { + 900.0 + }; + data[idx(i, j, k)] = (raw * gain) as i16; + } + } + } + let mut pr = params(); + pr.foreground = Foreground::MrRelative { + fraction: 0.12, + // Well above any anatomy, well below the 160 mm body — the window a + // bias estimate has to sit in. + sigma_mm: 25.0, + }; + pr.remove_devices = false; + let r = contour_body(&volume(data), &pr, nowhere(), &Progress::default()).expect("a body"); + let d = dice(&r.mask, &p.truth); + assert!(d > 0.97, "Dice on the MR phantom is {d:.4}"); + assert_eq!(r.mask[idx(12, 31, 20)], 1, "the bright end of the field"); + assert_eq!(r.mask[idx(87, 31, 20)], 1, "the dim end of the field"); +} + +#[test] +fn a_model_folder_that_cannot_exist_is_an_error_not_a_panic() { + // A regular file where the model folder should be: the download cannot + // even begin, whatever the network is doing, which is what makes this + // deterministic on a build machine that happens to be online. + let file = std::env::temp_dir().join("rds-body-not-a-folder"); + std::fs::write(&file, b"not a folder").expect("write the blocker"); + let p = phantom(); + let mut pr = params(); + pr.method = Method::ModelAssisted; + let err = contour_body(&p.volume, &pr, &file, &Progress::default()) + .expect_err("no weights, no answer"); + let text = format!("{err:#}"); + assert!( + text.to_lowercase().contains("body") + || text.contains("create") + || text.contains("os error"), + "the error should name what failed: {text}" + ); + let _ = std::fs::remove_file(&file); +} + +/// The real thing, end to end, against the published weights. Ignored by +/// default because it downloads 124 MB on first run: +/// +/// ```text +/// RDS_BODY_MODELS=path/to/models/totalsegmentator \ +/// cargo test --release --test body -- --ignored +/// ``` +#[test] +#[ignore] +fn the_model_assisted_method_runs_the_published_network() { + let dir = std::env::var("RDS_BODY_MODELS").expect("set RDS_BODY_MODELS"); + let p = phantom(); + let mut pr = params(); + pr.method = Method::ModelAssisted; + let r = contour_body( + &p.volume, + &pr, + std::path::Path::new(&dir), + &Progress::default(), + ) + .expect("a body"); + // The network sees a featureless ellipse rather than a person, so this + // asserts that the hybrid holds together — a body of a plausible size, + // with the skin still placed by the threshold — not a Dice figure that + // would only be meaningful on real anatomy. + assert!(r.voxels > 0, "the network found a patient"); + assert!(!r.device.is_empty(), "the device was reported"); + let d = dice(&r.mask, &p.truth); + assert!(d > 0.9, "Dice {d:.4}"); +} From 80d0b7c151da03f49e616f8a8c5337b93c7d8911 Mon Sep 17 00:00:00 2001 From: alexprotom Date: Thu, 27 Aug 2026 13:23:30 +0200 Subject: [PATCH 3/5] GUI fixes for multi line patient names + option to hide left panel and new Modules menu --- Cargo.lock | 2 +- docs/architecture.md | 5 +- docs/example-data.md | 4 +- docs/export-and-tools.md | 4 +- docs/registration.md | 6 +- docs/rt-objects.md | 2 +- docs/viewer.md | 23 +- src/app/chrome.rs | 135 +++++---- src/app/mod.rs | 27 +- src/app/panels.rs | 616 +++++++++++++++++++++++---------------- src/app/propagate_win.rs | 3 + src/app/reg_panel.rs | 16 +- src/settings.rs | 70 ++++- 13 files changed, 576 insertions(+), 337 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 1ce11f6..0b45d0c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6413,7 +6413,7 @@ dependencies = [ [[package]] name = "rust-dicom-station" -version = "0.5.0" +version = "0.1.0" dependencies = [ "anyhow", "burn", diff --git a/docs/architecture.md b/docs/architecture.md index af0a366..9e9dfb4 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -187,7 +187,8 @@ src/ plumbing (Job::spawn, poll_job, poll_tool_job), per-frame driver theme.rs theme-dependent colors chrome.rs menu bar, toolbar, status bar, help - panels.rs side panel and its per-dataset sections + panels.rs left panel: its show / hide, the optional modules and + the per-dataset Data tree sections views.rs central MPR viewports, interaction, texture caches d3.rs live 3D structure window planar.rs floating DX / CR / RTIMAGE viewers @@ -202,7 +203,7 @@ src/ auto-segmentation job starts dialogs.rs auto-segmentation window + results, generator, anonymizer, export, error dialog - reg_panel.rs the Registration section: method, region, parameters, + reg_panel.rs the Image registration section: method, region, parameters, landmarks, the run, the analytics, the vector field models_win.rs the model manager window propagate_win.rs structure propagation window and worker diff --git a/docs/example-data.md b/docs/example-data.md index 7d38561..49e7bf9 100644 --- a/docs/example-data.md +++ b/docs/example-data.md @@ -23,8 +23,8 @@ cargo run --release -- example_data/lung_p1_4DCT_phase_000 example_data/lung_p1_ That is a ready-made comparison-mode and registration test case with real respiratory motion: the tumor and the markers move visibly between the -phases, and *Registration ▶ Deformable* has something anatomically real to -recover. Equivalently, load the whole `example_data/` folder as dataset A +phases, and the deformable methods of the *Image registration* module +have something anatomically real to recover. Equivalently, load the whole `example_data/` folder as dataset A (both phases appear as two series of one study) and right-click one phase ▶ *Copy series to dataset B*. It is also the dataset the auto-segmentation was validated on ([auto-segmentation.md](auto-segmentation.md#validation)). diff --git a/docs/export-and-tools.md b/docs/export-and-tools.md index e945135..b724a66 100644 --- a/docs/export-and-tools.md +++ b/docs/export-and-tools.md @@ -149,7 +149,7 @@ matter of generating once more into another folder: cargo run --release -- test_data test_data_shifted ``` -*Registration ▶ Rigid* should then recover the (12, −9, 0) mm shift to -within a fraction of a millimeter. The whole phantom is analytically +a rigid run in the *Image registration* module should then recover the +(12, −9, 0) mm shift to within a fraction of a millimeter. The whole phantom is analytically known, which is what the integration tests assert against — see [architecture.md](architecture.md#testing). diff --git a/docs/registration.md b/docs/registration.md index d99956f..76686f9 100644 --- a/docs/registration.md +++ b/docs/registration.md @@ -124,10 +124,10 @@ Each pair shows its displacement, and after a run, its residual. ## Running a registration -With two datasets loaded, the *Registration* menu or the sidebar section -registers one onto the other — the direction is selectable (**B ▶ A** or +*Modules ▶ Image registration* puts the section in the left panel. With +two datasets loaded it registers one onto the other — the direction is selectable (**B ▶ A** or **A ▶ B**; the second-named dataset is the fixed image and receives the -fusion overlay). Everything else lives in the sidebar: method, region, +fusion overlay). Everything lives in that one section: method, region, parameters, landmarks, the result and the vector field. Runs happen on a background thread with progress and a **Cancel** button. diff --git a/docs/rt-objects.md b/docs/rt-objects.md index eb303b0..fed1b2c 100644 --- a/docs/rt-objects.md +++ b/docs/rt-objects.md @@ -104,7 +104,7 @@ which loaded dataset the grid's own frame of reference matches, so applying it the wrong way round is a deliberate act rather than an accident. A registration recovered here can be written back out as a Deformable -Spatial Registration (*Registration ▶ Vector field ▶ 💾 Save as DICOM…*). +Spatial Registration (*Image registration ▶ Vector field ▶ 💾 Save as DICOM…*). The IOD applies its grid after a pre-deformation matrix and before a post-deformation one; both are written as the identity and the grid carries the whole mapping, so another system has no composition rule to get wrong. diff --git a/docs/viewer.md b/docs/viewer.md index 65e560a..9a1cdc8 100644 --- a/docs/viewer.md +++ b/docs/viewer.md @@ -65,13 +65,25 @@ scrolling), the **3D A / 3D B** buttons and the segmentation tools. (x = width, y = center); the toolbar offers the numeric fields and the common CT presets: brain, subdural, stroke, head/neck soft tissue, temporal bone, lungs, mediastinum, abdomen, liver, spine, bone, CT angio, full -range. Window/level is shared between datasets A and B so both CTs are -windowed identically. +range. The list shows each preset's center and width; once one is chosen +the closed list carries its name alone (*Lungs*, *Liver*, …) — the two +numeric fields next to it already say what the numbers are. Any other +window, from a drag or the full range, leaves the list nameless again. +Window/level is shared between datasets A and B so both CTs are windowed +identically. **Status bar.** Patient coordinates, voxel indices, HU and dose (Gy and % of the reference dose) at the crosshair; in comparison mode the readouts for A and B are shown side by side. +**The left panel.** *View ▶ Left panel*, the **F9** key and the arrow on +the window's left edge hide and show it; dragging its inner edge past the +minimum width does the same, and the arrow stays on screen to bring it +back. Which sections it holds is up to the *Modules* menu: **Image +registration** and **Image simulation** are off until switched on, and the +choice is remembered between runs. The data tree of each loaded dataset is +always there. + ## Interaction reference | Input | Action | @@ -95,10 +107,13 @@ accumulated from any number of folders. *File ▶ Add DICOM folder to A/B…* merges a scanned folder into the slot without unloading what is already there; duplicates (by UID) are skipped and reported. -The sidebar shows each dataset as a full DICOM hierarchy — patient +The left panel shows each dataset as a **Data tree** — a full DICOM +hierarchy: patient (PatientName/PatientID) ▶ study (StudyInstanceUID, with date and description) ▶ image series — with the displayed series marked; clicking -another series loads it. The standard reference chain is parsed and shown +another series loads it. Long names, descriptions and IDs wrap over as +many lines as they need, so the panel can be dragged narrow without +cutting them off; only the section headers stay on one line. The standard reference chain is parsed and shown as links: each structure set displays the image series its contours were drawn on (RTReferencedSeriesSequence), each dose the plan it was computed for (ReferencedRTPlanSequence), and each plan the structure set it was diff --git a/src/app/chrome.rs b/src/app/chrome.rs index 19db3dc..cb05b56 100644 --- a/src/app/chrome.rs +++ b/src/app/chrome.rs @@ -9,13 +9,14 @@ impl ViewerApp { let mut open_b = false; let mut close_b = false; let mut reset_views = false; - let mut do_reg: Option<(RegMethod, bool)> = None; let mut open_gen = false; let mut open_models = false; let mut open_propagate = false; let mut open_drr = false; let mut open_export: Option = None; let mut new_theme: Option = None; + // A module was switched on or off — remember it for the next run. + let mut modules_changed = false; egui::Panel::top(egui::Id::new("menu_bar")).show(ui, |ui| { egui::MenuBar::new().ui(ui, |ui| { @@ -113,6 +114,13 @@ impl ViewerApp { ui.checkbox(&mut self.show_labels, "Orientation labels"); ui.checkbox(&mut self.show_isocenters, "Isocenters"); ui.separator(); + ui.checkbox(&mut self.side_open, "Left panel (F9)") + .on_hover_text( + "Hide the left panel and give the whole window to the views. \ + The arrow on the window's left edge brings it back, as does \ + F9.", + ); + ui.separator(); ui.label("Appearance:"); let before = self.theme; self.theme.radio_buttons(ui); @@ -128,56 +136,38 @@ impl ViewerApp { ui.close(); } }); - ui.menu_button("Registration", |ui| { - // Quick actions only — direction, region, parameters, - // landmarks, analytics, fusion and the vector field live - // in the sidebar Registration section. - let both = self.slots[0].study.is_some() && self.slots[1].study.is_some(); - let running = self.reg_job.is_some(); - let moving = SLOT_NAMES[1 - self.reg_fixed_slot.min(1)]; - let fixed = SLOT_NAMES[self.reg_fixed_slot.min(1)]; - ui.weak(format!("Register {moving} onto {fixed}:")); - for method in RegMethod::ALL { - if ui - .add_enabled( - both && !running, - egui::Button::new(format!( - "{} — {}", - method.family(), - method.short() - )), + // This menu is a set of switches, not a list of actions: + // it stays open until the pointer leaves it, so both + // modules can be turned on in one visit. + egui::containers::menu::MenuButton::new("Modules") + .config( + egui::containers::menu::MenuConfig::new() + .close_behavior(egui::PopupCloseBehavior::CloseOnClickOutside), + ) + .ui(ui, |ui| { + // Two optional side-panel sections, each one line of the + // menu. Everything they do — direction, method, region, + // parameters, landmarks, analytics, fusion, the vector + // field, the simulated motion — lives in the section + // itself, so the menu only decides whether it is there. + ui.weak("Sections of the left panel:"); + modules_changed |= ui + .checkbox(&mut self.module_registration, "Image registration") + .on_hover_text( + "Align two datasets: direction, method, region, parameters, \ + landmarks, analysis, fusion and the deformation vector field. \ + Needs two loaded datasets to run.", ) - .on_hover_text(method.hint()) - .clicked() - { - do_reg = Some((method, false)); - ui.close(); - } - } - ui.separator(); - let can_refine = self - .registration - .as_ref() - .is_some_and(|r| r.fixed_slot == self.reg_fixed_slot.min(1)); - if ui - .add_enabled( - both && !running && can_refine, - egui::Button::new("Refine the active registration"), - ) - .on_hover_text( - "Recover a correction on top of the active result with the \ - method and region chosen in the sidebar, and add the two \ - together", - ) - .clicked() - { - do_reg = Some((self.reg_method, true)); - ui.close(); - } - if !both { - ui.weak("Load two datasets (comparison mode) first"); - } - }); + .changed(); + modules_changed |= ui + .checkbox(&mut self.module_simulation, "Image simulation") + .on_hover_text( + "Registration QA: apply a known rigid motion and Gaussian \ + deformation to one dataset and generate the result into the \ + other — the ground truth a registration can be measured against.", + ) + .changed(); + }); ui.menu_button("Tools", |ui| { // The three segmentation engines, one block per dataset: // the same four entries, in the same order, for A and B. @@ -353,10 +343,6 @@ impl ViewerApp { if reset_views { self.reset_all_views(); } - if let Some((method, refine)) = do_reg { - self.reg_method = method; - self.start_registration(refine); - } if open_gen { self.gen_open = true; } @@ -381,6 +367,9 @@ impl ViewerApp { if let Some(theme) = new_theme { self.set_theme(ctx, theme); } + if modules_changed { + self.persist_settings(); + } } // -- Toolbar ---------------------------------------------------------- @@ -407,17 +396,36 @@ impl ViewerApp { .prefix("W "), ); let mut full_range = false; + // The closed combo carries the name of the preset in + // force. Its numbers are dropped there — the two drag + // values to the left already show them — but kept in the + // list, where they are what tells the presets apart. + // Any other window (a drag, a right-drag in a view, the + // full range) is nameless again. + self.wl_preset = self.wl_preset.filter(|i| { + WL_PRESETS.get(*i).is_some_and(|(_, c, w)| { + *c == self.window_center && *w == self.window_width + }) + }); + let selected = self + .wl_preset + .and_then(|i| WL_PRESETS.get(i)) + .map_or("CT presets", |(name, ..)| *name); + let mut pick: Option = None; + let current = self.wl_preset; egui::ComboBox::from_id_salt("wl_preset") - .selected_text("CT presets") - .width(110.0) + .selected_text(selected) + .width(150.0) .show_ui(ui, |ui| { - for (name, c, w) in WL_PRESETS { + for (i, (name, c, w)) in WL_PRESETS.iter().enumerate() { if ui - .button(format!("{name} (C {c:.0} / W {w:.0})")) + .selectable_label( + current == Some(i), + format!("{name} (C {c:.0} / W {w:.0})"), + ) .clicked() { - self.window_center = *c; - self.window_width = *w; + pick = Some(i); } } ui.separator(); @@ -425,7 +433,14 @@ impl ViewerApp { full_range = true; } }); + if let Some(i) = pick { + let (_, c, w) = WL_PRESETS[i]; + self.window_center = c; + self.window_width = w; + self.wl_preset = Some(i); + } if full_range { + self.wl_preset = None; if let Some(study) = &self.slots[self.hovered_slot.min(1)].study { let v = &study.volume; self.window_center = (v.min_value as f32 + v.max_value as f32) * 0.5; diff --git a/src/app/mod.rs b/src/app/mod.rs index 4936997..6391b08 100644 --- a/src/app/mod.rs +++ b/src/app/mod.rs @@ -885,6 +885,21 @@ pub struct ViewerApp { /// Bumped whenever ROI visibility / dose settings change → cache rebuild. settings_gen: u64, + /// The CT window preset last picked from the toolbar list (an index into + /// [`WL_PRESETS`]), so the closed combo can name it instead of reading + /// "CT presets". Dropped as soon as the window no longer matches it. + wl_preset: Option, + + /// *Modules ▶ Image registration*: the registration section is part of + /// the side panel. Persisted between runs. + module_registration: bool, + /// *Modules ▶ Image simulation*: the simulation section is part of the + /// side panel. Persisted between runs. + module_simulation: bool, + /// The side panel is expanded (View ▶ Left panel, F9, or the arrow on + /// the panel edge). Collapsed, the views have the whole window. + side_open: bool, + /// Light / dark / follow-the-system appearance, persisted between runs. theme: egui::ThemePreference, /// Non-fatal note shown in the View menu if the settings file could not @@ -1038,6 +1053,10 @@ impl ViewerApp { dose_threshold_pct: 15.0, iso_levels: default_iso_levels(), settings_gen: 0, + wl_preset: None, + module_registration: prefs.module_registration, + module_simulation: prefs.module_simulation, + side_open: true, theme: prefs.theme, settings_error: None, }; @@ -1069,6 +1088,8 @@ impl ViewerApp { match settings::save(&Settings { theme: self.theme, models_dir, + module_registration: self.module_registration, + module_simulation: self.module_simulation, }) { Ok(()) => self.settings_error = None, Err(e) => { @@ -1266,14 +1287,18 @@ impl eframe::App for ViewerApp { // focused): Ctrl+Z undo, Esc cancels a region-grow drag, [ ] resize // the brush. if !ctx.egui_wants_keyboard_input() { - let (undo, esc, smaller, bigger) = ctx.input(|i| { + let (undo, esc, smaller, bigger, toggle_side) = ctx.input(|i| { ( i.modifiers.command && i.key_pressed(egui::Key::Z), i.key_pressed(egui::Key::Escape), i.key_pressed(egui::Key::OpenBracket), i.key_pressed(egui::Key::CloseBracket), + i.key_pressed(egui::Key::F9), ) }); + if toggle_side { + self.side_open = !self.side_open; + } if undo { let slot = self.hovered_slot.min(1); self.undo_active_seg(slot); diff --git a/src/app/panels.rs b/src/app/panels.rs index 623461d..f5a60b9 100644 --- a/src/app/panels.rs +++ b/src/app/panels.rs @@ -12,13 +12,55 @@ impl ViewerApp { if self.slots[0].study.is_none() && self.slots[1].study.is_none() { return; } + // A thin strip along the window edge carries the show / hide arrow. + // It stays there when the panel is gone — otherwise nothing on + // screen would say how to get it back (View ▶ Left panel and F9 do + // the same). + let strip = egui::Frame::new() + .fill(ui.visuals().panel_fill) + .inner_margin(egui::Margin::symmetric(1, 4)); + egui::Panel::left(egui::Id::new("side_toggle")) + .exact_size(22.0) + .resizable(false) + .frame(strip) + .show(ui, |ui| { + let (glyph, hint) = if self.side_open { + ( + "◀", + "Hide the left panel — the views take the whole window (F9)", + ) + } else { + ("▶", "Show the left panel (F9)") + }; + if ui + .add( + egui::Button::new(glyph) + .frame(false) + .min_size(egui::vec2(20.0, 22.0)), + ) + .on_hover_text(hint) + .clicked() + { + self.side_open = !self.side_open; + } + }); + // `show_collapsible` also lets the panel be dragged shut and pulled + // back open by its edge; it wants the flag by reference, which the + // body's `&mut self` cannot share, so it travels via a local. + let mut open = self.side_open; egui::Panel::left(egui::Id::new("side")) .resizable(true) .default_size(280.0) - .show(ui, |ui| { + .show_collapsible(ui, &mut open, |ui| { egui::ScrollArea::vertical().show(ui, |ui| { - self.registration_section(ui); - self.simulation_section(ui); + // The two optional modules (Modules menu) come first, + // then one data tree per loaded dataset. + if self.module_registration { + self.registration_section(ui); + } + if self.module_simulation { + self.simulation_section(ui); + } for slot in 0..2 { if self.slots[slot].study.is_none() { continue; @@ -27,6 +69,55 @@ impl ViewerApp { } }); }); + self.side_open = open; + } + + /// A tree node whose title wraps over as many lines as it needs. + /// + /// [`egui::CollapsingHeader`] lays its title out on a single line and + /// lets it run past the panel edge, so a long patient name, a long + /// study description or a long ID would pin the panel open at that + /// width. Only the module headers keep that one-line behaviour. + /// + /// Returns the title's own response, so the caller can hang the node's + /// context menu on it. + fn wrapped_node( + ui: &mut egui::Ui, + id_salt: impl std::hash::Hash + std::fmt::Debug, + default_open: bool, + title: impl Into, + body: impl FnOnce(&mut egui::Ui) -> R, + ) -> egui::Response { + let id = ui.make_persistent_id(id_salt); + let state = egui::collapsing_header::CollapsingState::load_with_default_open( + ui.ctx(), + id, + default_open, + ); + let (_, header, _) = state + .show_header(ui, |ui| { + ui.add( + egui::Label::new( + egui::RichText::new(title.into()).text_style(egui::TextStyle::Button), + ) + .wrap() + .sense(egui::Sense::click()), + ) + }) + .body(body); + let resp = header.inner; + // Clicking the title opens and closes the node, as it does on a + // standard header; the arrow keeps working on its own. + if resp.clicked() { + let mut state = egui::collapsing_header::CollapsingState::load_with_default_open( + ui.ctx(), + id, + default_open, + ); + state.toggle(ui); + state.store(ui.ctx()); + } + resp } /// Study transform simulator: apply a known rigid motion + optional @@ -37,7 +128,7 @@ impl ViewerApp { return; } let mut do_generate = false; - egui::CollapsingHeader::new(egui::RichText::new("Simulation (registration QA)").strong()) + egui::CollapsingHeader::new(egui::RichText::new("Image simulation").strong()) .default_open(false) .show(ui, |ui| { if let Some(job) = &self.sim_job { @@ -114,7 +205,7 @@ impl ViewerApp { pub(super) fn study_section(&mut self, ui: &mut egui::Ui, slot: usize) { // Plain header — the patient(s) always appear as tree nodes below. - let header = format!("Dataset {}", SLOT_NAMES[slot]); + let header = format!("Data tree {}", SLOT_NAMES[slot]); let ch = egui::CollapsingHeader::new(egui::RichText::new(header).strong()) .id_salt(("study_hdr", slot)) .default_open(true) @@ -200,139 +291,139 @@ impl ViewerApp { } else { format!("{} ({})", pname, pinfo.patient_id) }; - let pch = egui::CollapsingHeader::new(ptitle) - .id_salt(("pat_hdr", slot, pi)) - .default_open(true) - .show(ui, |ui| { - // Studies of this patient, in first-seen order. - let mut studies: Vec<&str> = Vec::new(); - for s in &study.series { - if s.patient_key() == *pkey && !studies.contains(&s.study_uid.as_str()) - { - studies.push(&s.study_uid); - } + let pch = Self::wrapped_node(ui, ("pat_hdr", slot, pi), true, ptitle, |ui| { + // Studies of this patient, in first-seen order. + let mut studies: Vec<&str> = Vec::new(); + for s in &study.series { + if s.patient_key() == *pkey && !studies.contains(&s.study_uid.as_str()) { + studies.push(&s.study_uid); } - for (si, study_uid) in studies.iter().enumerate() { - let info = study - .series - .iter() - .find(|s| s.study_uid == *study_uid && s.patient_key() == *pkey) - .unwrap(); - let title = format!( - "Study {}{}", - if info.study_date.is_empty() { - format!("{}", si + 1) - } else { - info.study_date.clone() - }, - if info.study_description.is_empty() { - String::new() - } else { - format!(" — {}", info.study_description) - } - ); - let sch = egui::CollapsingHeader::new(title) - .id_salt(("study_tree", slot, pi, si)) - .default_open(true) - .show(ui, |ui| { - for (i, s) in study.series.iter().enumerate() { - if s.study_uid != *study_uid || s.patient_key() != *pkey { - continue; + } + for (si, study_uid) in studies.iter().enumerate() { + let info = study + .series + .iter() + .find(|s| s.study_uid == *study_uid && s.patient_key() == *pkey) + .unwrap(); + let title = format!( + "Study {}{}", + if info.study_date.is_empty() { + format!("{}", si + 1) + } else { + info.study_date.clone() + }, + if info.study_description.is_empty() { + String::new() + } else { + format!(" — {}", info.study_description) + } + ); + let sch = Self::wrapped_node( + ui, + ("study_tree", slot, pi, si), + true, + title, + |ui| { + for (i, s) in study.series.iter().enumerate() { + if s.study_uid != *study_uid || s.patient_key() != *pkey { + continue; + } + let resp = ui.add( + egui::Button::selectable(i == active, label(s)).wrap(), + ); + if resp.clicked() && i != active { + switch_to = Some(i); + } + resp.context_menu(|ui| { + if ui.button("✎ Rename series…").clicked() { + rename = Some(RenameTarget::Series { slot, idx: i }); + ui.close(); } - let resp = ui.selectable_label(i == active, label(s)); - if resp.clicked() && i != active { - switch_to = Some(i); + ui.separator(); + if ui + .button(format!("Copy series to dataset {other}")) + .clicked() + { + act_series = Some(TreeAction { + from: slot, + sel: TreeSel::Series(i), + op: TreeOp::Copy, + }); + ui.close(); + } + if ui + .button(format!("Move series to dataset {other}")) + .clicked() + { + act_series = Some(TreeAction { + from: slot, + sel: TreeSel::Series(i), + op: TreeOp::Move, + }); + ui.close(); + } + ui.separator(); + if ui.button("Remove series").clicked() { + act_series = Some(TreeAction { + from: slot, + sel: TreeSel::Series(i), + op: TreeOp::Remove, + }); + ui.close(); } - resp.context_menu(|ui| { - if ui.button("✎ Rename series…").clicked() { - rename = - Some(RenameTarget::Series { slot, idx: i }); - ui.close(); - } - ui.separator(); - if ui - .button(format!("Copy series to dataset {other}")) - .clicked() - { - act_series = Some(TreeAction { - from: slot, - sel: TreeSel::Series(i), - op: TreeOp::Copy, - }); - ui.close(); - } - if ui - .button(format!("Move series to dataset {other}")) - .clicked() - { - act_series = Some(TreeAction { - from: slot, - sel: TreeSel::Series(i), - op: TreeOp::Move, - }); - ui.close(); - } - ui.separator(); - if ui.button("Remove series").clicked() { - act_series = Some(TreeAction { - from: slot, - sel: TreeSel::Series(i), - op: TreeOp::Remove, - }); - ui.close(); - } - }); - resp.on_hover_text(format!( - "Series UID …{}\nright-click: rename, copy / move \ - to dataset {other}, or remove", - tail(&s.uid) - )); - } - }); - sch.header_response.context_menu(|ui| { - if ui.button("✎ Rename study…").clicked() { - rename = Some(RenameTarget::Study { - slot, - uid: study_uid.to_string(), - }); - ui.close(); - } - ui.separator(); - if ui - .button(format!("Copy study to dataset {other}")) - .clicked() - { - act_study = Some(TreeAction { - from: slot, - sel: TreeSel::Study(study_uid.to_string()), - op: TreeOp::Copy, - }); - ui.close(); - } - if ui - .button(format!("Move study to dataset {other}")) - .clicked() - { - act_study = Some(TreeAction { - from: slot, - sel: TreeSel::Study(study_uid.to_string()), - op: TreeOp::Move, - }); - ui.close(); - } - ui.separator(); - if ui.button("Remove study").clicked() { - act_study = Some(TreeAction { - from: slot, - sel: TreeSel::Study(study_uid.to_string()), - op: TreeOp::Remove, }); - ui.close(); + resp.on_hover_text(format!( + "Series UID …{}\nright-click: rename, copy / move \ + to dataset {other}, or remove", + tail(&s.uid) + )); } - }); - } - }); - pch.header_response.context_menu(|ui| { + }, + ); + sch.context_menu(|ui| { + if ui.button("✎ Rename study…").clicked() { + rename = Some(RenameTarget::Study { + slot, + uid: study_uid.to_string(), + }); + ui.close(); + } + ui.separator(); + if ui + .button(format!("Copy study to dataset {other}")) + .clicked() + { + act_study = Some(TreeAction { + from: slot, + sel: TreeSel::Study(study_uid.to_string()), + op: TreeOp::Copy, + }); + ui.close(); + } + if ui + .button(format!("Move study to dataset {other}")) + .clicked() + { + act_study = Some(TreeAction { + from: slot, + sel: TreeSel::Study(study_uid.to_string()), + op: TreeOp::Move, + }); + ui.close(); + } + ui.separator(); + if ui.button("Remove study").clicked() { + act_study = Some(TreeAction { + from: slot, + sel: TreeSel::Study(study_uid.to_string()), + op: TreeOp::Remove, + }); + ui.close(); + } + }); + } + }); + pch.context_menu(|ui| { if ui.button("✎ Rename patient…").clicked() { rename = Some(RenameTarget::Patient { slot, @@ -684,13 +775,16 @@ impl ViewerApp { } else { &set.label }; - let resp = ui.selectable_label( - i == active_set, - format!( - "▣ {name} ({} ROIs){}", - set.rois.len(), - Self::series_suffix(study, &set.referenced_series_uid) - ), + let resp = ui.add( + egui::Button::selectable( + i == active_set, + format!( + "▣ {name} ({} ROIs){}", + set.rois.len(), + Self::series_suffix(study, &set.referenced_series_uid) + ), + ) + .wrap(), ); if resp.clicked() && i != active_set { new_active = Some(i); @@ -902,14 +996,17 @@ impl ViewerApp { kind: SetKind::Segmentations, idx: i, }; - let resp = ui.selectable_label( - Some(i) == active_series, - format!( - "✎ {} ({} segments){}", - sr.label, - sr.segs.len(), - Self::series_suffix(study, &sr.referenced_series_uid) - ), + let resp = ui.add( + egui::Button::selectable( + Some(i) == active_series, + format!( + "✎ {} ({} segments){}", + sr.label, + sr.segs.len(), + Self::series_suffix(study, &sr.referenced_series_uid) + ), + ) + .wrap(), ); if resp.clicked() && Some(i) != active_series { new_active_series = Some(i); @@ -992,9 +1089,11 @@ impl ViewerApp { ui.color_edit_button_srgb(&mut row.1); ui.checkbox(&mut row.2, "") .on_hover_text("Show / select this segmentation"); - let resp = ui.selectable_label(i == active_seg, &name).on_hover_text( - "Click to make this the segmentation the tools edit", - ); + let resp = ui + .add(egui::Button::selectable(i == active_seg, name.clone()).wrap()) + .on_hover_text( + "Click to make this the segmentation the tools edit", + ); if resp.clicked() { activate = Some(i); } @@ -1254,114 +1353,117 @@ impl ViewerApp { return; } for (pi, plan) in study.plans.iter().enumerate() { - let plan_hdr = egui::CollapsingHeader::new(format!( - "Plan: {}", - if plan.label.is_empty() { - "unnamed" - } else { - &plan.label - } - )) - .id_salt(("plan", slot, pi)) - .default_open(pi == 0) - .show(ui, |ui| { - if !plan.name.is_empty() && plan.name != plan.label { - ui.weak(format!("Name: {}", plan.name)); - } - if !plan.plan_kind.is_empty() { - ui.weak(format!("Type: {}", plan.plan_kind)); - } - if let Some(fx) = plan.n_fractions { - ui.weak(format!("Fractions: {fx}")); - } - if let Some(rx) = plan.target_prescription_dose { - ui.weak(format!("Prescription: {rx:.2} Gy")); - } - if !plan.date.is_empty() { - ui.weak(format!("Date: {}", plan.date)); - } - // DICOM cross-reference: the structure set the plan was - // created on. - if !plan.referenced_structset_uid.is_empty() { - if let Some(ss) = study - .structure_sets - .iter() - .find(|s| s.sop_instance_uid == plan.referenced_structset_uid) - { - ui.weak(format!( - "▶ structures {}", - if ss.label.is_empty() { - &ss.file_name - } else { - &ss.label - } - )); + let plan_hdr = Self::wrapped_node( + ui, + ("plan", slot, pi), + pi == 0, + format!( + "Plan: {}", + if plan.label.is_empty() { + "unnamed" + } else { + &plan.label } - } - if !plan.beams.is_empty() { - egui::Grid::new(("beam_grid", slot, pi)) - .striped(true) - .min_col_width(10.0) - .show(ui, |ui| { - ui.strong("Beam"); - ui.strong("Type"); - ui.strong("G°"); - ui.strong("C°"); - ui.strong("E (MeV)"); - ui.strong("MU"); - ui.strong("CPs"); - ui.end_row(); - for b in &plan.beams { - ui.label(&b.name).on_hover_text(format!( - "Beam {} · {} · dose/fx {}", - b.number, - if b.delivery_type.is_empty() { - "TREATMENT" - } else { - &b.delivery_type - }, - b.beam_dose - .map(|d| format!("{d:.2} Gy")) - .unwrap_or_else(|| "n/a".into()), - )); - ui.label(format!( - "{}{}", - b.radiation_type, - if b.scan_mode.is_empty() { - String::new() - } else { - format!("/{}", b.scan_mode) - } - )); - ui.label( - b.gantry_angle - .map(|g| format!("{g:.0}")) - .unwrap_or_else(|| "–".into()), - ); - ui.label( - b.couch_angle - .map(|c| format!("{c:.0}")) - .unwrap_or_else(|| "–".into()), - ); - ui.label(match (b.energy_min, b.energy_max) { - (Some(a), Some(bb)) if (a - bb).abs() > 0.01 => { - format!("{a:.0}–{bb:.0}") - } - (Some(a), _) => format!("{a:.0}"), - _ => "–".into(), - }); - ui.label( - b.meterset - .map(|m| format!("{m:.1}")) - .unwrap_or_else(|| "–".into()), - ); - ui.label(format!("{}", b.n_control_points)); + ), + |ui| { + if !plan.name.is_empty() && plan.name != plan.label { + ui.weak(format!("Name: {}", plan.name)); + } + if !plan.plan_kind.is_empty() { + ui.weak(format!("Type: {}", plan.plan_kind)); + } + if let Some(fx) = plan.n_fractions { + ui.weak(format!("Fractions: {fx}")); + } + if let Some(rx) = plan.target_prescription_dose { + ui.weak(format!("Prescription: {rx:.2} Gy")); + } + if !plan.date.is_empty() { + ui.weak(format!("Date: {}", plan.date)); + } + // DICOM cross-reference: the structure set the plan was + // created on. + if !plan.referenced_structset_uid.is_empty() { + if let Some(ss) = study + .structure_sets + .iter() + .find(|s| s.sop_instance_uid == plan.referenced_structset_uid) + { + ui.weak(format!( + "▶ structures {}", + if ss.label.is_empty() { + &ss.file_name + } else { + &ss.label + } + )); + } + } + if !plan.beams.is_empty() { + egui::Grid::new(("beam_grid", slot, pi)) + .striped(true) + .min_col_width(10.0) + .show(ui, |ui| { + ui.strong("Beam"); + ui.strong("Type"); + ui.strong("G°"); + ui.strong("C°"); + ui.strong("E (MeV)"); + ui.strong("MU"); + ui.strong("CPs"); ui.end_row(); - } - }); - } - }); - plan_hdr.header_response.context_menu(|ui| { + for b in &plan.beams { + ui.label(&b.name).on_hover_text(format!( + "Beam {} · {} · dose/fx {}", + b.number, + if b.delivery_type.is_empty() { + "TREATMENT" + } else { + &b.delivery_type + }, + b.beam_dose + .map(|d| format!("{d:.2} Gy")) + .unwrap_or_else(|| "n/a".into()), + )); + ui.label(format!( + "{}{}", + b.radiation_type, + if b.scan_mode.is_empty() { + String::new() + } else { + format!("/{}", b.scan_mode) + } + )); + ui.label( + b.gantry_angle + .map(|g| format!("{g:.0}")) + .unwrap_or_else(|| "–".into()), + ); + ui.label( + b.couch_angle + .map(|c| format!("{c:.0}")) + .unwrap_or_else(|| "–".into()), + ); + ui.label(match (b.energy_min, b.energy_max) { + (Some(a), Some(bb)) if (a - bb).abs() > 0.01 => { + format!("{a:.0}–{bb:.0}") + } + (Some(a), _) => format!("{a:.0}"), + _ => "–".into(), + }); + ui.label( + b.meterset + .map(|m| format!("{m:.1}")) + .unwrap_or_else(|| "–".into()), + ); + ui.label(format!("{}", b.n_control_points)); + ui.end_row(); + } + }); + } + }, + ); + plan_hdr.context_menu(|ui| { if ui.button("✎ Rename plan…").clicked() { rename = Some(RenameTarget::Plan { slot, idx: pi }); ui.close(); diff --git a/src/app/propagate_win.rs b/src/app/propagate_win.rs index 552ebf2..6fc77ca 100644 --- a/src/app/propagate_win.rs +++ b/src/app/propagate_win.rs @@ -227,6 +227,9 @@ impl ViewerApp { region: refined.region, }); self.reg_gen += 1; + // The refinement is now the active registration — show the + // section that reports and clears it. + self.module_registration = true; } let Some(study) = &self.slots[dst_slot].study else { return; diff --git a/src/app/reg_panel.rs b/src/app/reg_panel.rs index 78e7ca6..525fe97 100644 --- a/src/app/reg_panel.rs +++ b/src/app/reg_panel.rs @@ -227,6 +227,9 @@ impl ViewerApp { let Some(study) = &self.slots[fixed_slot].study else { return; }; + // A transform installed from elsewhere (a REG object in the tree, a + // propagation) needs the section that shows and clears it. + self.module_registration = true; let vol = study.volume.clone(); let analysis = analysis::analyse(&vol, &transform, None); let field = VectorField::sample(&vol, &transform, None, self.field_step_mm); @@ -314,6 +317,17 @@ impl ViewerApp { // result is on display, and while a run is in flight — the last one // because that is where its progress and its Cancel button live. if !both && self.registration.is_none() && self.reg_job.is_none() { + // The module is switched on, so the section says what it is + // waiting for rather than leaving an empty panel. + egui::CollapsingHeader::new(egui::RichText::new("Image registration").strong()) + .default_open(true) + .show(ui, |ui| { + ui.weak( + "Load a second dataset (File ▶ Add DICOM folder to B…) — \ + registration aligns one onto the other", + ); + }); + ui.separator(); return; } let mut run: Option = None; @@ -325,7 +339,7 @@ impl ViewerApp { let mut clear_landmarks = false; let mut save_field = false; - egui::CollapsingHeader::new(egui::RichText::new("Registration").strong()) + egui::CollapsingHeader::new(egui::RichText::new("Image registration").strong()) .default_open(true) .show(ui, |ui| { if let Some(job) = &self.reg_job { diff --git a/src/settings.rs b/src/settings.rs index ced63bd..d66bd8e 100644 --- a/src/settings.rs +++ b/src/settings.rs @@ -9,6 +9,10 @@ const APP_NAME: &str = "RustDICOMStation"; /// Settings key of the model root; the installer writes it too. pub const MODELS_DIR_KEY: &str = "models_dir"; +/// Settings keys of the two optional side-panel modules. +const MODULE_REG_KEY: &str = "module_image_registration"; +const MODULE_SIM_KEY: &str = "module_image_simulation"; + /// User preferences that survive a restart. #[derive(Clone, Debug, PartialEq, Eq)] pub struct Settings { @@ -20,6 +24,14 @@ pub struct Settings { /// `None` means use the platform-specific default returned by /// [`default_models_dir`]. pub models_dir: Option, + + /// *Modules ▶ Image registration*: the registration section is shown in + /// the side panel. + pub module_registration: bool, + + /// *Modules ▶ Image simulation*: the simulation section is shown in the + /// side panel. + pub module_simulation: bool, } impl Default for Settings { @@ -29,6 +41,10 @@ impl Default for Settings { Settings { theme: ThemePreference::Dark, models_dir: None, + // Both optional modules start hidden; the Modules menu turns them + // on and the choice is remembered. + module_registration: false, + module_simulation: false, } } } @@ -167,6 +183,22 @@ fn theme_to_str(t: ThemePreference) -> &'static str { } } +fn bool_to_str(b: bool) -> &'static str { + if b { + "on" + } else { + "off" + } +} + +fn bool_from_str(s: &str) -> Option { + match s.trim().to_ascii_lowercase().as_str() { + "on" | "true" | "yes" | "1" => Some(true), + "off" | "false" | "no" | "0" => Some(false), + _ => None, + } +} + fn theme_from_str(s: &str) -> Option { match s.trim().to_ascii_lowercase().as_str() { "dark" => Some(ThemePreference::Dark), @@ -218,6 +250,14 @@ fn parse(text: &str) -> Settings { if !v.is_empty() { s.models_dir = Some(PathBuf::from(v)); } + } else if key.eq_ignore_ascii_case(MODULE_REG_KEY) { + if let Some(b) = bool_from_str(value) { + s.module_registration = b; + } + } else if key.eq_ignore_ascii_case(MODULE_SIM_KEY) { + if let Some(b) = bool_from_str(value) { + s.module_simulation = b; + } } } s @@ -233,6 +273,13 @@ fn render(s: &Settings) -> String { if let Some(dir) = &s.models_dir { out.push_str(&format!("{MODELS_DIR_KEY} = {}\n", dir.display())); } + out.push_str(&format!( + "# optional side-panel modules (Modules menu) = on | off\n\ + {MODULE_REG_KEY} = {}\n\ + {MODULE_SIM_KEY} = {}\n", + bool_to_str(s.module_registration), + bool_to_str(s.module_simulation) + )); out } @@ -249,7 +296,7 @@ mod tests { ] { let s = Settings { theme, - models_dir: None, + ..Settings::default() }; assert_eq!(parse(&render(&s)), s, "round trip of {theme:?}"); } @@ -265,7 +312,7 @@ mod tests { parse("unknown = 3\nTHEME = Light \n"), Settings { theme: ThemePreference::Light, - models_dir: None + ..Settings::default() }, "case-insensitive key and value, surrounding space ignored" ); @@ -273,14 +320,31 @@ mod tests { parse("theme = white"), Settings { theme: ThemePreference::Light, - models_dir: None + ..Settings::default() }, "\"white\" accepted as an alias for light" ); let with_dir = Settings { theme: ThemePreference::Dark, models_dir: Some(PathBuf::from("D:/models")), + ..Settings::default() }; assert_eq!(parse(&render(&with_dir)), with_dir, "model dir round trip"); } + + #[test] + fn round_trips_the_module_flags() { + for (reg, sim) in [(false, false), (true, false), (false, true), (true, true)] { + let s = Settings { + module_registration: reg, + module_simulation: sim, + ..Settings::default() + }; + assert_eq!(parse(&render(&s)), s, "round trip of ({reg}, {sim})"); + } + assert!( + parse(&format!("{MODULE_REG_KEY} = TRUE")).module_registration, + "case-insensitive alias" + ); + } } From ffa19c47dafef6156878421fa8a0a0872d361c0b Mon Sep 17 00:00:00 2001 From: alexprotom Date: Thu, 27 Aug 2026 13:24:36 +0200 Subject: [PATCH 4/5] fix version to merge with main --- Cargo.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Cargo.toml b/Cargo.toml index 8db9542..c63582f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "rust-dicom-station" -version = "0.1.0" +version = "0.5.0" edition = "2021" description = "Fast, robust DICOM / RT DICOM viewer in pure Rust (CT/MR volumes, RTSTRUCT, RTDOSE, RTPLAN) with a three-view MPR layout" license = "MIT" From 319c79db088c2fd73ae654dc61f34ec2c784a15c Mon Sep 17 00:00:00 2001 From: alexprotom Date: Thu, 27 Aug 2026 14:19:05 +0200 Subject: [PATCH 5/5] Bump version to 0.6.0 --- Cargo.lock | 2 +- Cargo.toml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 0b45d0c..2db2d91 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6413,7 +6413,7 @@ dependencies = [ [[package]] name = "rust-dicom-station" -version = "0.1.0" +version = "0.6.0" dependencies = [ "anyhow", "burn", diff --git a/Cargo.toml b/Cargo.toml index c63582f..4e02f93 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "rust-dicom-station" -version = "0.5.0" +version = "0.6.0" edition = "2021" description = "Fast, robust DICOM / RT DICOM viewer in pure Rust (CT/MR volumes, RTSTRUCT, RTDOSE, RTPLAN) with a three-view MPR layout" license = "MIT"