predict normalizes nnU-Net checkpoints with instance z-score regardless of the scheme in their plans.json

[Jinhojeong]

Running an nnU-Net checkpoint through vesuvius.predict normalizes the input with per-volume z-scoring even when the checkpoint's own plans.json declares something else. The plans are loaded and stay in hand. They are just never consulted when the normalization is chosen, and nothing warns, so the run looks normal and the outputs look plausible.

Where it happens, on main at a6c8ad5.

• --normalization defaults to instance_zscore in the CLI parser, vesuvius/src/vesuvius/models/run/inference.py:1165, and the constructor default is the same at :252.
• self.model_normalization_scheme and self.model_intensity_properties both start as None at :319-320.
• Two branches fill them in. The resnet branch sets the scheme at :405, and the train.py checkpoint path sets both at :542-549.
• The nnU-Net branches do not. --model_type nnunet at :419-426 and the auto fallback at :440-448 both call load_model_for_inference, and the model_info it returns carries plans_manager and configuration_manager but no normalization keys, vesuvius/src/vesuvius/utils/models/load_nnunet_model.py:297-308. The only place that loader builds those keys is the HuggingFace train.py branch at :386-387.
• The hf:// path at inference.py:378-384 calls the same loader and lands in the same hole, unless the repo turns out to be a train.py checkpoint, which is the branch tested right after at :387. That is the form vesuvius/docs/inference.md leads with, so it is not an edge case.
• So self.model_normalization_scheme or self.normalization_scheme at inference.py:696 resolves to the CLI default, and VCDataset is constructed with it at :727, inside the call that runs :721-749.
• The scheme actually in use is printed only under verbose, vesuvius/src/vesuvius/data/vc_dataset.py:187, so a default run says nothing about normalization at all. Across 16 non-verbose runs of the checkpoint below, no log contains the string normaliz.

The information is already in hand at that point. configuration_manager.normalization_schemes and plans_manager.foreground_intensity_properties_per_channel are both populated by the loader, and inference.py reads neither.

There is also no way to ask for the right scheme from the outside. The --normalization help lists instance_zscore, global_zscore, instance_minmax and none. Volume does implement a ct scheme that is nnU-Net's CTNormalization verbatim, vesuvius/src/vesuvius/data/volume.py:970-980, but inference.py never passes intensity_props down to VCDataset, so --normalization ct raises at volume.py:234-237. --normalization global_zscore fails from the other side, since global_mean and global_std are read off model_intensity_properties, which is None on this path, inference.py:715-717.

That ct failure is also reported badly, which is worth fixing on its own. The ValueError from Volume is re-raised by VCDataset at vesuvius/src/vesuvius/data/vc_dataset.py:271, and Inferer.infer() catches it at inference.py:1106-1109, prints it, and then falls off the end of the method with an implicit None. main() unpacks that return at :1294 and dies on a TypeError, so the last line printed is cannot unpack non-iterable NoneType object rather than the missing intensity properties. The output directory was created back in the constructor at :352, so what is left on disk is an empty logits directory. The else branch at :1366-1368 that exists to report a None result cannot be reached, because the unpack raises first. The exit status is still 1, so a caller that checks it is not misled.

What it costs, measured with vesuvius.predict itself on the public surface_m7 checkpoint. Its plans.json declares CTNormalization for both 2d and 3d_fullres, with mean 87.5442, std 47.7438 and clipping to [0, 212]. I ran that checkpoint through the wrapper at fixed geometry two ways over 16 volumes of the public Kaggle 892 set, with everything but the normalization held constant. The first run was the wrapper as it comes. The second handed it a block already CT-normalized with the checkpoint's own plans values and asked for --normalization none, which is the only way to get CT normalization through this path today.

Median recall is 0.757 on the default and 0.934 on the CT-normalized input. Median precision is 0.442 against 0.739, and the predicted-positive fraction is 0.284 against 0.241. Recall and precision are higher on the CT side on all 16 volumes, and the predicted-positive fraction is lower on all 16. The median per-volume difference is 0.157 in recall, 0.228 in precision and 0.049 in predicted-positive fraction. Both arms also reproduce outside villa on a plain nnUNetPredictor loop, so this is not something my harness is doing. That loop matches the CT-normalized wrapper run to four decimals on all 16 volumes, and the default wrapper run to 0.0007 in median recall.

