Fix nnU-Net output filename matching & Add --only_backbone option

[FrancescoCorrenti]

Fix nnU-Net output filename matching & Add --only_backbone option

Description

This PR introduces two main improvements to the segmentation pipeline logic in run_segment.py.

1. Fix: Robust filename matching for nnU-Net outputs

The Problem:
Previously, the script attempted to locate the intermediate output files (generated by inference_nnunet.py) by deriving the filename from the output_folder path. However, the underlying nnU-Net inference script generates output filenames based on the parent directory of the input file, not the output directory.

The Fix:
I updated the filename derivation logic to use:

f_n = os.path.split(os.path.dirname(input_path))[1]

This ensures that run_segment.py correctly locates the .npz and .pkl files generated by the backbone inference step. This fixes FileNotFoundError crashes that occurred when the input folder name and output folder name did not match (specifically when BFC was disabled).

2. Feature: Added --only_backbone argument

I added a new flag --only_backbone to the argument parser.

• Behavior: When this flag is set, the script terminates immediately after the "Backbone AI" (nnU-Net) inference and data loading steps.
• Use Case: This is useful for users who only need the deep learning prediction and want to save time by skipping the computationally expensive multi-atlas fallback and Trustworthy AI merging steps.

Changes

• run_segment.py:
  • Updated f_n variable assignment to use os.path.dirname(input_path) to align with nnU-Net naming convention.
  • Added --only_backbone to ArgumentParser.
  • Added logic to return early if args.only_backbone is true, skipping the fallback and trustworthy sections.

Testing

• Verified that the script now correctly finds intermediate files even when the Output Folder name differs completely from the Input Folder name.
• Verified that passing --only_backbone successfully produces the Backbone AI output and exits without errors before running the atlas registration.