PredictSubOrgansnUnet.py

"""
Use this to run the nnU-Net inference. You can run multiple instances of this code in parallel!
Just change the `part_id` and `gpu` arguments to split the task into multiple parts.
Example of splitting it in 4 parts:
    python PredictSubOrgansnUnet.py --num_parts 4 --part_id 0 --gpu 0
    python PredictSubOrgansnUnet.py --num_parts 4 --part_id 1 --gpu 1
    python PredictSubOrgansnUnet.py --num_parts 4 --part_id 2 --gpu 2
    python PredictSubOrgansnUnet.py --num_parts 4 --part_id 3 --gpu 3
"""


import argparse
from nnunetv2.paths import nnUNet_results, nnUNet_raw
import torch
from batchgenerators.utilities.file_and_folder_operations import join
from nnunetv2.inference.predict_from_raw_data import nnUNetPredictor
import os
import pandas as pd

def split_files(files_input, files_output, num_parts, part_id):
    """
    Splits the files_input and files_output into num_parts and selects the partition for the given part_id.

    Args:
        files_input (list): List of input files.
        files_output (list): List of output files.
        num_parts (int): Number of partitions.
        part_id (int): Partition ID to select (starting from 0).

    Returns:
        tuple: (files_input_partition, files_output_partition) for the given part_id.
    """
    assert len(files_input) == len(files_output), "files_input and files_output must have the same length"
    assert 0 <= part_id < num_parts, "Invalid part_id. It must be in the range [0, num_parts-1]"

    # Split the files based on num_parts
    total_files = len(files_input)
    files_per_part = (total_files + num_parts - 1) // num_parts  # Ensure rounding up
    start_idx = part_id * files_per_part
    end_idx = min(start_idx + files_per_part, total_files)

    # Return the partition
    return files_input[start_idx:end_idx], files_output[start_idx:end_idx]

def parse_args():
    parser = argparse.ArgumentParser(description="Run nnUNet prediction with custom arguments.")
    parser.add_argument('--pth', type=str, default='/projects/bodymaps/Data/image_only/AbdomenAtlasPro/AbdomenAtlasPro/', 
                        help="Path to input images folder")
    parser.add_argument('--outdir', type=str, default='/projects/bodymaps/Pedro/data/JHH_liver_segments/', 
                        help="Path to output folder for predictions")
    parser.add_argument('--checkpoint', type=str, 
                        default='./nnUNetOrgansAndSubSegments/', 
                        help="Path to model checkpoint folder")
    parser.add_argument('--num_parts', type=int, default=1, 
                        help="Number of parts to split the task into")
    parser.add_argument('--part_id', type=int, default=0, 
                        help="ID of the current part (0-indexed)")
    parser.add_argument('--gpu', type=int, default=0, 
                        help="ID of the gpu to be used")
    parser.add_argument('--workers', type=int, default=2, 
                        help="Number of worker processes for preprocessing and segmentation export")
    parser.add_argument('--BDMAP_format', action='store_true', 
                    help="Enable BDMAP format")
    parser.add_argument('--ids', default=None, help='Path to csv with a BDMAP ID column and the cases you want to inference')            
    parser.add_argument('--reset', action='store_true', 
                    help="Overwrites old cases")
    return parser.parse_args()


def filter_existing_outputs(files_input, files_output):
    """
    Filter out cases where the output file already exists.
    
    Parameters:
        files_input (list): A list of input file paths (or lists of file paths).
        files_output (list): A list of output file paths.
        
    Returns:
        tuple: Two lists, (filtered_files_input, filtered_files_output), where
               each pair corresponds to an output that does not already exist.
    """
    filtered_input = []
    filtered_output = []
    
    for inp, out in zip(files_input, files_output):
        if not os.path.exists(out+'.nii.gz'):
            filtered_input.append(inp)
            filtered_output.append(out)
        else:
            print(f"Skipping already saved output: {out}")
    
    return filtered_input, filtered_output

def main():
    args = parse_args()

    # Instantiate the nnUNetPredictor
    predictor = nnUNetPredictor(
        tile_step_size=0.5,
        use_gaussian=True,
        use_mirroring=True,
        perform_everything_on_device=True,
        device=torch.device('cuda', args.gpu),
        verbose=False,
        verbose_preprocessing=False,
        allow_tqdm=True
    )

    if args.ids is not None:
        ids = pd.read_csv(args.ids)
        ids = sorted(ids['BDMAP ID'].to_list())
    else:
        ids = sorted(os.listdir(args.pth))
        
    print(f'Number of loaded ids: {len(ids)}')
    

    # Collect input and output files
    if args.BDMAP_format:
        print('Path:',args.pth)
        print('Files in path:',len(ids))

        files_input = [[os.path.join(args.pth, folder, 'ct.nii.gz')] 
                        for folder in ids 
                        if ('BDMAP' in folder)]
        print('Input files before split:',len(files_input))
        files_output = [os.path.join(args.outdir, folder) 
                        for folder in ids 
                        if ('BDMAP' in folder)]
        files_input, files_output = filter_existing_outputs(files_input, files_output)
        files_input, files_output = split_files(files_input, files_output, args.num_parts, args.part_id)
    else:
        if args.ids is None:
            files_input = [[os.path.join(args.pth, file)] for file in ids if (file.endswith('.nii.gz') and file[0]!='.')]
            files_output = [os.path.join(args.outdir, file.replace('.nii.gz','')) for file in ids if (file.endswith('.nii.gz') and file[0]!='.')]
        else:
            files_input = [[os.path.join(args.pth, file+'.nii.gz')] for file in ids]
            files_output = [os.path.join(args.outdir, file) for file in ids]
        files_input, files_output = filter_existing_outputs(files_input, files_output)
        
        files_input, files_output = split_files(files_input, files_output, args.num_parts, args.part_id)
        

    print('Input:', files_input[:10])
    print('Output:', files_output[:10])
    print(f'Cases to predict in part {args.part_id}: {len(files_input)}')


    # Initialize the network architecture and load the checkpoint
    predictor.initialize_from_trained_model_folder(
        join(args.checkpoint),
        use_folds=('all',),
        checkpoint_name='checkpoint_final.pth',
    )

    # Run the prediction
    predictor.predict_from_files(
        files_input, files_output,
        save_probabilities=False,
        overwrite=args.reset,
        num_processes_preprocessing=args.workers,
        num_processes_segmentation_export=args.workers,
        folder_with_segs_from_prev_stage=None,
        num_parts=1,
        part_id=0
    )


if __name__ == '__main__':
    main()