Those 16 are not a neutral draw and I would not quote 0.157 as a typical figure. They are 16 volumes out of the sample_00161 to sample_00178 block, all from the located population, and their centred 256-cube means have a median of 132.70 against a plans mean of 87.5442, about 0.95 plans standard deviations up. That is the part of the corpus where the two schemes have the most room to disagree.

How large it gets depends on how far a volume's own statistics sit from the training fingerprint, since that is what instance z-score substitutes for them. The 16 above sit about one sigma up. A volume whose own mean and spread happen to land near the plans values would come out close either way, so a corpus with more than one intensity population can read as though the model behaves differently on each part of it when some of the difference is the normalization. I have a larger comparison that shows that pattern and will add it here if it is useful.

A second defect sits in front of this one on a current stack. load_model calls trainer_class.build_network_architecture at vesuvius/src/vesuvius/utils/models/load_nnunet_model.py:241-248, passing the architecture class name, its kwargs and the required-import list positionally. nnunetv2 2.8.1 declares that staticmethod as (plans_manager, configuration_manager, num_input_channels, num_output_channels, enable_deep_supervision), so the call raises TypeError and :249-252 re-raises it as a RuntimeError. The fallback that would otherwise cover a missing trainer is initialize_network, and it is broken on its own. At :122-124 it runs exec(f"import {i}") over configuration_manager.network_arch_init_kwargs_req_import, whose entries are the names of architecture kwargs, conv_op, norm_op, dropout_op and nonlin for this checkpoint, rather than module paths. nnU-Net's own loader instead pydoc.locates the values those names hold. So an nnU-Net checkpoint does not load against nnunetv2 2.8.1 by either route, and I had to substitute nnU-Net's own construction to run the two arms above. On 0.2.4 the trainer failure is swallowed rather than re-raised, so the run reaches the initialize_network bug instead, which means both halves are live depending on the version. One side effect is that on a current nnunetv2 the normalization default is unreachable, since the run dies before any data is read, and that may be part of why it has gone unreported.

Three changes, and the first one is the normalization fix.

1. Fill the scheme and the intensity properties in the nnU-Net loading path, and pass the properties down to VCDataset in the same change. self.model_normalization_scheme and self.model_intensity_properties can come from configuration_manager.normalization_schemes and plans_manager.foreground_intensity_properties_per_channel, the way the resnet and train.py branches already do for their own checkpoints, but neither value transfers as it stands. normalization_schemes is a per-channel list in nnU-Net's vocabulary, ['CTNormalization'] for this checkpoint, so it has to be mapped to the string Volume expects, which is ct. foreground_intensity_properties_per_channel is keyed by channel index as a string, so the single input channel lives under "0". The VCDataset half is not optional either. Setting the scheme alone turns a silent wrong answer into a crash, because :696 would then resolve to ct while VCDataset is still built without intensity_props at :727, and Volume raises at volume.py:234-237. While in there, ct is worth listing in the --normalization help.
2. Separately from that, warn when the checkpoint's plans declare a scheme other than the one about to be used, rather than taking the CLI default silently. That does not fix a run, but it stops the failure from being invisible, and it is cheap enough to land on its own. The ct crash path deserves the same treatment, so give Inferer.infer() an explicit return in its except branch and let main() reach the None branch it already has.
3. In the nnU-Net loader, call build_network_architecture with the signature the installed nnunetv2 declares, since the call at :241-248 is the pre-2.8.1 order. That is the half a current stack hits, because the trainer resolves and :249-252 re-raises rather than falling back. The fallback is worth repairing in the same pass, since initialize_network at :122-124 imports the kwarg names instead of locating the values they hold.

I am happy to send the repro. It is one public checkpoint, one public volume and the two wrapper runs, printing recall and precision against the same labels both ways. On nnunetv2 2.8.1 the loader has to be worked around first, for the reason above. Say if a PR would be more useful than the report.

[Jinhojeong]

