{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# *Proof of concept for a segmentator that could be run at predicton time*\n\n#### Sources\n* https://github.com/wasserth/TotalSegmentator\n* https://github.com/MIC-DKFZ/nnUNet/\n* https://github.com/icometrix/dicom2nifti\n\n#### Datasets\n* TotalSegmentator dataset with the segmentations provided by the host of the competiton.\n* https://www.kaggle.com/competitions/rsna-2023-abdominal-trauma-detection/discussion/428538\n\n#### Objetives\n* Load a dicom serie\n* Return the segmentation and the dicom array used\n* Do it under 15 seconds per patient\n* Standardize the output for easy feeding to other NNs\n    * always return an array with 128 channels, image of size 256 x 256\n    * if the serie have under 128 dicoms, it gets interpolate with nearest to 128\n\n#### Constrains\n* Need to work offline\n* Reusing the dicom array used in the segmentation for the prediction is a must since its already in memory\n* Using GPU T4 x2 is probably the way to go, to keep the segmentation model loaded in a GPU\n* Can only segment a small subset of a dicom serie (128)\n* Can only segment a serie per patient to have enough time for the prediction\n\n#### The code of nnUNet was modified to resample in the GPU and to load the dicoms from a fork of dicom2nifti, returns a list with the original dicoms, two arrays for the segmentation and the 2.5d image, 128 channels, 256x256\n\n#### The segmentation classes: \n* liver: 1\n* spleen: 2\n* lkidney: 3\n* rkidney: 4\n* bowel: 5\n","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19"}},{"cell_type":"markdown","source":"## Install dependencies","metadata":{}},{"cell_type":"code","source":"!cp -r ../input/atd-rsna-segmentator-poc/ATD_RSNA_Segmentator/* .\n!pip install --no-index --no-deps ./wheels/*.whl --quiet","metadata":{"execution":{"iopub.status.busy":"2023-09-15T12:19:37.453985Z","iopub.execute_input":"2023-09-15T12:19:37.454309Z","iopub.status.idle":"2023-09-15T12:19:48.626835Z","shell.execute_reply.started":"2023-09-15T12:19:37.454278Z","shell.execute_reply":"2023-09-15T12:19:48.62556Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Config","metadata":{}},{"cell_type":"code","source":"BASE_PATH = \"../input/rsna-2023-abdominal-trauma-detection/train_images\"","metadata":{"execution":{"iopub.status.busy":"2023-09-15T12:19:48.629564Z","iopub.execute_input":"2023-09-15T12:19:48.629892Z","iopub.status.idle":"2023-09-15T12:19:48.636224Z","shell.execute_reply.started":"2023-09-15T12:19:48.629864Z","shell.execute_reply":"2023-09-15T12:19:48.635143Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Check the execution time for 40 patients","metadata":{}},{"cell_type":"code","source":"from os.path import join\nfrom os import makedirs, listdir\n\nimport numpy as np\nimport torch\nimport matplotlib.pyplot as plt\n\nfrom infer import volume_and_seg","metadata":{"execution":{"iopub.status.busy":"2023-09-15T12:19:48.637865Z","iopub.execute_input":"2023-09-15T12:19:48.63821Z","iopub.status.idle":"2023-09-15T12:19:55.475473Z","shell.execute_reply.started":"2023-09-15T12:19:48.638179Z","shell.execute_reply":"2023-09-15T12:19:55.474476Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def run_time(path_serie):\n    \"\"\" Print the inference execution time in seconds \"\"\"\n    start = torch.cuda.Event(enable_timing=True)\n    end = torch.cuda.Event(enable_timing=True)\n\n    start.record()\n\n    dicom_list, seg, vol = volume_and_seg(path_serie, clear_mempool=True, cuda=1)\n\n    end.record()\n    torch.cuda.synchronize()\n\n    print(f\"{start.elapsed_time(end)/1_000:.2f} seconds.\")","metadata":{"execution":{"iopub.status.busy":"2023-09-15T12:19:55.477879Z","iopub.execute_input":"2023-09-15T12:19:55.478469Z","iopub.status.idle":"2023-09-15T12:19:55.487165Z","shell.execute_reply.started":"2023-09-15T12:19:55.478433Z","shell.execute_reply":"2023-09-15T12:19:55.485072Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"patients = sorted(listdir(BASE_PATH))\n\nfor patient in patients[:40]:  # 40 patients for testing purposes\n    path = join(BASE_PATH, patient)\n    series = sorted(listdir(path))\n    for serie in series:\n        path_serie = join(path, serie)\n        run_time(path_serie)\n        break  # One serie per patient","metadata":{"execution":{"iopub.status.busy":"2023-09-15T12:19:55.48883Z","iopub.execute_input":"2023-09-15T12:19:55.489665Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"#### Took a long time for the first inference, but after the model is loaded this could be viable, keep in mind that now the segmentation and the array that contains the image are in memory, so it doesn't need to be loaded from the disk again to make the prediction.\n\n#### Now that we have the execution time, lets see if it actually segment any organ at all, the model in this dataset is far away from what should be doable, it was trained without CV during a night in a NVIDIA 3090.","metadata":{"execution":{"iopub.status.busy":"2023-09-15T09:08:17.129706Z","iopub.execute_input":"2023-09-15T09:08:17.130313Z","iopub.status.idle":"2023-09-15T09:08:18.270313Z","shell.execute_reply.started":"2023-09-15T09:08:17.130282Z","shell.execute_reply":"2023-09-15T09:08:18.26915Z"}}},{"cell_type":"code","source":"path = \"../input/rsna-2023-abdominal-trauma-detection/train_images/10004/21057\"\n\n# :dicom_list: the dicoms used for the segmentation in case you want to process it in another way\n# :seg: segmentation \n# :vol: image of size 256x256 with 128 channels\n# :path_serie: path to the dicom serie\n# :clear_mempool: clear the GPU memory\n# :cuda: id of the cuda device to use\ndicom_list, seg, vol = volume_and_seg(path_serie, clear_mempool=True, cuda=1)\n\ndef plot_image_with_seg(volume, volume_seg=[], num_subplots=20):\n    # Source: https://www.kaggle.com/code/parhammostame/construct-3d-arrays-from-dcm-nii-3-view-angles\n    slices = np.linspace(0, volume.shape[0]-1, num_subplots).astype(np.int16)\n           \n    rows = np.max( [np.floor(np.sqrt(num_subplots)).astype(int) - 2, 1])\n    cols = np.ceil(num_subplots/rows).astype(int)\n    \n    fig, ax = plt.subplots(rows, cols, figsize=(cols * 2, rows * 4))\n    fig.tight_layout(h_pad=0.01, w_pad=0)\n    \n    ax = ax.ravel()\n    for this_ax in ax:\n        this_ax.axis('off')\n\n    for counter, this_slice in enumerate( slices ):\n        plt.sca(ax[counter])\n        \n        image = volume[this_slice, :, :]\n        plt.imshow(image, cmap='gray')\n        \n        mask = np.where(volume_seg[this_slice, :, :], volume_seg[this_slice, :, :], np.nan)\n        plt.imshow(mask, cmap='Set1', alpha=0.5)        \n        \n# seg = seg == 4  # Change the number to the organ that you want to visualize\n# 0 is the background of the segmentation, if you want to use it as a mask of the combined segmentations you need to invert it first\n# mask = 1 - (seg == 0).astype(np.uint8)  # Change type from boolean to uint8 and swap the 0s and 1s (now the background is 0 and the segmentation 1)\nplot_image_with_seg(vol, seg, num_subplots=seg.shape[0])","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# How to train your own segmentation model in a local workstation or Google Colab\n    \n* Install pytorch\n    * Follow the instructions from the [official page](https://pytorch.org/get-started/locally/)\n\n* Install nnUNet\n    * Instructions from [GitHub](https://github.com/MIC-DKFZ/nnUNet/blob/master/documentation/installation_instructions.md)\n> git clone https://github.com/MIC-DKFZ/nnUNet.git &\n> cd nnUNet &\n> pip install -e .\n    * Create 3 folders in the nnUNet directory\n        * raw\n        * preprocess\n        * results\n    * Export the global variables with the path to the folders, example\n>         export nnUNet_raw=\"/workspace/nnUNet/raw\"\n>         export nnUNet_preprocessed=\"/workspace/nnUNet/preprocessed\"\n>         export nnUNet_results=\"/workspace/nnUNet/results\"\n\n* The dataset\n    * The files can be download from [here](https://www.kaggle.com/competitions/rsna-2023-abdominal-trauma-detection/discussion/428538)\n    * The affine of the segmentations provided by the host needs to be fixed\n    * Need to move the files into 3 folders with the structure\n        * imagesTr\n        * imagesTs\n        * labelsTr\n    * Name schema for the files:\n        * imagesTr -> sample0000_0000.nii.gz, sample0001_0000.nii.gz ... sample9999_0000.nii.gz\n        * labelsTr -> sample0000.nii.gz, sample0001.nii.gz ... sample9999.nii.gz\n    * Create a json file with the modality (CT) and the labels for the dataset, example:\n>     {\n>     \"channel_names\": {\n>         \"0\": \"CT\"\n>     },\n>     \"labels\": {\n>         \"background\": 0,\n>         \"liver\": 1,\n>         \"spleen\": 2,\n>         \"left kidney\": 3,\n>         \"right kidney\": 4,\n>         \"bowel\": 5\n>     },\n>     \"numTraining\": 1080,\n>     \"file_ending\": \".nii.gz\",\n>     \"overwrite_image_reader_writer\": \"NibabelIO\"\n>    }\n    * numTraining is the total number of training samples, is a sanity check by nnUNet\n    * Nibabel reader_writer instead of simpleitk\n    * Remove the original segmentations from the TotalSegmentator dataset and add the ones provided by the host\n    * Move the folders to a folder with the name DatasetID_X where ID is a number for example, Dataset020_ATD,\n      move this folder to the raw folder\n    \n    * [I have uploaded the dataset with all this done](https://drive.proton.me/urls/H4YHF8X6K4#s7JqIbcZKyJs), so is ready to be used in nnUNet\n    \n## Plotting to check if the segmentations are aligned\n* -d dataset_id -o png_folder\n> nnUNetv2_plot_overlay_pngs -d 020 -o /workspace/nnUNet/pngs\n\n## Plan and preprocess the dataset, verify intengrity\n* Configurations: 3d_fullres, 3d_lowres, 2d\n* the default GPU memory target is 8GB, you can change it with -gpu_memory_target (try first with the default values) \n* id of the dataset (id of the dataset in the raw folder, 020 in this case) \n> nnUNetv2_plan_and_preprocess -d 020 --verify_dataset_integrity -c 3d_fullres\n\n## Training\n* The TotalSegmentator trainer is called nnUNetTrainer_Xepochs_NoMirroring, is in the folder /nnunetv2/training/nnUNetTrainer/variants/training_length/, you should create a class with the number of epochs that you want to train, using the 250 epoch version for the tutorial, nnUNetTrainer_250epochs_NoMirroring.\n* the structure is -tr trainer dataset_id configuration fold\n    * fold is 0 to 4 or -f all (no fold)\n> nnUNetv2_train -tr nnUNetTrainer_250epochs_NoMirroring 020 3d_fullres 0\n\n## Useful information\n* export nnUNet_compile=True if you want pytorch to compile the model\n* You can continue training from the last checkpoint with the argument -c\n* 3d_lowres isn't worth it in my experiments\n* -overwrite_target_spacing allows to force a voxel_spacing value, worth checking it if you want to reuse the 2.5d array at end for the prediction model too, so you have some consistency between models.","metadata":{}},{"cell_type":"markdown","source":"## Things you probably want to modify\n* In the folder dicom2nifti, the file common.py, you may want to modify the function read_dicom_directory to improve the selection of the slices from a serie.\n* the pixel standardization and the window is done by the dicom2nifti, if you want to force the window in each dicom, change the code in the function apply_scaling in common.py\n* If you want the original affine, return it in the volume_and_seg function, infer.py file\n* You can create a nifti object at any point with nibabel, nib.Nifti1Image(img, affine) or nib.Nifti1Image(img)\n* ProxyArrays with nibabel can help you process the images faster, https://nipy.org/nibabel/images_and_memory.html\n* Cucim and cupy are next level, probably can add it in more places.\n* Interpolation nearest isn't very good, linear is also fast, cubic would be amazing but that isn't doable with the target hardware, you may want to experiment with other methods.","metadata":{}}]}