segment.py

import sys
import argparse
import os
from glob import glob
import json
import SimpleITK as sitk

# importing utils and 
from utils.logger import logger, pprint
from utils.dataset import segment_lungs_and_remove_trachea
from enums.dtype import DataTypes


if __name__ == "__main__":
    # optional arguments from the command line 
    parser = argparse.ArgumentParser()

    parser.add_argument('--dataset_path', type=str, default='dataset/train', help='root dir for nifti training data')

    # parse the arguments
    args = parser.parse_args()

    # check if the dataset_path exists
    if not os.path.exists(args.dataset_path):
        logger.error(f"Path {args.dataset_path} does not exist")
        sys.exit(1)

    # get the list of exhale and inhale files from the dataset_path
    logger.info(f"Reading nifti data from '{args.dataset_path}'")
    exhale_volumes = [path.replace('\\', '/') for path in sorted(glob(os.path.join(args.dataset_path, "***" , "*eBHCT.nii.gz"), recursive=True))]
    inhale_volumes = [path.replace('\\', '/') for path in sorted(glob(os.path.join(args.dataset_path, "***" , "*iBHCT.nii.gz"), recursive=True))]

    # log the number of exhale and inhale files
    logger.info(f"Found {len(exhale_volumes)} exhale volumes: ({[subject.split('/')[-2] for subject in exhale_volumes]})")
    logger.info(f"Found {len(inhale_volumes)} inhale volumes: ({[subject.split('/')[-2] for subject in inhale_volumes]})\n")
    pprint(exhale_volumes, inhale_volumes)
    print('\n')

    # read the data dictionary
    with open(os.path.join(args.dataset_path.replace("train", "", 1).replace("test", "", 1), 'description.json'), 'r') as json_file:
        dictionary = json.loads(json_file.read())

    # iterate over all of the nifti inhale and exhale volumes and segment the lungs
    for volume in exhale_volumes + inhale_volumes:
        # get the subject name and information
        subject_name = volume.split('/')[-2]
        subject_information = dictionary[args.dataset_path.replace('\\', '/').split("/")[-1]][subject_name]

        logger.info(f"Segmenting {volume}")
        sitk_image = sitk.ReadImage(volume)
        np_image = sitk.GetArrayFromImage(sitk_image)

        # logs
        print(subject_information)
        print("sitk:\t\t", sitk_image.GetSize(), sitk_image.GetPixelIDTypeAsString(), sitk_image.GetOrigin(), sitk_image.GetSpacing())
        print("np:\t\t", np_image.shape, np_image.dtype)

        # segment the lungs
        if subject_name == 'copd2':
            # set a specific threshold to copd2
            threshold = 430
            fill_holes_before_trachea_removal = True
        else:
            threshold = 700 
            fill_holes_before_trachea_removal = False

        print("thresh:\t\t", threshold)
        print("fill_holes:\t", fill_holes_before_trachea_removal)

        _, _, _, lung_segmentation = \
            segment_lungs_and_remove_trachea(np_image, 
                                            threshold=threshold, structure=(7, 7, 7), fill_holes_before_trachea_removal=fill_holes_before_trachea_removal)
        
        lung_segmentation_sitk = sitk.GetImageFromArray(lung_segmentation)
        lung_segmentation_sitk.CopyInformation(sitk_image)

        # logs
        print("lung:\t\t", lung_segmentation.shape, lung_segmentation.dtype)
        print("lung_sitk:\t", lung_segmentation_sitk.GetSize(), lung_segmentation_sitk.GetPixelIDTypeAsString(), lung_segmentation_sitk.GetOrigin(), lung_segmentation_sitk.GetSpacing(), "\n")

        # save the lung segmentation
        sitk.WriteImage(lung_segmentation_sitk, volume.replace(".nii.gz", "_lung.nii.gz"))

    print("Segmentation complete!")