@TAUIL-Abd-Elilah has given me permission to add his rows to this issue, in #193 (comment). He asked only that it stay visible the rows are the broken ones, since that is what they are for, and that is what the rest of this does. None of his data appears here without that consent.

The most useful part is that he reproduced the defect inside his own stack rather than accepting mine. He ran this issue's CT-normalized recipe on his own path, a pre-normalized block with --normalization none, over 4 volumes of his cohort against his cached rows. Median recall moves from 0.803 to 0.940 and median precision from 0.467 to 0.742. That is a different install and a different driver script reaching the same direction and roughly the same size as the 16 volumes above. He has since re-scored his full 60-volume cohort the same way and gets 0.9109 recall and 0.6680 precision, where the same code on the same volumes had been reporting 0.7740 and 0.4459. That also closes something I had left open. When I ran villa's wrapper myself I could reproduce the mechanism but not his numbers, and just under 40 percent of the recall gap stayed unexplained. On the corrected path his 0.9109 and my 0.9138 are 0.003 apart on the same 60 volumes, so the residual was a property of the broken path rather than a second difference between our stacks.

Every row of his that I use below comes from the default path and is wrong in the way this issue describes. He has put a correction at the top of his BENCHMARK.md, marked the affected reference table, and left the old numbers in place rather than deleting them so the size of the error stays checkable. The per-volume recall files for all 855 of his scored volumes are still from the default path and are labelled as such. He also found and reported villa's --bbox problems himself in #1241 and #1247, and he reproduced this one within hours of my raising it, so none of this should read as a complaint about his work.

What his rows add is the population dependence, which this issue asserts and does not show. The 16-volume ablation above is all located volumes whose centred 256-cube means sit about 0.95 plans sigma above the training fingerprint, which is near the end of the range where the two schemes disagree most. His cohort and a matched non-located draw cover the other end.

On his 60 located volumes my plans-normalization recall runs a median 0.134 per volume above his published rows and is higher on all 60, with medians of 0.914 against 0.774. I then drew 60 non-located volumes and ran them at his geometry, 58 of which carry labelled sheet inside the scored crop. There the same comparison is 0.002 with mine on the low side of it and higher on 26 of the 58, with medians of 0.923 against 0.924. Recall is the only endpoint I compare on that side, because the precision field in his per-volume files is invalid for an unrelated reason that he documents himself.

The intensities line up with the split. The centred 256-cube mean has a median of 134.95 over the located 60 and 70.57 over the non-located 60, against the plans intensity mean of 87.5442 and standard deviation 47.7438. One population sits 0.99 plans sigma above the fingerprint and the other 0.36 below it. Instance z-score puts a volume's own mean and spread where the training ones belong, so the two schemes converge as a volume approaches the fingerprint and separate as it leaves. A corpus holding more than one intensity population will read as though the model behaves differently on each part of it, when some of what is being measured is the normalization.

That is the practical cost of the default. His benchmark reported 0.777 recall on 174 located volumes and 0.918 on 681 non-located ones. He had ruled out leakage, label artifact, fused geometry and density and concluded the gap was unexplained, which is what he says in the comment linked above. On the correctly normalized path the same split gives 0.914 and 0.923 on my arms.

Two limits on that. My non-located arm compares my plans path against his default-path rows rather than running both schemes over those volumes, so the agreement there is evidence that the schemes converge near the fingerprint and not a second both-ways ablation. And my located 60 are the first 60 of his test split, which by my own measurement are the bright, low-contrast end of the located population, so the 0.914 against 0.923 pair is a probe rather than a fair estimate for either side.

The rows are in the diagnostic repo. results/shell_split/control_his60.jsonl and control_nonloc60.jsonl carry my per-volume counts on the two cohorts, results/shell_split/divergence_diagnosis.json puts my located rows next to his, and the cohort lists are results/histmatch_confound/located60.json and results/shell_split/nonlocated60.json. The block means by cohort are in results/histmatch_confound/per_volume.csv. The four wrapper arms and his published rows for the 16 are in results/wrapper_check/wrapper_four_way.json.

I have the rest of the 892 on disk and will post correctly normalized rows here when I have them, so the size of the effect should end up with a better estimate than this one.