Initialize predictor from individual path

[MattTheCuber]

This PR adds the method initialize_from_paths as an alternative to initialize_from_trained_model_folder to the nnUNetPredictor class. With this method, instead of passing the entire training directory, you can just pass the files.

[FabianIsensee]

@claude please do a brief review of this. Seems useful!

[ghost]

🔍 nnU-Net Code Review

Summary

This PR adds initialize_from_paths to nnUNetPredictor, allowing users to initialize the predictor by passing individual file paths (dataset JSON, plans JSON, checkpoint) instead of a structured training output directory.

Contribution Guidelines

• General applicability: The use case (initializing from individual paths) is reasonable — it enables prediction when models are stored outside the standard nnU-Net directory layout (e.g., shipped as standalone artifacts). This aligns with the framework's goal of broad utility.
• No linked issue: For a feature addition to a core class, prior discussion in an issue would have been advisable per CONTRIBUTING.md ("For larger features and refactors, please open a GitHub issue to discuss the idea before starting work").
• No benchmarks needed: This is a functional addition, not a performance claim.

Key Findings

• Critical — Missing build_network_architecture signature check: The existing initialize_from_trained_model_folder (lines 106–128) uses inspect.signature to detect whether the trainer uses the new signature (plans_manager, configuration_manager, ...) or the old/deprecated one (network_arch_class_name, network_arch_init_kwargs, ...). The new initialize_from_paths method only calls the old signature, which will fail for any trainer that has been updated to the new signature. This is a correctness bug that will cause a TypeError at initialization time for modern trainers.

  The relevant block from the existing method that is missing:

  sig = inspect.signature(trainer_class.build_network_architecture)
  if 'plans_manager' in sig.parameters:
      network = trainer_class.build_network_architecture(
          plans_manager, configuration_manager,
          num_input_channels, num_output_channels,
          enable_deep_supervision=False,
      )
  else:
      # old signature with deprecation warning
      ...

• Single-fold only: The method loads exactly one checkpoint. The existing initialize_from_trained_model_folder supports loading multiple folds for ensemble prediction, which is a core part of nnU-Net's inference pipeline. This limitation should at minimum be clearly documented in the docstring so users understand they are trading multi-fold ensembling for flexibility.

• Substantial code duplication: The new method duplicates ~40 lines from initialize_from_trained_model_folder. Consider refactoring the shared logic (checkpoint parsing, network construction, attribute assignment) into a private helper that both methods call. This would also prevent the signature-check divergence bug noted above.

Domain-Specific Notes

• Inference pipeline correctness: Because the method uses only the old build_network_architecture signature, the constructed network may not match the architecture that produced the checkpoint weights, which could lead to silent prediction errors or a state_dict mismatch crash. This is the highest-priority fix.

Minor Suggestions

• Parameter naming: dataset_path suggests a path to the dataset directory, but it actually expects the path to dataset.json. Consider renaming to dataset_json_path for clarity (consistent with the existing codebase's dataset_json variable naming).
• Docstring: The current docstring ("A more flexible method for initialize_from_trained_model_folder") is quite sparse. Documenting the expected file formats and the single-fold limitation would help users.

[FabianIsensee]

This PR seems to have some issues that would need to be resolved for merging. I would also like to understand why this is needed. It adds complexity and if in the future we decide to have different (or more) file we need to initialization we would need to change the signature of this function. nnU-Net puts everything in one folder already, so why not keep using that?

[MattTheCuber]

My main reasoning was to reduce the number of files and file size of deploying models. We deploy our models to customers. Zipping up the results folder is pretty large. For example, one of our training output folders is 5.8 GB, where pulling out the three necessary files ends up being under 0.5 GB. This PR just makes extracting the files easier. A better solution might be to add a single ZIP/TAR file mode and some utilities for pulling relevant files from the training folder to create them. Then these could be deployed as a single file and read by nnUNet.

[FabianIsensee]

You don't need to zip the entire output folder. You can create a custom script for zipping checkpoints that only includes the necessary files. Those should be the three files your PR links and nothing else (but preserve the folder structure)

[FabianIsensee]

A small note on deploying to customers: You should delete the optimizer states from the checkpoints. This is a free halving of the checkpoint size :-)

[MattTheCuber]

Our use case would find it easier not to have to run a script on a folder that creates a zip with a specific hierarchy, but I understand if this PR isn't the optimal solution. Mostly just wanted to give more flexibility in model loading rather than a specific directory structure with three files.

Your note would be very useful, what key is the optimizer states stored under?

[FabianIsensee]

:-)