diff --git a/AGENTS.md b/AGENTS.md index f7df93a6..97a50bb5 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -1,303 +1,93 @@ # KonfAI — Agent Guide -This is the canonical reference for humans and AI agents working on KonfAI. Read it before making changes. For the current code-quality audit and known issues, see [`AUDIT.md`](AUDIT.md). +Canonical reference for humans and AI agents. Read it before changing code. Known issues and the backlog live in [`AUDIT.md`](AUDIT.md); user-facing detail lives in `docs/` and `examples/`. ## 1. What KonfAI is -KonfAI is a **modular, extensible, and fully configurable deep-learning framework for medical imaging** (Boussot & Dillenseger, 2025 — arXiv:2508.09823). The central idea: a model, its data pipeline, losses/metrics, optimizers/schedulers, augmentations, post-processing, and the whole train/predict/evaluate workflow are described in **structured YAML config files, without modifying the underlying code**. YAML is mapped onto Python objects by a **reflection engine** (`konfai/utils/config.py`). You describe *what* you want and KonfAI instantiates it; the config doubles as a complete, self-contained, reproducible record of the experiment. +KonfAI is a modular, fully-configurable deep-learning framework for medical imaging (Boussot & Dillenseger, 2025 — arXiv:2508.09823). A model, its data pipeline, losses/metrics, optimizer/schedulers, augmentations, and the whole train/predict/evaluate workflow are described in **YAML** and mapped onto Python objects by a reflection engine — *without editing code*. The config is a complete, reproducible record of the experiment. KonfAI has produced top-ranking MICCAI-challenge results (SynthRAD, TrackRAD, CURVAS, PANTHER) across segmentation, registration, and synthesis. -**Why it exists:** to separate *experimental intent* (the YAML) from *implementation details* (the code), improving reproducibility, transparency, and traceability while still allowing advanced, non-standard workflows that other frameworks (PyTorch Lightning, MONAI, nnU-Net, NiftyNet) make hard. KonfAI has produced top-ranking results in MICCAI challenges (SynthRAD 2025, TrackRAD 2025, CURVAS 2024/2025, PANTHER) across segmentation, registration, and image synthesis. +Three pillars run through the codebase: -**Three design principles** (from the paper): -1. **Declarative configuration** — every component is identified by its Python class path and *all* constructor arguments live in the YAML; the pipeline is reconstructable from config alone (no hardcoded values). -2. **Modularity & controlled extensibility** — interchangeable components (models, losses, transforms, schedulers, augmentations) are instantiated at runtime through a unified registry; each inherits a shared base (`Network`, `Criterion`/loss, `Transform`, …). Extend by subclassing and referencing the new class by path in YAML — no core edits. -3. **Experiment traceability & workspace management** — each experiment lives in a dedicated workspace (config snapshot, logs, predictions, checkpoints, evaluation JSON). - -**Advanced capabilities** it abstracts natively (the reason it's worth the complexity): patch-based learning in **2D / 2.5D / 3D**, **test-time augmentation (TTA)**, **model ensembling**, **deep supervision / multi-head / intermediate-feature access**, and **multi-model setups** (GANs, teacher–student/adversarial). The key enabler is that models declare *every* submodule explicitly via `add_module`, so any layer/head/skip/attention block is **referenceable by name** in YAML (e.g. for a loss, a perceptual feature, or a deep-supervision head) — standard PyTorch `forward()` hides these. - -Two design pillars run through the entire codebase: - -1. **Config-by-reflection.** `apply_config(path)` inspects a callable's signature at call time and recursively builds its arguments from the YAML subtree it owns (`@config("Key")`). Defaults are *materialised back* into the YAML file so a run leaves a fully-resolved config on disk. -2. **Lazy, patch-based imaging.** Volumes are never loaded whole into RAM. Data is read as overlapping patches, optionally streamed chunk-by-chunk, and predictions are reassembled with overlap blending. **Preserving this invariant is mandatory.** - -A third, newer pillar is emerging: - -3. **Declarative models.** Networks are routed module graphs assembled via `add_module` (with `in_branch`/`out_branch` routing metadata). Models can be written as Python classes in `konfai/models/` *or* described entirely in a `.yml` file via the **YAML model builder** (`konfai/utils/model_builder.py`). See §7. +1. **Config-by-reflection.** `apply_config(path)` reads a callable's signature and builds its arguments from the YAML subtree it owns (`@config("Key")`), recursing into nested `@config` objects. Resolved defaults are written *back* to the file, so a run leaves a fully-resolved config on disk. **Reading a config mutates it.** +2. **Lazy, patch-based imaging.** Volumes are never loaded whole into RAM: data is read as overlapping patches (optionally streamed) and predictions are reassembled with overlap blending. **Mandatory invariant.** +3. **Declarative models.** Networks are routed `add_module` graphs — written as Python classes in `konfai/models/`, or entirely as a `.yml` via the YAML model builder. ## 2. Repository layout | Path | Role | |---|---| -| `konfai/` | Core package (config, data, network, workflows, utils) | -| `konfai-apps/` | **Separate** installable package `konfai_apps` (app management, HF repos, FastAPI server) — own `pyproject.toml` | +| `konfai/` | Core package (config, data, network, metric, workflows, utils) | +| `konfai-apps/` | **Independent** package `konfai_apps` (app management, HF repos, FastAPI server) — own `pyproject.toml`, deps, and CI | | `apps/` | Ready-to-use model app bundles (excluded from the `konfai` wheel) | -| `examples/` | Runnable `Segmentation` and `Synthesis` workflows (assume CWD = the example dir) | -| `docs/` | Sphinx documentation site | -| `tests/` | Core test suite (`tests/unit`, `tests/integration`) | -| `konfai-apps/tests/` | Apps test suite (run separately — see §8) | - -### Core package map - -| Module | Role | -|---|---| -| `konfai/utils/config.py` | **The reflection engine.** `Config` context manager + `@config` + `apply_config`. Read before any config change. | -| `konfai/utils/errors.py` | Typed exception hierarchy (`KonfAIError` → `ConfigError`, `TrainerError`, `DatasetManagerError`, …). Use these; do not invent new ones. | -| `konfai/utils/dataset.py` | Storage abstraction. `Dataset` + `AbstractFile` backends (`SitkFile`, `H5File`, `OmeZarrFile`, `DicomFile`) + the `Attribute` geometry sidecar. | -| `konfai/utils/dicom.py` | DICOM series reader/writer (pydicom). | -| `konfai/utils/ome_zarr.py` | OME-Zarr / OME-NGFF reader/writer (zarr; ngff-zarr is a declared but currently *unused* dep — see AUDIT.md). | -| `konfai/utils/ITK.py` | SimpleITK helpers (transform composition, resampling, mask-bbox cropping). | -| `konfai/utils/runtime.py` | `DistributedObject`, distributed setup/teardown, GPU/VRAM, TensorBoard process, logging. | -| `konfai/utils/utils.py` | `SUPPORTED_EXTENSIONS`, `get_module`, patch-slice math, `split_path_spec`. | -| `konfai/data/data_manager.py` | `DataManager`/`DatasetIter`: builds the flat (dataset, augmentation, patch) index, subset/validation selection, DDP sharding, DataLoaders. | -| `konfai/data/patching.py` | `ModelPatch`/`Accumulator`: sliding-window slicing + overlap-blended reassembly; `DatasetManager` per-case cache/stream layer. | -| `konfai/data/transform.py` | Preprocessing/postprocessing transforms (`Normalize`, `Standardize`, `Resample…`, `Mask`, `Crop`, …) + `TransformLoader`. | -| `konfai/data/augmentation.py` | Data augmentation primitives + ADA for DiffusionGAN. | -| `konfai/network/network.py` | `ModuleArgsDict` (routed graph), `Network` (training machinery), `ModelLoader`, `Measure`. The heart of the model system. | -| `konfai/network/blocks.py` | Reusable blocks (`ConvBlock`, `ResBlock`, `Attention`, `Concat`, `get_torch_module`, …). | -| `konfai/utils/model_builder.py` | The declarative YAML→`Network` builder with a safe module registry. | -| `konfai/models/` | Concrete architectures (segmentation, classification, generation, registration, representation). | -| `konfai/metric/measure.py` | `Criterion` hierarchy (losses + metrics) loaded by name from YAML. | -| `konfai/metric/schedulers.py` | Scalar weight schedulers for criteria. | -| `konfai/trainer.py` / `predictor.py` / `evaluator.py` | The `train()` / `predict()` / `evaluate()` pipelines. | -| `konfai/main.py` | CLI: `konfai` → `main()`, `konfai-cluster` → `cluster()`. Subcommands `TRAIN`, `RESUME`, `PREDICTION`, `EVALUATION`. | -| `konfai/__init__.py` | Env-backed directory/state accessors, device utilities. | - -## 3. Architecture in depth - -### 3.1 Config reflection (`konfai/utils/config.py`) -- `@config("Key")` tags a class/callable with the YAML branch it owns (sets `function._key`). -- `apply_config(path)(Cls)()` opens a `Config` on the resolved dot-path, then walks `inspect.signature` params: primitives are coerced, `Path` resolved, `list`/`dict` expanded, and nested `@config` objects recursively instantiated. -- **Reading config mutates it**: resolved defaults are merged back into the YAML file on `Config.__exit__`. A successful run rewrites the active config with all defaults filled in. -- Defaults use the `"default|"` marker; Python `None` round-trips through YAML as the string `"None"`. - -### 3.2 Data pipeline -`Dataset` (`utils/dataset.py`) maps a logical `(group, name)` onto on-disk artefacts via a backend chosen by `file_format`: `h5`→`H5File`, `omezarr`/`ome-zarr`/`zarr`→`OmeZarrFile`, `dicom`→`DicomFile`, anything else→`SitkFile` (`.mha`/`.nii`/npy/xml/vtk/transforms). Geometry travels in an `Attribute` sidecar (`Origin`, `Spacing`, `Direction` + arbitrary metadata). - -`DataManager`/`DatasetIter` (`data/data_manager.py`) build a flat `(x=dataset, a=augmentation, p=patch)` index, apply subset/validation selection, shard across DDP ranks, and feed PyTorch DataLoaders. `DatasetManager`/`ModelPatch`/`Accumulator` (`data/patching.py`) do the patch slicing and overlap-blended reassembly. Transforms (`data/transform.py`) and augmentations (`data/augmentation.py`) are config-bound and run per case. - -**Data convention:** arrays are always **channel-first** `[C, (Z), Y, X]`. Geometry/spacing is **`(x, y, z)`** order (SimpleITK convention). +| `examples/` | Runnable `Segmentation` / `Synthesis` workflows (assume CWD = the example dir) | +| `docs/` · `tests/` | Sphinx site · core test suite (`tests/unit`, `tests/integration`) | -### 3.3 Network graph (`konfai/network/network.py`) -- `ModuleArgsDict` is a `torch.nn.Module` whose ordered children each carry routing metadata. `add_module(name, module, in_branch=[...], out_branch=[...], alias=..., ...)` wires a dataflow graph (a string-keyed branch register; branch `'0'` is the implicit input/main path). -- `Network` adds optimizer/scheduler/criterion wiring, patch-based forward, gradient/cross-GPU checkpointing, alias-based state-dict load/save, and loss/metric aggregation via `Measure`. -- `ModelLoader.get_model()` resolves the configured `classpath`: a `.yml`/`.yaml` suffix → YAML builder (§7); otherwise import the class from `konfai.models` and build it via `apply_config`. Non-`Network` classes are wrapped in `MinimalModel`. +**Core modules worth knowing:** `utils/config.py` (the reflection engine — read before any config change); `utils/dataset.py` (storage backends `SitkFile`/`H5File`/`OmeZarrFile`/`DicomFile` + the `Attribute` geometry sidecar); `data/data_manager.py` + `data/patching.py` (lazy patch index, DDP sharding, overlap-blended reassembly); `network/network.py` (`ModuleArgsDict`/`Network`/`ModelLoader`/`Measure` — the heart of the model system); `metric/measure.py` (`Criterion` = losses + metrics); `data/{transform,augmentation}.py`; `trainer.py`/`predictor.py`/`evaluator.py` (the pipelines); `main.py` (CLI); `utils/{errors,runtime,model_builder}.py`. -### 3.4 Workflows & runtime -`train()`/`predict()`/`evaluate()` are wrapped by `run_distributed_app` (`utils/runtime.py`), which builds the configured `DistributedObject`, sets the `KONFAI_*` env vars, forces `KONFAI_CONFIG_MODE='Done'`, and spawns one process per GPU (or submits to SLURM via `submitit`). All disk/logging side effects are gated on `global_rank == 0`. +## 3. How it fits together -## 4. Execution flow (end to end) +**Commands → config files.** KonfAI is command-driven; four CLI states map to three YAML files: -``` -CLI (main.py) ─▶ train()/predict()/evaluate() [@run_distributed_app] - └▶ build_*() ─▶ configure_workflow_environment() # sets KONFAI_config_file/ROOT/STATE + dirs - └▶ KONFAI_CONFIG_MODE = 'Done' - └▶ apply_config()(Trainer/Predictor/Evaluator)() # YAML → object graph - └▶ execute_distributed_object() - └▶ mp.spawn → DistributedObject.__call__(rank) - └▶ setup_gpu() → model.init()/init_outputs_group()/_compute_channels_trace() - └▶ DataManager → DataLoader (lazy patches) - └▶ training loop / inference / evaluation -``` - -## 4b. Commands, config files & workspace - -KonfAI is command-driven. Four CLI states map to three YAML config files: - -| Command | Config file | Top-level key | Purpose | +| Command | File | Root key | Purpose | |---|---|---|---| | `TRAIN` / `RESUME` | `Config.yml` | `Trainer:` | Model + dataset + losses + augmentations + optimizer/schedulers + training params | -| `PREDICTION` | `Prediction.yml` | `Predictor:` | Load model(s), patch/TTA/ensemble inference, `outputs_dataset` post-processing | -| `EVALUATION` | `Evaluation.yml` | `Evaluator:` | Compare predictions vs ground truth, per-case + aggregate metric JSON | - -Each run creates/updates a **workspace** organised by `train_name` (model name): +| `PREDICTION` | `Prediction.yml` | `Predictor:` | Load model(s), patch/TTA/ensemble inference, output post-processing | +| `EVALUATION` | `Evaluation.yml` | `Evaluator:` | Predictions vs ground truth → per-case + aggregate metric JSON | -``` -Workspace/ - Checkpoints//.pt # checkpoints (BEST or ALL) - Setups//Config_0.yml # full resolved config snapshot - Statistics//events.out.* # TensorBoard logs - Predictions//Dataset/... # predicted volumes + Prediction.yml - Evaluations//Metric_*.json # per-case + aggregate metrics + Evaluation.yml - Dataset//. # input data tree (e.g. CT.mha, MASK.mha) -``` - -### Config anatomy (high-value keys an agent will touch) -- **`Trainer.Model`**: `classpath` (e.g. `segmentation.UNet.UNet` or `UNet.yml`), `Optimizer`, `schedulers`, `outputs_criterions` (loss/metric attached to a **named** module output, e.g. `UNetBlock_0:Head:Softmax`), plus the model's own constructor args. -- **`Dataset`**: `groups_src` → `groups_dest` (each group has `transforms`, `is_input`, `patch_transforms`); `dataset_filenames` (e.g. `./Dataset/:a:mha` — the `:a:`/`:i:` flags mean append/intersect and the suffix is the format); `use_cache`, `batch_size`, `subset`, `validation`, `validation_augmentations` (default `true`; set `false` to validate only on base samples — train and validation datasets are then prepared separately so augmented variants never leak into validation), `shuffle`, `inline_augmentations`. -- **`Patch`**: `patch_size`, `overlap`, `extend_slice` (>0 ⇒ 2.5D, requires `patch_size[0]==1`), `pad_value`. -- **`augmentations`**: `DataAugmentation_*` → `data_augmentations` (Flip/Rotate/…); `nb` augmentations per sample. -- **Training params**: `epochs`, `it_validation`, `autocast` (AMP), `gradient_checkpoints`, `gpu_checkpoints`, `ema_decay`, `data_log`, `save_checkpoint_mode` (`BEST`/`ALL`), `EarlyStopping`, `manual_seed`. (Training also stops automatically once the schedulers decay the optimizer learning rate to `≤ 0`, via `EarlyStoppingBase.stop()`.) -- **`Predictor.outputs_dataset`**: per output module, `OutputDataset` with `after_reduction_transforms` (Argmax/…), `final_transforms` (TensorCast), `patch_combine`, `reduction` (TTA: Mean/Median/…), `inverse_transform`; top-level `combine` does model ensembling. -- **`Evaluator.metrics`**: per target group, `targets_criterions` → `criterions_loader` (e.g. `Dice`). +Each run writes a **workspace** keyed by `train_name`: `Checkpoints/`, `Setups/` (resolved config snapshot), `Statistics/` (TensorBoard), `Predictions/`, `Evaluations/` (metric JSON), `Dataset/`. -## 5. Environment variables +**Conventions.** Arrays are **channel-first** `[C,(Z),Y,X]`; geometry/spacing is **`(x,y,z)`** (SimpleITK). `Attribute` geometry keys are `Origin`/`Spacing`/`Direction`. -These are process-global state set by `configure_workflow_environment`; accessors in `konfai/__init__.py` read them. +**Network graph.** `add_module(name, module, in_branch=[...], out_branch=[...], alias=...)` wires a string-keyed branch register (branch `'0'` = input; execution = insertion order). **Named module outputs are referenceable in YAML** — e.g. an `outputs_criterions` key is a module's dotted path like `UNetBlock_0:Head:Softmax` (the `:`/`.` separators are load-bearing). `out_branch:[-1]` marks a terminal/deep-supervision head; `alias` lists are positional and load-bearing for pretrained-weight remapping. -| Variable | Meaning | -|---|---| -| `KONFAI_config_file` | Path to the active YAML config (note the **lowercase** `config_file` suffix — easy to mistype). Read directly from `os.environ` with no fallback. | -| `KONFAI_CONFIG_MODE` | Config behaviour: `Done` (workflows require this), `default`, `Import` (suppresses `apply_config` side effects during module import), `interactive`, `remove`. **Note:** `interactive`/`remove` are effectively dead in the current product (see AUDIT.md). | -| `KONFAI_ROOT` / `KONFAI_STATE` | Workflow root dir / current `State` (TRAIN/RESUME/PREDICTION/EVALUATION). | -| `KONFAI_CONFIG_PATH` | Hidden channel set by `apply_config`, read back by `PerceptualLoss`. | -| `CUDA_VISIBLE_DEVICES`, `KONFAI_MASTER_PORT`, `KONFAI_CLUSTER`, `KONFAI_OVERWRITE` | Distributed/cluster control. | - -## 6. Imaging formats (DICOM & OME-Zarr) +**Runtime.** Workflows run under `run_distributed_app` (`utils/runtime.py`): it builds the configured `DistributedObject`, sets the `KONFAI_*` env vars, forces `KONFAI_CONFIG_MODE='Done'`, and spawns one process per GPU (or submits to SLURM via `submitit`). Disk/log side effects are gated on `global_rank == 0`. -Both are first-class dataset backends, dispatched by `file_format` and reachable from config (`SUPPORTED_EXTENSIONS` validates the format). Verified working end-to-end (round-trips, patch reads, geometry): +For the full config-key catalogue and a concrete end-to-end trace, read the `docs/` config guides and `examples/`. -- **DICOM** (`utils/dicom.py`): groups by `SeriesInstanceUID`, sorts slices by `ImagePositionPatient` projected on the `ImageOrientationPatient` normal, derives `(origin, spacing(x,y,z), direction(9))`, applies HU rescale, supports header-only metadata + lazy per-slice patch reads, and writes uncompressed series. Geometry round-trips **byte-identically to SimpleITK's GDCM reader**. Caveats: write quantises floats to int16 and always uses the CT SOP class; left-handed directions normalise to right-handed on round-trip (standard DICOM behaviour); multi-frame/enhanced DICOM and irregular slice spacing are not handled. See AUDIT.md. -- **OME-Zarr** (`utils/ome_zarr.py`): a thin adapter over **`ngff-zarr`** (no hand-rolled NGFF parsing) — `ngff_zarr.from_ngff_zarr` for lazy chunk-wise reads, `to_ngff_zarr` for writes; output is `ngff-zarr`-interoperable (verified). **Resolution selection:** a multiscale pyramid level can be chosen per source via the dataset-spec suffix **`omezarr@`** (e.g. `./Dataset:a:omezarr@2` reads pyramid level 2, coarser), parsed by `split_format_level()` and threaded to `OmeZarrFile` — independent of any transform; level 0 (full res) by default. The chosen level's `scale` becomes the KonfAI `Spacing`. Caveat: `Direction` is not representable in NGFF and round-trips only via a proprietary `konfai` group attr; KonfAI's own writer emits a single level (multi-level pyramids come from `ngff_zarr.to_multiscales`). +## 4. Extending KonfAI -## 7. Declarative YAML model builder +Every extension point is **"subclass a base, reference it by classpath in YAML"** — no core edits: -`konfai/utils/model_builder.py` builds a routed `Network` from a `.yml` file (`name`/`parameters`/`network`/`modules`). It is **safe by construction**: node types must come from two curated registries (`_MODULE_REGISTRY` for `nn.Module` factories, `_OBJECT_REGISTRY` for config objects) — no `eval`/import injection. Features: `${param}` exact-string references (preserve Python values), `{$multiply: [...]}`, `{$object: BlockConfig, args: {...}}`, nested graphs, full `add_module` routing. +- **Model:** subclass `network.Network`, build the graph in `__init__` via `add_module`. Reference `classpath: module.MyNet`, a local `Model:MyNet`, or a `.yml`. +- **Loss / metric:** subclass `metric.measure.Criterion`; `forward` returns a `Tensor` (loss) or a `(value, dict)` tuple (metric — consumers `isinstance`-branch). Attach under `outputs_criterions`/`metrics` to a **named module output**. Optional-dep criteria import lazily via `_require_optional(...)` and raise an actionable `MeasureError` — never a bare top-level import. +- **Transform:** subclass `data.transform.Transform`; implement `__call__` **and** `transform_shape()` (must predict the output spatial shape *exactly* — patch planning depends on it). Pair `inverse()` if `apply_inverse`. +- **Augmentation:** subclass `data.augmentation.DataAugmentation`; `_state_init` (sample params per case index) + `_compute` (apply lazily). Only `Mask`/`Permute` may change shape. +- **Imaging format:** add a `Dataset.AbstractFile` backend, dispatch it in `File.__enter__`, register aliases in `SUPPORTED_EXTENSIONS`; import-guard the heavy lib. -`ModelLoader` activates it when `classpath` ends in `.yml`/`.yaml` (resolved relative to `KONFAI_config_file`, or under `konfai_root()/models` for `default|…`). See `examples/Segmentation/UNet.yml`. +**Classpaths:** a bare name (e.g. `Dice`) resolves inside that kind's package; `module:Class` imports *any* module — a local file (`Loss:MyWrapper`) or an installed library (`monai.losses:DiceLoss`, `torch:nn:L1Loss`). -**Status (verified):** the example `UNet.yml` builds a model **identical** to the Python `UNet` (1,934,299 params, matching forward). **It does not yet replace `konfai/models/`:** -- The registry has only 13 module types. Replacing the existing models needs many more (`Linear`, `Add`/`Multiply`, `Attention`, `Upsample`, `View`/`Select`, `LatentDistribution`, norm/activation factories, …) and the enum config objects (`DownsampleMode`/`UpsampleMode`/`NormMode`). -- Roughly 70–80% of `models/` is pure `add_module` wiring (UNet, NestedUNet, ResNet, GAN, VAE) and is migratable once the registry grows. The rest carries genuine Python logic that the builder cannot express: custom `forward`/sampling/training loops in `ddpm`, `diffusionGan`, `cStyleGan`, `registration` (spatial-transformer grid math), `convNeXt` (LayerNorm/DropPath), `representation`. +**YAML model builder** (`utils/model_builder.py`): builds a `Network` from a `.yml`, **safe by construction** (node types must come from two curated registries — no `eval`/import injection). It *complements* `models/` today and can replace the feed-forward subset once the registry grows (`UNet`/`NestedUNet`/`ResNet` are migratable; custom-`forward` models like DDPM/DiffusionGAN/ConvNeXt are not). See AUDIT.md. -So: **the YAML builder complements `models/` today and can replace the feed-forward subset after the registry is expanded.** See AUDIT.md for the migration plan. +## 5. Apps (`konfai-apps`) -**Model → YAML migration matrix** (how much of each zoo model the builder can express today): +A separate package layered on KonfAI's **public** API (core never imports it). An "app" bundles a config + custom `.py` + `.pt` weights, resolved from a Local dir, a HuggingFace repo, or a Remote server; the `apps/*` bundles are thin CLI wrappers. -| Model | Migratable now? | Blocker | -|---|---|---| -| UNet, NestedUNet/UNet++, ResNet | ✅ Yes | Pure `add_module` wiring; registry already covers it. | -| GAN/CycleGAN, VAE, VoxelMorph | ⚠️ Partial | Needs registry growth (`Linear`/`Add`/`LatentDistribution`/…) + channel-math/factory constructs. | -| ConvNeXt, DDPM, DiffusionGAN, cStyleGAN, Representation | ⛔ No | Custom Python `forward`/sampling/training logic the builder has no construct for. | +> ⚠️ **Trust model.** Resolving an app **copies and imports its `.py` files** and **pip-installs its `requirements.txt`** → it runs arbitrary code and dependency installs. **Only resolve apps from sources you trust.** -## 7b. How a run actually works (mental model) +## 6. Running things -Tracing one config key end-to-end makes the reflection engine concrete. Given `Config.yml`: -```yaml -Trainer: - Model: - classpath: segmentation.UNet.UNet - UNet: - dim: 2 - channels: [1, 32, 64, 128, 256] -``` -1. `train()` → `build_train()` → `configure_workflow_environment()` sets `KONFAI_config_file=Config.yml`, then forces `KONFAI_CONFIG_MODE='Done'`. -2. `apply_config("Trainer")(Trainer)()` opens a `Config` on the `Trainer` subtree and walks `Trainer.__init__`'s signature. The `model: ModelLoader` param (a `@config`-tagged type) is recursively built. -3. `ModelLoader.get_model()` reads `classpath`. No `.yml` suffix → `get_module("segmentation.UNet.UNet", "konfai.models")` imports the class, then `apply_config("Trainer.Model.UNet")(UNet)()` binds `dim`, `channels`, … from the `UNet:` subtree (defaults materialised back into the file). -4. `UNet.__init__` issues `add_module(...)` calls → a routed `ModuleArgsDict` graph. -5. On `Config.__exit__`, the resolved subtree is merged back to `Config.yml` (so the file now records every default). - -Swap `classpath: UNet.yml` and step 3 instead calls `build_model_from_yaml` (§7) — same `add_module` graph, no Python class. - -## 7c. Extending KonfAI (custom components) - -Every extension point is "subclass a base, reference it by class path in YAML" — no core edits. The base classes inherit the config + (for modules) the routing machinery. - -- **Custom model:** subclass `konfai.network.network.Network`; build the graph in `__init__` via `add_module(name, module, in_branch=[...], out_branch=[...])`. Reference with `classpath: my_module.MyNet` (importable) or, for a pure feed-forward graph, write a `.yml` and register any new block via `model_builder.register_module`. -- **Custom loss/metric:** subclass the criterion base in `konfai/metric/measure.py`; `forward` returns a `Tensor` (loss) or a tuple (metric). Reference under `outputs_criterions`/`metrics` by class path, attached to a **named module output** (e.g. `UNetBlock_0:Head:Softmax`). -- **Custom transform:** subclass `konfai.data.transform.Transform`; implement `__call__(name, tensor, cache_attribute)` **and** `transform_shape()` (must predict the output spatial shape exactly — patch planning depends on it). Pair `inverse()` symmetrically if `apply_inverse`. -- **Custom augmentation:** subclass `konfai.data.augmentation.DataAugmentation`; implement `_state_init` (sample params per case index) and `_compute` (apply lazily). Return one shape per input; only `Mask`/`Permute` may change the shape. -- **New imaging format:** add a `Dataset.AbstractFile` backend in `konfai/utils/dataset.py`, dispatch it in `File.__enter__`, and register its aliases in `SUPPORTED_EXTENSIONS` (`utils/utils.py`). Keep the heavy reader lib import-guarded. - -**Routing rules to respect** (`add_module`): branch `'0'` is the implicit input; an `in_branch` must be produced by an earlier module (execution = insertion order, no topo-sort); `out_branch: [-1]` marks a terminal/deep-supervision head; module names must contain no `.`; `alias` lists are positional and load-bearing for pretrained weight remapping. - -## 7d. Apps & packaged models (`konfai-apps`, `apps/*`) - -`konfai-apps` is a **separate package** layering remote/app/packaged-model functionality on top of the core public API (it never reaches into core internals; core never imports it). An "app" bundles `app.json` metadata + a KonfAI config (`Prediction.yml`, …) + custom `.py` + `.pt` weight checkpoints. Apps are resolved from a **Local** directory, a **HuggingFace** repo, or a **Remote** server. The `apps/*` bundles (`totalsegmentator`, `mrsegmentator`, `impact_synth`) are thin CLI wrappers that resolve an HF app and call `KonfAIApp.pipeline()`. Pretrained models are distributed as `.pt` checkpoints downloaded on demand. There is also a FastAPI server (`app_server.py`) with job lifecycle, GPU-semaphore scheduling, SSE log streaming, and TTL'd results. - -> ⚠️ **Trust model (read before resolving any app).** Resolving an app **copies the app's `.py` files into the working directory and imports them**, and installs the app's `requirements.txt` via a `pip install` subprocess. A downloaded app therefore runs **arbitrary code and arbitrary dependency installs** on your machine. This is inherent to "packaged model = code + weights + config". **Only resolve apps from sources you trust** (your own repos, vetted HF orgs). Do not point the loader at untrusted HuggingFace IDs or remote servers. See AUDIT.md §4b. - -## 7e. Metrics & criteria - -Criteria live in `konfai/metric/measure.py` (`Criterion` hierarchy, loaded by class path from `outputs_criterions`/`metrics`, weight-scheduled via `konfai/metric/schedulers.py`). Notes for agents: - -- **Optional-dependency criteria** import heavy packages lazily through `_require_optional(module, criterion=…, extra=…)`, which raises an actionable `MeasureError` (with the `pip install konfai[]` hint) at construction. `SSIM` needs `konfai[ssim]` (`scikit-image`); `FID` needs `konfai[fid]` (`scipy`+`torchvision`); `LPIPS` needs `konfai[lpips]`. Add new optional-dep criteria the same way — never a bare `import` that fails mid-run. -- **`Criterion.forward` is typed `-> Tensor` but several subclasses return a `(loss, dict)` tuple** (metrics); consumers `isinstance`-branch. Follow the existing pattern of the criterion you extend. -- **`update_scheduler`** selects the active weight scheduler for the current iteration; an empty schedule raises `ConfigError`, and iterations past the last window clamp to the last scheduler. - -## 8. Running things - -### Tests -The dev environment must have the imaging extras installed to exercise the real DICOM/OME-Zarr/ITK paths (the Pixi `dev` env does; a bare `pip install .[dev]` does **not** — see AUDIT.md tooling drift). - -```bash -pixi run test # core unit + integration (tests/) — pytest -q -pixi run --environment dev python -m pytest tests/unit -q # core unit only -pip install -e ./konfai-apps && pixi run --environment dev python -m pytest konfai-apps/tests # apps suite -pixi run test-cov # with coverage -``` - -Baseline: `tests/unit` green, `tests/integration` green, `konfai-apps/tests/unit` 25 passed. - -> **Caveat:** root `pytest testpaths=['tests']` excludes `konfai-apps/tests`, so `pixi run test` does **not** run the apps suite. Run it explicitly. `konfai-apps` is an **independent package**: the core dev env no longer carries its runtime deps (fastapi/uvicorn/python-multipart) — install the package itself with `pip install -e ./konfai-apps` (which pulls them) before running its suite, exactly as its CI does. - -### Lint / format / types / build / docs ```bash -pixi run --environment dev lint # ruff check konfai konfai-apps/konfai_apps -pixi run --environment dev format-check # ruff format --check -pixi run --environment dev typecheck # mypy konfai -pixi run --environment dev build # python -m build -pixi run --environment docs build-docs # Sphinx HTML -pixi run check # lint + format-check + test (run before finalising) +pixi run check # lint + format-check + test (run before finalising) +pixi run test # core unit + integration (tests/) +pixi run --environment dev typecheck # mypy konfai +pip install -e ./konfai-apps && pixi run --environment dev python -m pytest konfai-apps/tests # apps suite (separate) ``` -(`lint`/`format`/`format-check` exist in both the `dev` and `lint` envs; pass `--environment`.) - -## 9. Optional dependencies - -Install via `pip install konfai[]` or Pixi (the `dev` env bundles the runtime extras; docs deps live in the `docs` env). -| Extra | Packages | Use | -|---|---|---| -| `itk` | `SimpleITK` | ITK image I/O + transforms | -| `hdf5` | `h5py` | HDF5 datasets | -| `dicom` | `pydicom` | DICOM series backend | -| `omezarr` | `zarr`, `ngff-zarr` | OME-Zarr / NGFF backend | -| `imaging` | `SimpleITK`, `h5py`, `pydicom`, `zarr`, `ngff-zarr` | All imaging backends | -| `monitoring` | `nvidia-ml-py` (imports as `pynvml`) | GPU VRAM monitoring | -| `tensorboard` | `tensorboard` | Training visualisation | -| `vtk` / `lpips` / `cluster` | `vtk` / `lpips` / `submitit` | Mesh I/O / perceptual loss / SLURM | -| `ssim` / `fid` | `scikit-image` / `scipy`+`torchvision` | SSIM criterion / FID criterion | -| `all` | every runtime extra | Full install | -| `dev` | pytest, ruff, build, Sphinx, … | Development | +The Pixi `dev` env carries the imaging extras; a bare `pip install .[dev]` does not. `pixi run test` does **not** run `konfai-apps/tests` — install that package first (it pulls its own runtime deps), exactly as its CI does. Install runtime extras with `pip install konfai[]` (`itk`, `hdf5`, `dicom`, `omezarr`, `imaging`, `tensorboard`, `lpips`, `ssim`, `fid`, `cluster`, …). -> Heavy deps (`SimpleITK`, `h5py`, `pydicom`, `zarr`, `pynvml`, `tensorboard`) are **optional** and import-guarded — code must fail at point-of-use with an install hint, not at import. - -## 10. Invariants — do NOT break these +## 7. Invariants — do NOT break - **Never load a full volume into RAM.** Use lazy/patch/streaming access (`can_stream_patch`, `read_data_slice`). -- **Channel-first arrays** `[C,(Z),Y,X]`; **spacing/geometry in `(x,y,z)`** order. `Attribute` geometry keys are `Origin`/`Spacing`/`Direction`. -- **`Attribute` stringifies every value** and reparses geometry via `np.fromstring(s[1:-1], sep=" ")` — only flat scalars/1-D arrays round-trip (nested arrays break; this is a real bug, see AUDIT.md). Read via `__getitem__`/`get_np_array`; do not pre-suffix keys with `_`. -- **`KONFAI_config_file` and `KONFAI_CONFIG_MODE` must be set** before any `Config()`; tests must `monkeypatch.setenv` both. Workflows require `KONFAI_CONFIG_MODE='Done'`. -- **Patch ordering** must match between `disassemble` (read) and `Accumulator` (write); for PREDICTION/EVALUATION, all patches of a case must stay on the same DDP rank. -- **`outputs_criterions` keys** must equal a module's dotted path (e.g. `UNetBlock_0:Head:Argmax`); the `:`/`.` separators are load-bearing. -- **`state_dict` load/save deliberately does not recurse into nested `Network`s** (each owns its optimizer/state); alias lists are positional and load-bearing for pretrained weights. -- **YAML model builder is the trusted-untrusted boundary**: only registry types may be instantiated; module names contain no `.`. -- **Format aliasing**: `ome-zarr`/`ome_zarr`/`zarr` → `omezarr`; keep `SUPPORTED_EXTENSIONS` consistent. +- **Channel-first `[C,(Z),Y,X]`; spacing `(x,y,z)`.** `Attribute` stringifies every value and reparses geometry via `np.fromstring` — only flat scalars / 1-D arrays round-trip (see #5 in AUDIT.md). Read via `__getitem__`/`get_np_array`. +- **`KONFAI_config_file` + `KONFAI_CONFIG_MODE` must be set before any `Config()`** (tests must `monkeypatch.setenv` both); workflows require `KONFAI_CONFIG_MODE='Done'`. Reading a config rewrites it on disk. +- **Patch ordering** must match between read (`disassemble`) and write (`Accumulator`); for PREDICTION/EVALUATION all patches of a case stay on the same DDP rank. +- **`outputs_criterions` keys equal a module's dotted path**; the `:`/`.` separators are load-bearing. +- **`state_dict` load/save does not recurse into nested `Network`s** (each owns its optimizer/state); alias lists are positional. +- **The YAML model builder is the trusted/untrusted boundary** — only registry types; module names contain no `.`. - **`konfai-apps` is a separate package**; `apps/` is excluded from the `konfai` wheel. -## 11. Coding conventions -- Line length 120 (Ruff). Type annotations on new public functions. Apache-2.0 SPDX header on every new source file. -- No wildcard imports; prefer `pathlib.Path`. Use existing error classes from `konfai/utils/errors.py`. -- Do not import heavy optional deps (`SimpleITK`, `h5py`, `pydicom`, `zarr`) at module top level in code paths that don't need them — guard with `try/except ImportError` + a `_require_*()` helper. - -## 12. Commit conventions -- **Conventional Commits** are enforced in CI (`cz check`): `type(scope): subject` (`feat`, `fix`, `perf`, `docs`, `build`, `ci`, `refactor`, `test`, `chore`). -- A `commit-msg` hook + CI reject AI-agent branding (`maestro`, `claude`, `codex`, "generated by/with") — **avoid these words even in file names referenced in the subject**. -- Imperative present tense, subject < 72 chars. No AI co-author trailers. - -## 13. Rules for AI agents -1. **Read before editing.** Open every file you change. -2. **Keep diffs small.** One logical change per PR; no unrelated reformats. -3. **Run `pixi run check`** (and the apps suite if you touched `konfai-apps`) before finalising. -4. **No new runtime dependencies** without explicit request + matching `pyproject.toml` update (declare the dep in the *same* commit as the code that uses it). -5. **Never load imaging datasets fully into RAM.** -6. **Update docs** when changing user-facing CLI/config behaviour, and update `tests/unit/test_config.py` when changing config binding. -7. **Use existing error types**; do not invent exceptions. -8. **Do not skip pre-commit hooks** with `--no-verify`. +## 8. Conventions & rules -## 14. Common pitfalls -- `Config.__init__` reads `KONFAI_config_file` straight from `os.environ`; tests must set both env vars via `monkeypatch.setenv`. -- `nvidia-ml-py` imports as `pynvml` — not the same string. -- Format readers (DICOM/OME-Zarr/ITK) live in `konfai/utils/` and are imported by `konfai/data/` — do not move them into `data/`. -- `pixi run test` does not run `konfai-apps/tests`; run them explicitly. -- The pip `[dev]` extra and the Pixi `dev` feature define *different* dev environments (see AUDIT.md) — prefer Pixi for full coverage. -- Reading a config rewrites it on disk (defaults materialised on `Config.__exit__`). +- **Code:** line length 120 (Ruff); type annotations on new public functions; Apache-2.0 SPDX header on every new source file; prefer `pathlib.Path`; use the error classes in `utils/errors.py` (do not invent exceptions); import-guard heavy optional deps (`SimpleITK`/`h5py`/`pydicom`/`zarr`) — fail at point-of-use with an install hint, not at import. +- **Commits:** Conventional Commits (`cz check`): `type(scope): subject`, imperative, < 72 chars. A `commit-msg` hook + CI **reject AI-agent branding** (`claude`/`codex`/"generated by/with") and AI co-author trailers — avoid them. +- **For agents:** read before editing; keep diffs small (one logical change per PR, no unrelated reformats); run `pixi run check` (and the apps suite if you touched `konfai-apps`) before finalising; no new runtime dependency without an explicit request + a matching `pyproject.toml` update in the same commit; update docs and `tests/unit/test_config.py` when changing config binding; do not skip pre-commit with `--no-verify`. diff --git a/AUDIT.md b/AUDIT.md index 2b431b41..d4ea1775 100644 --- a/AUDIT.md +++ b/AUDIT.md @@ -1,206 +1,64 @@ # KonfAI — Code Audit -**Scope.** The whole repository (`konfai/`, `konfai-apps/`, tests, docs, packaging), with deep focus on three areas: OME-Zarr support, DICOM support, and the declarative YAML model builder. +**Scope.** The whole repository, with deep focus on OME-Zarr, DICOM, and the declarative YAML model builder. -**Method.** Every subsystem was read in depth, then high-severity findings were verified against the code and, where runnable, with a minimal reproduction in the Pixi `dev` environment (torch 2.12, SimpleITK 2.5.5, pydicom 3.0.2, zarr 3.2.1, ngff-zarr 0.37, numpy 2.5). Each finding carries a verdict (`confirmed` / `partially-confirmed` / `overstated`); two claims were downgraded on verification (see §8). +**Method.** Every subsystem was read in depth; high-severity findings were verified against the code and, where runnable, reproduced in the Pixi `dev` environment. Each finding carries a verdict (`confirmed` / `partially-confirmed` / `overstated`). -**Test baseline.** `tests/` 188 passed, `konfai-apps/tests/unit` 25 passed; ruff lint + format clean. +**Status.** The modernization stack (PRs #6–#9) and the audit follow-up (PR #10) are merged to `main`. About half the confirmed bugs are fixed with regression tests; the rest are a prioritized backlog (below). `tests/` and `konfai-apps/tests` are green; lint/format clean. ---- +## Focus areas (all work end-to-end) -## 1. Executive summary +| Area | Verdict | +|---|---| +| **OME-Zarr** (`utils/ome_zarr.py`) | ✅ Round-trips verified, `ngff-zarr`-interoperable. ⚠️ Re-implements parts of `ngff-zarr`; some multiscale read paths are dead code; `Direction` rides a proprietary attr (not NGFF-standard). | +| **DICOM** (`utils/dicom.py`) | ✅ Geometry round-trips **byte-identical to SimpleITK's GDCM reader** (incl. left-handed → right-handed normalization). Caveats: lossy int16 write, CT-only SOP class, single-gap z-spacing. | +| **YAML model builder** (`utils/model_builder.py`) | ✅ `examples/Segmentation/UNet.yml` builds a model **identical** to the Python `UNet` (1,934,299 params, matching forward). ⛔ Cannot yet replace `models/`: registry too small (13 types) and custom-`forward` models are inexpressible. | -KonfAI is a genuinely capable, well-architected framework: the config-by-reflection engine, the routed `ModuleArgsDict` graph, and the lazy patch/stream data layer are coherent and powerful, and the three focus areas **work end-to-end**. The main risks are not in the headline features but in (a) a handful of **confirmed correctness bugs** in augmentation/transform/VAE/CLI/trainer code, (b) **pervasive low-level smells** (mutable default args, stringly-typed `Attribute`, `os.environ` used as control-flow state, name-string coupling), (c) **dead code & duplication** (notably `ITK.py`, the dead OME-Zarr multiscale path, repeated SimpleITK transform-serialization), and (d) **tooling/test-coverage gaps** (apps tests not in the default run; imaging happy-paths not in CI; pip-`dev` vs Pixi-`dev` divergence). +## Fixed (merged to `main`) -| Focus area | Status | Headline | -|---|---|---| -| **OME-Zarr** | ✅ Works; ⚠️ reinvents ngff-zarr | Round-trips verified; output is ngff-zarr–interoperable; but `ngff-zarr` is a declared-yet-unused dep and ~half the module (multiscale read/`select_level`) is not wired into the pipeline. | -| **DICOM** | ✅ Works (matches SimpleITK) | Geometry round-trips **byte-identical to SimpleITK's GDCM reader**, incl. left-handed/flipped-z normalization. Caveats: lossy int16 write, CT-only SOP class, no multi-frame, single-gap z-spacing. | -| **YAML model builder** | ✅ Works; ⛔ cannot yet replace `models/` | `examples/Segmentation/UNet.yml` builds a model **identical** to the Python `UNet` (1,934,299 params). Registry (13 types) is too small and custom-`forward` models can't be expressed — see §4. | +**Modernization (PRs #6–#9):** OME-Zarr + DICOM dataset backends, the declarative YAML model builder, performance/elegance pass, and docs. ---- +**Audit follow-up (PR #10):** -## 2. OME-Zarr (`konfai/utils/ome_zarr.py`) +- **#1** per-epoch augmentation re-sampling — `DataAugmentation.reset_state()` clears per-case sampling so params are re-drawn each epoch. +- **#2** `load_state_dict` warm-start checks `isinstance(child, (Conv, Linear))`, guards on the checkpoint key, and `continue`s instead of `return`ing. +- **#10** `Crop.transform_shape` reuses a persisted box to skip the full-volume read (mitigated; a fully-lazy variant is out of scope). +- **#14** `Accumulator.assemble` seeds the output from the first present patch and raises a typed `PatchError`. +- **#15** `update_scheduler` raises `ConfigError` on an empty schedule. +- **#16** SSIM/FID/LPIPS declared as `konfai[ssim]`/`konfai[fid]`/`konfai[lpips]` extras + an actionable `MeasureError` at construction. +- **#17** `LinearVAE` rebuilt on a real `LatentDistribution` bottleneck; `VAE` documented as a deterministic autoencoder. +- **#18** `Representation.Adaptation` sets `requires_grad` in `__init__` (forward is pure again). +- **#19** debug blocks (`Print`/`Write`/`Exit`) gated + documented; `Write` takes an explicit `path`. +- **DICOM hardening:** non-destructive `write_dicom_series`; multi-frame + irregular-spacing detection. +- **`os.environ` device control-flow removed:** `Network.to` threads an explicit GPU-index counter. +- **#11** reclassified as a reserved future feature (documented). -**Verified working:** write→read round-trips for 3-D `(C,Z,Y,X)`, 2-D `(C,Y,X)`, patch/slice reads, axis canonicalization to channel-first, and scale/translation↔Spacing/Origin conversion are all correct against **zarr 3.2.1** (a `create_array`/`create_dataset` shim handles v2/v3). Direction (not representable in NGFF) round-trips via a proprietary `konfai` attrs key. Covered by `tests/unit/test_imaging_roundtrip.py`. +Each fixed item landed with a regression test where the behaviour is testable (`tests/unit/test_audit_fixes.py`, `test_patching.py`, `test_named_forward.py`, `test_imaging_formats.py`, `test_early_stopping.py`). -**Issues** -- ⚠️ **Reinvents ngff-zarr.** `ngff_zarr` is imported and probed (`_NGFF_ZARR_AVAILABLE`) but **never called** anywhere; `_parse_ngff_axes/_scale/_translation/_canonical_shape/select_level` hand-parse `multiscales[0]` — exactly what ngff-zarr exposes as typed objects. ngff-zarr *can* read KonfAI's output (`from_ngff_zarr` returns correct dims/scale/translation/data), so the formats are compatible and a migration is low-risk. → backlog PR (§7); an interop regression test guards the format compatibility in the meantime. -- 🔵 **Dead code:** `read_ome_zarr_slice` (documented as the "primary entry point") and `select_level` are used only by tests, never by `dataset.py` (which uses `read_ome_zarr_data_slice`). Multiscale pyramids and the time axis are effectively unsupported despite the API surface; write only emits a single level. -- 🟡 **Standards/interop gaps:** `Direction` is invisible to any standards-compliant reader (only the proprietary key carries it); store is written as zarr-v3 storage carrying an NGFF "0.4" (a v2-era) version label — a hybrid that happens to read back. -- 🟢 **Error-type consistency (fixed):** `read_ome_zarr_data_slice` raised a bare `ValueError` on slice-arity mismatch → now `DatasetManagerError` (matches AGENTS.md rule), with a test. -- 🔵 `_open_level` re-opens the group and re-reads attrs on every slice/info call — no store/group caching; costly for remote stores the docstring advertises. +## Open backlog — confirmed bugs, not yet fixed -## 3. DICOM (`konfai/utils/dicom.py`) +Severity: 🔴 high · 🟠 medium · 🟡 low. Each is a small, well-scoped, test-backed fix. -**Verified working:** series discovery by `SeriesInstanceUID`, position-based slice ordering, geometry extraction `(origin, spacing(x,y,z), direction(9))`, HU rescale, lazy per-slice patch reads, and write. **Cross-validated:** KonfAI's reader output is **byte-identical to SimpleITK's `ImageSeriesReader`** for origin/spacing/direction/data, including a **left-handed (feet-first, z-down) input**, which both normalize to the same right-handed frame with a flipped array — i.e. KonfAI is consistent with the reference DICOM reader. +| # | Bug | Sev | Where | +|---|---|---|---| +| 3 | VAE latent uses `torch.rand_like` (U[0,1]) where the reparameterization trick needs `randn_like` (N(0,1)) | 🔴 | `network/blocks.py` | +| 4 | Early-stopping/BEST score keys on a leaked loop variable, so it scores the **EMA** model when EMA is on, not the base model | 🔴 | `trainer.py` | +| 5 | Explicit `Standardize` `mean`/`std` crashes (`torch.tensor([list])` → nested tensor, then the stringly-typed `Attribute` reparse fails) | 🟠 | `data/transform.py` | +| 6 | `Rotate` feeds degrees straight to `cos/sin` with no `deg2rad` (a "90°" rotation is 90 rad) | 🟠 | `data/augmentation.py` | +| 7 | CLI `-tb/--tensorboard` is silently dropped for predict/eval (`tb` vs `tensorboard` param name) | 🟠 | `main.py` / `predictor.py` / `evaluator.py` | +| 8 | `Unsqueeze.forward(*tensor)` errors on a tensor input (`torch.unsqueeze` gets a tuple) | 🟠 | `network/blocks.py` | +| 9 | `ResampleToShape` mutates its own config across cases (`new_shape = self.shape` aliasing) | 🟠 | `data/transform.py` | +| 12 | `Select.forward` squeezes by index, not by size | 🟡 | `network/blocks.py` | +| 13 | `ITK._open_transform` double-appends displacement-field transforms (currently in dead code) | 🟡 | `utils/ITK.py` | -**Issues** -- 🟡 **Lossy/limited write:** floats are quantised to int16 via a derived slope/intercept; the SOP class is **always `CTImageStorage`** even when `Modality` is `OT`; `ImageType` is hardcoded `DERIVED/PRIMARY/AXIAL`. -- 🟢 **Edge cases (hardened):** inter-slice spacing still uses the first gap (to stay byte-identical to SimpleITK) but the **whole series is now checked for uniform spacing** and raises `DatasetManagerError` on irregular series; **multi-frame / enhanced DICOM** (`NumberOfFrames > 1`) is now rejected with a clear error instead of being mis-stacked. Single-slice series still default z-spacing to 1.0 mm. Tests added. -- 🟢 **Destructive write (fixed):** `write_dicom_series` previously deleted **all** `*.dcm` in the target directory; it now removes only the slices it owns (its `NNNNNN.dcm` naming), preserving unrelated DICOM files. Test added. -- 🟢 **Error-type consistency (fixed):** `read_dicom_series_slice` raised `ValueError` on slice-arity mismatch → now `DatasetManagerError`, with a test. -- 🔵 **Performance:** a single patch read parses the whole series' headers 3–4× (`get_dicom_info` + two `_select_series_files` + sorts); `file_to_data_statistics` calls the slice reader per z-slice → O(N²) header reads for one volume's statistics. `discover_series` swallows all exceptions (incl. corrupt-but-present series → misleading "no DICOM found"). -- 🔵 Dead imports: `DicomSequence`, the empty `if TYPE_CHECKING: pass`. +## Cross-cutting (deferred refactors, each PR-sized) -## 4. YAML model builder (`konfai/utils/model_builder.py`) +- **Typed `Attribute` sidecar.** Every value is `str()`'d and geometry reparsed via `np.fromstring(s[1:-1])` — lossy and the root cause of #5. The active bug is patched; the serialization rework is its own change. +- **Perceptual-loss duplication.** `IMPACTReg`/`IMPACTSynth`/`SAM_Perceptual` are ~70-line near-duplicates with *different* mask/normalization — needs characterization tests first. +- **Smells:** pervasive mutable default args (masked by the Config layer); `Criterion.forward` typed `-> Tensor` but several subclasses return tuples; `ITK.py` is largely dead; the `interactive`/`remove` config modes are unreachable. +- **Tests still thin** on DDP/multi-rank, RESUME, and explicit-arg transforms (where #5/#6 slipped through). -**Verified working:** safe-registry build, `${param}`/`$multiply`/`$object` resolution, nested routed graphs, and `add_module` routing all function; all 25 `test_model_builder.py` tests pass. **Equivalence:** building `examples/Segmentation/UNet.yml` yields a model with **the same parameter count (1,934,299) and forward output** as the Python `UNet` configured identically. +## Claims downgraded by verification -**Can it replace `konfai/models/`? Not yet — here is exactly what's missing:** -1. **Registry is too small (13 types):** `Conv/ConvTranspose/MaxPool/AvgPool/Conv1d-3d/Softmax/Identity/ArgMax/ConvBlock/ResBlock/Concat`. Shipped models also need (all exist in `blocks.py`/torch but are **not** registered): `Linear`, `Add`, `Multiply`, `Attention`, `Upsample`, `View`/`Select`/`Subset`, `LatentDistribution`, `NormalNoise`, `Const`, `Unsqueeze`/`Permute`/`ToChannels`/`ToFeatures`, plus norm/activation factories. The object registry has only `BlockConfig` — the heavily-used `DownsampleMode`/`UpsampleMode`/`NormMode` enums are not expressible. -2. **Custom-logic models can't be expressed at all.** ~70–80% of `models/` is pure `add_module` wiring (UNet, NestedUNet, ResNet, GAN, VAE — migratable once the registry grows), but the rest carries genuine Python `forward`/sampling/training logic the builder has no construct for: `ddpm` (7 custom methods — diffusion schedule/sampling), `diffusionGan`, `cStyleGan`, `registration` (spatial-transformer grid math), `convNeXt` (LayerNorm/DropPath/LayerScaler), `representation`. - -**Verdict:** the builder is a sound, safe foundation that **complements** `models/` today and can **replace the feed-forward subset** after (1). A migration plan is in §7. Other notes: no looping/recursion construct, so deep nets are hand-unrolled and verbose (UNet.yml is 196 lines vs a few in Python); `register_module` mutates a **process-global** registry with no isolation/unregister (a hazard for the long-lived app server); the YAML feature has **no user-doc statement** of these limitations. - ---- - -## 4b. `konfai-apps`, packaged models & external bundles - -`konfai-apps` is a separate package (local runner `KonfAIApp`, remote `KonfAIAppClient` + FastAPI `app_server.py`, and `Local`/`HuggingFace`/`Remote` repository adapters). The `apps/*` bundles (`totalsegmentator`, `mrsegmentator`, `impact_synth`) are thin CLI wrappers that resolve a HuggingFace app and call `KonfAIApp.pipeline()`. The package boundary is **clean** (core never imports apps; apps use only the documented `konfai` public API). - -- 🔴 **Trust model is undocumented (security).** Resolving an app **copies its `.py` files into the working directory and imports them**, and an app's `requirements.txt` is installed via a `pip install` subprocess (`app_repository.py`). A HuggingFace/remote app therefore runs **arbitrary code and arbitrary dependency installs** on the host. This is inherent to the "packaged model = code + weights + config" design and is not a fixable bug — but it **must be stated plainly** so users only resolve apps from sources they trust. Documented in AGENTS.md §"Apps & packaged models"; a future hardening PR may add checksum/lockfile pinning. -- 🟠 **GPU scheduler race (server).** In auto device mode the check-then-acquire of a free GPU is non-atomic (`app_server.py`), so two concurrent jobs can select the same device under load. Backlog (PR G). -- 🟡 **Result TTL hardcoded** (~120 s) → a slow client download can 404 mid-stream; broad `except Exception` across server/repository; `download_/install_evaluation` ≈ `…_uncertainty` (duplication); some French comments / missing SPDX headers in `cli.py`/`__init__.py`. Backlog (PR G). -- 🟡 **#11 (MC-dropout ignored locally)** lives here: `LocalAppRepository.install_inference` never consumes the `mc` count, so `--mc` has no local effect. Backlog (PR B). - -## 4c. Model zoo (`konfai/models/`) - -The zoo (UNet, NestedUNet/UNet++, ResNet, ConvNeXt, GAN/CycleGAN, VAE, DDPM, DiffusionGAN, cStyleGAN, VoxelMorph registration, Representation) is the least-tested surface. Findings (backlog — PR C; they touch model behavior and need careful, test-backed fixes): - -- 🔴 **`LinearVAE` is non-functional** (`vae.py`): hard-coded dims (`23343→5→23343`), commented-out modules, and **no reparameterization**. The main `VAE` also performs no latent sampling (no reparameterization trick) — only the `LatentDistributionZ` block (fixed in #3) draws noise where wired. -- 🔴 **`forward` has side effects (non-reproducible / not thread-safe).** `DiffusionGAN.UpdateP` mutates `_it`/`p` and reads `measure` inside `forward` and **never resets `_it`**; `Representation.Adaptation` calls `requires_grad_()` on every `forward`. Forward passes should be pure. -- 🟠 **Debug blocks shipped in the production block library:** `Print`, `Write` (writes `.mha` to disk), and `Exit` (raises) live ungated in `blocks.py`. They should be gated behind a debug flag or removed. -- 🟠 **Heavy duplication:** `Generator`/`Discriminator`/`ResBlock`/sinusoidal `TimeEmbedding` are re-implemented across `gan.py`/`diffusionGan.py`/`ddpm.py`; `representation.py` defines its own `ConvBlock` duplicating `blocks.ConvBlock`. -- 🟡 Hardcoded dims (`registration.Rigid` 512×512); `get_torch_module` typed `-> Module` but returns a class; `MSE.forward` ignores `*targets` (`ddpm.py`); `NormalNoise.forward` shape mismatch. -- **Model→YAML migration matrix** (for the model-builder, §4): **UNet / NestedUNet / ResNet are 100 % migratable today**; UNet++/GAN/VAE/VoxelMorph are partial (need registry growth + channel-math/factory constructs); **ConvNeXt / DDPM / DiffusionGAN / cStyleGAN / Representation are not migratable** without a custom-`forward` construct the builder does not have. - -## 4d. Metrics & criteria runtime - -- 🟢 **Undeclared optional metric deps — fixed.** `SSIM` (`scikit-image`), `FID` (`scipy` + `torchvision`), and `LPIPS` (`lpips`) were lazily imported with no extras and no actionable error → a raw `ImportError` mid-run. Now declared as `konfai[ssim]` / `konfai[fid]` / `konfai[lpips]` extras and routed through `_require_optional(...)`, which raises a `MeasureError` **at criterion construction** stating exactly what to install. Covered by a regression test. -- 🟢 **`update_scheduler` crash on an empty schedule — fixed.** It previously raised a raw `NameError` when the scheduler dict was empty (a misconfiguration); it now raises a clear `ConfigError`. The past-last-window case already clamps to the last scheduler (not a crash — see §8). Both cases are covered by tests. -- 🟠 **Copy-paste drift in perceptual losses:** `IMPACTReg._compute` / `IMPACTSynth._loss_compute` / `SAM_Perceptual._compute` are ~70-line near-duplicates with **different** mask-resampling/normalization (SAM hardcodes 2-D `[512,512]`). Extract a shared helper. Backlog (PR E). -- 🟠 **`PerceptualLoss.forward` does `del os.environ['device']`** without restore, and the process-global `models_register` is never cleared (leak in the long-lived server). Backlog (PR E). - -## 4e. Patch reconstruction & overlap blending - -The patch-assembly path (`Accumulator`, `PathCombine`/`Mean`/`Cosinus`) is covered by `tests/unit/test_patching.py`. Findings: - -- 🟢 **#14 fixed** (typed `PatchError`, seed from first present patch — see §5). -- 🟡 **Blending windows are not a partition of unity at volume borders.** Summing the `Mean`/`Cosinus` weight window over a clamped tiling (from `get_patch_slices_from_shape`) yields per-voxel weight sums ranging from ~0.25 (outer corners) to 1.0 (centre), not a uniform 1.0. The windows assume **uniform** overlap multiplicity, but the last patch in each dim is clamped to the volume edge, so border voxels are **under-weighted** (attenuated) in the assembled output. Likely acceptable for prediction averaging, but overlap-blended assembly does **not** exactly reconstruct a known field at borders. Out of scope for the current fix set. The tests assert the safe invariants (exact reconstruction without blending; windows bounded in `[0,1]`; unit at centre; cosine tapers more than mean). - ---- - -## 5. Confirmed bugs (adversarially verified) - -Severity: 🔴 high · 🟠 medium · 🟡 low. "Status" is whether the bug is fixed or queued as a follow-up PR; fixes are limited to clear, well-covered changes (no large refactor). - -| # | Bug | Sev | Verdict | Status | -|---|---|---|---|---| -| 1 | **Per-epoch augmentation re-sampling is broken.** `DataAugmentation.state_init` short-circuits when an index is already in `who_index`, and `who_index` is **never cleared**, so random transform params are frozen for the object's lifetime. `docs/.../training.md:97` promises re-sampling each epoch — false. (`augmentation.py:199-202`) | 🔴 | confirmed (repro: identical params across 3 epochs) | **Fixed** — `reset_state()` clears `who_index`/`shape_index` per case; `DatasetManager.reset_augmentation` calls it before each `state_init`; test added. | -| 2 | **`load_state_dict` warm-start checks the wrong object.** `isinstance(module, torch.nn.Linear)` should be `isinstance(child, …)` (so resized `Linear` never warm-starts), and an early `return` inside the per-child loop aborts loading the rest of the subtree. Live checkpoint-load path. (`network.py:898,913`) | 🔴 | confirmed | **Fixed** — checks `isinstance(child, (Conv, Linear))`, `continue` instead of `return`, guarded on the checkpoint key being present; resized-layer + sibling-load test added. | -| 3 | **VAE latent uses uniform noise.** `LatentDistributionZ.forward` uses `torch.rand_like` (U[0,1]) where the reparameterization trick needs `torch.randn_like` (N(0,1)); the KL term assumes a unit Gaussian. One-token fix. (`blocks.py:442`) | 🔴 | partially-confirmed (claim conflated `NormalNoise`, which is fine) | PR | -| 4 | **Early-stopping/BEST score uses the EMA model when EMA is on.** `trainer._log` reuses a leaked loop variable `label` (= `_EMA` after the loop) to key `measures`, so the returned score is the EMA model's, not the base model's. (`trainer.py:469,519-520`) | 🔴 | confirmed | PR | -| 5 | **`Standardize`/explicit `mean`/`std` crash.** `torch.tensor([self.mean])` on a list makes a nested tensor; the stringly-typed `Attribute` reparse (`np.fromstring(s[1:-1])`) then fails. Only the computed-scalar path works. (`transform.py:281-288`) | 🟠 | confirmed (repro) | PR | -| 6 | **`Rotate` treats degrees as radians.** Angles sampled in `[a_min,a_max]` (default `[0,360]`) are fed straight to `cos/sin` with no `deg2rad`; a "90°" rotation is actually 90 rad. Used by `diffusionGan`. (`augmentation.py:315-331,90-135`) | 🟠 | confirmed (repro) | PR | -| 7 | **CLI `-tb/--tensorboard` is silently dropped for predict/eval.** CLI dest is `tensorboard`, but `predict()`/`evaluate()` declare the param as `tb`; `run_distributed_app` filters kwargs to the signature, so TensorBoard cannot be enabled from the CLI for PREDICTION/EVALUATION. `train()` works. (`main.py:108` vs `predictor.py:1075`/`evaluator.py:528`) | 🟠 | confirmed | PR | -| 8 | **`Unsqueeze.forward(*tensor)` errors on a tensor input** (`torch.unsqueeze` gets a tuple). Used in `resnet`/`convNeXt`. (`blocks.py:254`) | 🟠 | confirmed (repro) | PR | -| 9 | **`ResampleToShape` mutates its own config across cases.** `new_shape = self.shape` (alias) with sentinel-0 substitution writes the first case's dims into the shared instance, leaking into later cases. (`transform.py:459,466`) | 🟠 | partially-confirmed (`ResampleToResolution` is fine) | PR | -| 10 | **`Crop.transform_shape` does a full volume read** (`read_data` + percentile) at `DatasetManager.__init__`, violating the never-load-full-volume rule. Acknowledged `TODO(perf)`. (`transform.py:1043-1055`) | 🟠 | confirmed | **Mitigated** — the crop box is content-dependent, so the *shape* genuinely needs the data; now a persisted box is reused to skip the read, the box parsing is de-duplicated, and the constraint is documented. A fully-lazy variant (deferring patch planning) is out of scope. | -| 11 | ~~MC-dropout count ignored in local inference.~~ **Not a bug** — `number_of_mc_dropout` is a *reserved* parameter for a planned MC-dropout feature (stochastic passes for models with dropout layers); it is plumbed through but intentionally not applied yet. Documented in `install_inference`. (`app_repository.py`) | — | reclassified | **Reclassified** (future feature, documented). | -| 12 | **`Select.forward` squeezes by index, not size** (`enumerate(range(...))` → tests `i==1`, not `shape[i]==1`). Dead-ish but wrong. (`blocks.py:374`) | 🟡 | confirmed | PR | -| 13 | **`ITK._open_transform` double-appends displacement-field transforms** (append inside the branch + the unconditional append) → double application in `compose_transform`. Currently in dead code. (`ITK.py:83,88`) | 🟡 | confirmed (repro) | PR | -| 14 | **`Accumulator.assemble` can `UnboundLocalError`** if patch index 0 was never filled (`result` only bound inside the index-0 branch). Live path fills index 0, so latent. (`patching.py:184-203`) | 🟡 | partially-confirmed (repro) | **Fixed** — seeds the output from the first present patch; raises a typed `PatchError` when nothing was added; tests added. | -| 15 | **`update_scheduler` raises a raw `NameError` on an empty scheduler dict** (a misconfiguration). (`network.py:458-473`) | 🟡 | confirmed | **Fixed** — raises `ConfigError`; past-last-window already clamps to the last scheduler (not a crash, see §8); tests added. | -| 16 | **Optional metric deps undeclared** (SSIM/`scikit-image`, FID/`scipy`+`torchvision`, LPIPS/`lpips`) → raw `ImportError` mid-run. (`measure.py`) | 🟠 | confirmed | **Fixed** — `konfai[ssim]`/`konfai[fid]`/`konfai[lpips]` extras + `_require_optional` raising an actionable `MeasureError` at criterion construction; test added. | -| 17 | **`LinearVAE` non-functional** (hard-coded dims, no reparameterization); `VAE` no latent sampling. (`vae.py`) | 🔴 | confirmed (read) | **Fixed** — `LinearVAE` rebuilt as a parameterized linear VAE with a real `LatentDistribution` bottleneck (mu/log_std/z, KL-ready named outputs); `VAE` documented as a deterministic autoencoder. Forward/stochasticity test added. | -| 18 | **Model `forward` side effects** (`DiffusionGAN.UpdateP` mutates `_it`, never resets; `Representation.Adaptation` `requires_grad_()` each forward). | 🔴 | confirmed (read) | **Fixed** (`Adaptation`: `requires_grad` moved to `__init__`, test added). `UpdateP` is intrinsic ADA running state (documented; registering buffers would break existing checkpoints). | -| 19 | **Debug blocks shipped ungated** (`Print`/`Write`/`Exit` in `blocks.py`). | 🟠 | confirmed (read) | **Fixed** — grouped and documented as debug-only; `Write` no longer writes to a hardcoded path (explicit `path` arg) and warns on use. | - -Each fixed row landed **with** a regression test where behavior is testable. Two cross-cutting refactors remain deliberately deferred (§6): **de-duplicating the IMPACT/SAM perceptual losses** (untested challenge-loss code with subtle per-variant mask/normalization differences — needs characterization tests first) and a **typed `Attribute` sidecar** (deep serialization-backbone change; its active bug #5 is already fixed and the byte-identical geometry guarantees are test-asserted). Both are PR-sized efforts of their own. - -## 6. Cross-cutting audit (by dimension) - -**Performance** -- Confirmed perf-positive work already in tree (in-memory best-checkpoint tracking, O(1) index-cache, once-per-batch predict logging, `get_names` cache). -- DICOM statistics are O(N²) header reads (§3); OME-Zarr re-opens the group per access (§3); `Crop` reads full volumes (#10); `data_manager.__getitem__` recomputes `needs_full_load` every sample (cheap after caching but allocates). - -**Elegance / duplication / unnecessary complexity** -- **Stringly-typed `Attribute`** (`dataset.py:68-95`): every value is `str()`'d and geometry reparsed via `np.fromstring(s[1:-1], sep=" ")` — lossy, locale/printoptions-sensitive, and the root cause of bug #5. `startswith`-based key counting also cross-contaminates prefix-sharing keys (e.g. `Spacing` vs a `SpacingExtra` metadata key). -- **Duplicated SimpleITK transform serialization** appears 3× (`dataset.py` H5 write / Sitk read / `read_transform`) with a latent `UnboundLocalError` on unknown transform types. -- **`os.environ` as control-flow state (largely fixed):** `Network.to`'s GPU-index counter no longer lives in `os.environ['device']` — it is threaded explicitly through the recursion as a mutable box that resets per top-level call (also fixing a latent model/EMA mis-placement when the counter leaked across calls), and `PerceptualLoss.forward`'s `del os.environ['device']` reset is gone. Remaining: `get_layers` still appends to `KONFAI_DEBUG_LAST_LAYER`, but it is a benign opt-in debug trace (only active if the var is pre-set) read by nothing for control flow. -- **Name-string coupling:** channel tracing keys `ToChannels/ToFeatures` and `ReduceLROnPlateau` by `__class__.__name__` rather than `isinstance` — renames silently break routing. -- **`named_forward` nested out-branch propagation** (`network.py:701-708`) is the highest-complexity, untested, string-surgery (`split('.')`, `;accu;`) function in the codebase. -- Dead config machinery: the `interactive` and `remove` `KONFAI_CONFIG_MODE` modes are read but never set anywhere (~⅓ of `get_value`/`_get_input*` is unreachable). `check_konfai_install`/`KonfAIPackagesError`/`_KONFAI_DEPS` (~60 lines) have no callers. -- **`ITK.py`** is ~13/15 functions unused in-repo (large untested dead surface). - -**Architecture / API** -- `Criterion.forward` is typed `-> torch.Tensor` but many subclasses return tuples — the base contract is effectively a lie; consumers `isinstance`-check. -- `SitkFile.file_to_data` returns an lxml `Element` for `.xml` (violates the declared `tuple[np.ndarray, Attribute]`). -- Backend ABC is not honored uniformly (`SitkFile.get_names/get_group` raise `NotImplementedError`; `OmeZarr/Dicom.get_names` ignore the `group` arg). -- Two different `run_distributed_app` decorators (apps vs core) with different semantics are in scope together — confusing; rename one. -- `read_data` opens H5 with `read=False` (writable handle) even on read paths — concurrency/perf hazard. - -**Typing** -- Pervasive `image: sitk.Image = None` / `h5_group: h5py.Group = None` non-Optional defaults (typing lies); `get_torch_module`/`get_conv` annotated `-> torch.nn.Module` but return classes/`None`; `Network.forward` annotated `-> torch.Tensor` but returns the last yielded value. - -**Mutable default arguments** — pervasive across the public config API (`schedulers={...}`, `outputs_criterions={...}`, `alias=[[],[],[]]`, `patch_size=[128,128,128]`, `ModelPatch()` defaults, `data_augmentations={'default|Flip': Prob(1)}`). Mostly masked because the Config layer overrides them and they aren't mutated — but a real footgun and Ruff-B006 trap. - -**Missing tests** (highest value first) -- The patch math (`Accumulator`, `PathCombine/Mean/Cosinus`, overlap blending, 2.5-D `extend_slice`) has **no direct unit tests** — the riskiest reconstruction code is only exercised indirectly. -- `named_forward` branch routing has no direct test. -- Transforms are largely untested with explicit args (`Standardize`/`Clip` percentile/lists — bug #5 slipped through), `Rotate`/affine augmentations untested (bug #6), DDP/multi-rank paths untested. -- No test that `--tensorboard` reaches predict/eval (bug #7); no RESUME test. - -**Documentation** -- `docs/source/konfai.utils.rst` omits `dicom`, `ome_zarr`, `model_builder`, `runtime`, `errors`; there is **no `konfai.models.rst`** at all. `docs/.../apps.rst` autodocs `konfai_apps`, which Read the Docs doesn't install → build warnings/empty output. `examples/README.md` lists a non-existent `UNetpp.py`. `installation.md` under-describes the `imaging` extra (omits pydicom/zarr/ngff-zarr). - -**Packaging / tooling** (see also the `pixi`/`ruff` discussion) -- **Two divergent dev environments:** pip `[project.optional-dependencies].dev` has Sphinx but **no ruff/mypy/build/imaging deps**; the Pixi `dev` feature has those but **no Sphinx**. `pip install -e .[dev]` cannot run ruff/mypy or imaging tests. -- **`ruff==0.15.2` pinned in 4 places** (Pixi `dev` + `lint` features, pre-commit, CI) that must stay in lockstep. Pixi-idiomatic fix: declare ruff + tasks once in the `lint` feature and compose `dev = { features = ["dev","lint"] }`; prefer conda `dependencies` over `pypi-dependencies` for conda-forge tools. -- **`pixi run test` does not run `konfai-apps/tests`** (root `testpaths=['tests']`); only the apps CI does. Add a Pixi task for the apps suite. -- CI installs `.[dev]` (no imaging extras) then runs `pytest tests`, so the imaging happy-paths run **only locally** — add an imaging-extra CI job. - -## 7. Recommended PRs (prioritized) - -**Status.** The modernization stack (performance/elegance, OME-Zarr+DICOM backends, YAML model builder, docs) merged to `main` as PRs #6–#9. The audit-follow-up work fixed bugs #1, #14, #15, #16 with tests and extended this document and `AGENTS.md` to cover the whole repository (apps trust model, model zoo, metrics, patch blending). The list below is the remaining prioritized backlog. - -1. **Correctness bug-fix PR (high):** bugs #1–#4 (aug re-sampling, `load_state_dict`, VAE noise, EMA early-stopping) — each with a regression test. Small diffs, high impact. *(#1 done; #3/#4 done in #6; #2 remains → backlog PR B.)* -2. **Transform/augmentation correctness PR (medium):** bugs #5, #6, #9 (`Standardize` explicit stats, `Rotate` deg→rad, `ResampleToShape` aliasing) + a typed/robust `Attribute` serialization (fixes the #5 root cause) — with tests for explicit-arg transforms. -3. **CLI/apps PR (medium):** bug #7 (`tb`→`tensorboard`), bug #11 (`--mc` local), + a `--tensorboard reaches predict/eval` test and a RESUME test. -4. **OME-Zarr ↔ ngff-zarr migration (medium):** replace the hand-rolled NGFF parse/write with `ngff-zarr` (`from_ngff_zarr`/`to_ngff_zarr`), keep zarr for chunked array access; drop dead `read_ome_zarr_slice`/`select_level` or wire multiscale into the live path. Interop already verified, so risk is bounded; gate behind the existing round-trip + interop tests. -5. **Model-builder registry expansion (medium):** register the missing blocks/torch layers + enum objects (§4.1), add per-build registry isolation/unregister, document the custom-`forward` limitation. Then migrate ResNet/NestedUNet/GAN/VAE to `.yml` one at a time, each guarded by a param-count/forward-equivalence test like the one added here. -6. **Patch-math test PR (high value, no code change):** direct unit tests for `Accumulator`/overlap blending/`PathCombine`/2.5-D before any refactor; then bugs #13, #14 with coverage. -7. **Tooling PR (low):** unify the dev environments and the single-source ruff pin; add a Pixi task + CI job for `konfai-apps/tests` and an imaging-extra CI job; clean `ITK.py` dead code; fix the docs module reference. -8. **DICOM robustness PR (medium):** irregular-spacing detection, multi-frame rejection/support, configurable SOP class, non-destructive write — each with a test. - -## 8. Claims downgraded by verification -- **DICOM "geometrically wrong" on flipped-z — overstated.** KonfAI is **byte-identical to SimpleITK**; left-handed → right-handed normalization is correct DICOM behavior, not a bug. (Now locked by a regression test.) -- **`Config.get_value` `next()` StopIteration — overstated.** Real in isolation but the branch is unreachable given how `key_tmp` is built; defensive-only. -- **VAE noise / `ResampleToResolution` / `Accumulator` — partially confirmed:** the defect is real but narrower than first claimed (sibling classes/paths are fine); see #3, #9, #14. - -## 9. Changes applied - -**Modernization stack (PRs #6–#9, merged):** -- `konfai/utils/ome_zarr.py`, `konfai/utils/dicom.py`: bare `ValueError` on slice-arity mismatch → `DatasetManagerError` (AGENTS.md error-type rule), with tests. -- `tests/unit/test_imaging_roundtrip.py`: DICOM↔SimpleITK geometry consistency, left-handed normalization, OME-Zarr↔ngff-zarr interop, non-identity-direction backend round-trip, error-type checks. -- `tests/unit/test_yaml_model_equivalence.py`: `UNet.yml` builds + forwards, and matches the Python `UNet` parameter count. -- `AGENTS.md`: rewritten/expanded from the whole-codebase + paper understanding. - -**Audit-follow-up:** -- **#1** `DataAugmentation.reset_state()` clears per-case sampling; `DatasetManager.reset_augmentation` calls it before each `state_init` so augmentation parameters are re-drawn every epoch. (`augmentation.py`, `patching.py`) -- **#2** `load_state_dict` warm-start checks `isinstance(child, (Conv, Linear))` (not the parent), guards on the checkpoint key, and `continue`s instead of `return`ing so siblings of a resized layer still load. (`network.py`) -- **#10** `Crop.transform_shape` reuses a persisted box to skip the read, de-duplicates the box parsing (`_parse_box`), and documents the content-dependent-shape constraint. (`transform.py`) -- **#14** `Accumulator.assemble` seeds the output from the first present patch and raises a typed `PatchError` (new in `errors.py`) when nothing was added, instead of `UnboundLocalError`. (`patching.py`, `errors.py`) -- **#15** `update_scheduler` raises `ConfigError` on an empty schedule instead of `NameError`. (`network.py`) -- **#16** SSIM/FID/LPIPS dependencies declared as `konfai[ssim]`/`konfai[fid]`/`konfai[lpips]` extras; `_require_optional` raises an actionable `MeasureError` at criterion construction. (`pyproject.toml`, `measure.py`) -- **#17** `LinearVAE` rebuilt as a parameterized linear VAE on a real `LatentDistribution` bottleneck; `VAE` documented as a deterministic autoencoder. (`vae.py`) -- **#18** `Representation.Adaptation` sets `requires_grad` in `__init__` instead of every forward; `DiffusionGAN.UpdateP` documented as intrinsic ADA state. (`representation.py`, `diffusionGan.py`) -- **#19** `Print`/`Write`/`Exit` grouped and documented as debug-only; `Write` takes an explicit `path` and warns. (`blocks.py`) -- **#11** reclassified as a reserved future feature (documented in `app_repository.py`). -- **DICOM hardening:** non-destructive `write_dicom_series`; multi-frame and irregular-spacing detection in `extract_geometry`. (`dicom.py`) -- **`os.environ` device control-flow removed:** `Network.to` threads an explicit GPU-index counter; `PerceptualLoss` drops the `del os.environ['device']` reset. (`network.py`, `measure.py`) -- Tests: `tests/unit/test_patching.py`, `tests/unit/test_named_forward.py`, new cases in `tests/unit/test_audit_fixes.py` (#1, #2, #15, #16, #17, #18), and new DICOM-hardening cases in `tests/unit/test_imaging_formats.py`. -- `AGENTS.md`: added "Apps & packaged models" (incl. the trust-model warning), a metrics/criteria note, and the model→YAML migration matrix. - -Two cross-cutting refactors are deliberately deferred to their own test-backed PRs: de-duplicating the IMPACT/SAM perceptual losses (untested challenge-loss code with subtle per-variant differences) and a typed `Attribute` sidecar (deep serialization change; its active bug #5 is already fixed). +- **DICOM "geometrically wrong on flipped-z" — overstated.** KonfAI is byte-identical to SimpleITK; left-handed → right-handed normalization is correct DICOM behaviour (locked by a regression test). +- **VAE noise / `ResampleToResolution` / `Accumulator` / `Config.get_value` StopIteration — partially confirmed:** the defect is real but narrower than first claimed (sibling classes/paths are fine).