{"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":"# Vertebrae segmentation with TotalSegmentator","metadata":{}},{"cell_type":"code","source":"! pip install -q pydicom\n! pip install -q git+https://github.com/wasserth/TotalSegmentator.git\n! wget https://gist.githubusercontent.com/dgobbi/ab71f5128aa43f0d33a41775cb2bcca6/raw/c09283cc94d9faa882bd125ab6dc300a1f749153/dicom_to_nifti.py","metadata":{"execution":{"iopub.status.busy":"2022-09-11T12:27:06.088592Z","iopub.execute_input":"2022-09-11T12:27:06.089025Z","iopub.status.idle":"2022-09-11T12:27:59.948853Z","shell.execute_reply.started":"2022-09-11T12:27:06.088937Z","shell.execute_reply":"2022-09-11T12:27:59.947444Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os, sys, shutil, glob, pathlib, gc\nimport subprocess\nimport re\nimport numpy as np\nimport pandas as pd\nimport torch\nimport cv2\nimport matplotlib\nimport matplotlib.pyplot as plt\nimport pydicom\nfrom pydicom.pixel_data_handlers.util import apply_voi_lut\nimport nibabel","metadata":{"execution":{"iopub.status.busy":"2022-09-11T13:11:04.113163Z","iopub.execute_input":"2022-09-11T13:11:04.113938Z","iopub.status.idle":"2022-09-11T13:11:04.122836Z","shell.execute_reply.started":"2022-09-11T13:11:04.113889Z","shell.execute_reply":"2022-09-11T13:11:04.121789Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"COMP_DATA_DIR = '/kaggle/input/rsna-2022-cervical-spine-fracture-detection'\nCSECS = ['C1', 'C2', 'C3', 'C4', 'C5', 'C6', 'C7']\nTSECS = ['T1', 'T2', 'T3', 'T4', 'T5', 'T6', 'T7', 'T8', 'T9', 'T10', 'T11', 'T12']\nSEG_LABEL2IDX = {l: i + 1 for i, l in enumerate(CSECS + TSECS)}\nSEG_IDX2LABEL = {idx: l for l, idx in SEG_LABEL2IDX.items()}","metadata":{"execution":{"iopub.status.busy":"2022-09-11T12:57:24.287121Z","iopub.execute_input":"2022-09-11T12:57:24.287531Z","iopub.status.idle":"2022-09-11T12:57:24.294018Z","shell.execute_reply.started":"2022-09-11T12:57:24.287496Z","shell.execute_reply":"2022-09-11T12:57:24.292856Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"os.listdir(f'{COMP_DATA_DIR}/segmentations')[:5]","metadata":{"execution":{"iopub.status.busy":"2022-09-11T12:28:32.33791Z","iopub.execute_input":"2022-09-11T12:28:32.338737Z","iopub.status.idle":"2022-09-11T12:28:32.361757Z","shell.execute_reply.started":"2022-09-11T12:28:32.338684Z","shell.execute_reply":"2022-09-11T12:28:32.360775Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"study_id = '1.2.826.0.1.3680043.21321'","metadata":{"execution":{"iopub.status.busy":"2022-09-11T12:28:38.290527Z","iopub.execute_input":"2022-09-11T12:28:38.291146Z","iopub.status.idle":"2022-09-11T12:28:38.296118Z","shell.execute_reply.started":"2022-09-11T12:28:38.291107Z","shell.execute_reply":"2022-09-11T12:28:38.294972Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"! python dicom_to_nifti.py \\\n-i {COMP_DATA_DIR}/train_images/{study_id} \\\n-o {study_id}.nii.gz","metadata":{"execution":{"iopub.status.busy":"2022-09-11T12:28:42.073532Z","iopub.execute_input":"2022-09-11T12:28:42.074724Z","iopub.status.idle":"2022-09-11T12:29:10.615631Z","shell.execute_reply.started":"2022-09-11T12:28:42.074678Z","shell.execute_reply":"2022-09-11T12:29:10.614335Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"! TotalSegmentator \\\n-i {study_id}.nii.gz \\\n-o segmentations/{study_id} \\\n--preview \\\n--ml","metadata":{"execution":{"iopub.status.busy":"2022-09-11T12:37:01.043213Z","iopub.execute_input":"2022-09-11T12:37:01.043611Z","iopub.status.idle":"2022-09-11T12:43:42.087205Z","shell.execute_reply.started":"2022-09-11T12:37:01.043571Z","shell.execute_reply":"2022-09-11T12:43:42.085829Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"TOTALSEG_IDX2LABEL = {\n    23: \"T12\",\n    24: \"T11\",\n    25: \"T10\",\n    26: \"T9\",\n    27: \"T8\",\n    28: \"T7\",\n    29: \"T6\",\n    30: \"T5\",\n    31: \"T4\",\n    32: \"T3\",\n    33: \"T2\",\n    34: \"T1\",\n    35: \"C7\",\n    36: \"C6\",\n    37: \"C5\",\n    38: \"C4\",\n    39: \"C3\",\n    40: \"C2\",\n    41: \"C1\"}\n\nTOTALSEG_LABEL2IDX = {l: idx for idx, l in TOTALSEG_IDX2LABEL.items()}","metadata":{"execution":{"iopub.status.busy":"2022-09-11T12:59:41.043609Z","iopub.execute_input":"2022-09-11T12:59:41.044017Z","iopub.status.idle":"2022-09-11T12:59:41.050528Z","shell.execute_reply.started":"2022-09-11T12:59:41.043982Z","shell.execute_reply":"2022-09-11T12:59:41.049585Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"seg = nibabel.load(f'segmentations/{study_id}/s01.nii.gz').get_fdata()\n\nnew_seg = np.zeros_like(seg)\nfor l in (CSECS + TSECS):\n    new_seg[seg == TOTALSEG_LABEL2IDX[l]] = SEG_LABEL2IDX[l]\nseg = new_seg","metadata":{"execution":{"iopub.status.busy":"2022-09-11T12:46:42.803591Z","iopub.execute_input":"2022-09-11T12:46:42.804695Z","iopub.status.idle":"2022-09-11T12:46:47.182256Z","shell.execute_reply.started":"2022-09-11T12:46:42.804645Z","shell.execute_reply":"2022-09-11T12:46:47.181236Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"seg_gt = nibabel.load(f'{COMP_DATA_DIR}/segmentations/{study_id}.nii').get_fdata()","metadata":{"execution":{"iopub.status.busy":"2022-09-11T13:13:38.468041Z","iopub.execute_input":"2022-09-11T13:13:38.468657Z","iopub.status.idle":"2022-09-11T13:13:39.980787Z","shell.execute_reply.started":"2022-09-11T13:13:38.46862Z","shell.execute_reply":"2022-09-11T13:13:39.979666Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"seg_gt.shape, seg.shape","metadata":{"execution":{"iopub.status.busy":"2022-09-11T13:13:53.717315Z","iopub.execute_input":"2022-09-11T13:13:53.72019Z","iopub.status.idle":"2022-09-11T13:13:53.730173Z","shell.execute_reply.started":"2022-09-11T13:13:53.720147Z","shell.execute_reply":"2022-09-11T13:13:53.728956Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"cmap = matplotlib.cm.get_cmap('tab20', len(CSECS))\ncmap.set_over('yellow')\ncmap.set_under('black')\nkwargs = dict(\n    cmap=cmap,\n    vmin=1 - .5, vmax=len(CSECS) + .5\n)","metadata":{"execution":{"iopub.status.busy":"2022-09-11T13:14:17.719604Z","iopub.execute_input":"2022-09-11T13:14:17.720262Z","iopub.status.idle":"2022-09-11T13:14:17.727773Z","shell.execute_reply.started":"2022-09-11T13:14:17.720224Z","shell.execute_reply":"2022-09-11T13:14:17.726533Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"islice = 450\n\nfig, axs = plt.subplots(nrows=1, ncols=2, figsize=(12, 6))\naxs = axs.flatten()\nh0 = axs[0].imshow(seg_gt[:,::-1,::-1].transpose(2, 1, 0)[islice], **kwargs)\nfig.colorbar(h0, ax=axs[0], ticks=range(1, len(CSECS)+1), extend='both')\naxs[0].set_title('GT segmentation')\n\nh1 = axs[1].imshow(seg[::-1,::-1].transpose(2, 1, 0)[islice], **kwargs)\nfig.colorbar(h1, ax=axs[1], ticks=range(1, len(CSECS)+1), extend='both')\naxs[1].set_title('TotalSegmentator');\n\nfig.suptitle(f'Study {study_id} islice={islice}')\nplt.tight_layout();","metadata":{"execution":{"iopub.status.busy":"2022-09-11T13:21:46.499753Z","iopub.execute_input":"2022-09-11T13:21:46.500259Z","iopub.status.idle":"2022-09-11T13:21:47.311307Z","shell.execute_reply.started":"2022-09-11T13:21:46.500213Z","shell.execute_reply":"2022-09-11T13:21:47.309874Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}