{"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":"# Setup and helper functions","metadata":{}},{"cell_type":"code","source":"!pip install pydicom\n!pip install python-gdcm\n!pip install pylibjpeg pylibjpeg-libjpeg pydicom","metadata":{"execution":{"iopub.status.busy":"2023-04-08T20:57:32.839593Z","iopub.execute_input":"2023-04-08T20:57:32.840198Z","iopub.status.idle":"2023-04-08T20:58:12.423841Z","shell.execute_reply.started":"2023-04-08T20:57:32.839966Z","shell.execute_reply":"2023-04-08T20:58:12.422479Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nimport re\nimport csv\nimport cv2\nimport random\nimport shutil\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nfrom matplotlib.patches import Rectangle\n\nimport pydicom as dicom\n\nimport nibabel as nib","metadata":{"execution":{"iopub.status.busy":"2023-04-08T20:58:12.426469Z","iopub.execute_input":"2023-04-08T20:58:12.426969Z","iopub.status.idle":"2023-04-08T20:58:13.204917Z","shell.execute_reply.started":"2023-04-08T20:58:12.426915Z","shell.execute_reply":"2023-04-08T20:58:13.203682Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"rsna_root = '/kaggle/input/rsna-2022-cervical-spine-fracture-detection'","metadata":{"execution":{"iopub.status.busy":"2023-04-08T20:58:13.206654Z","iopub.execute_input":"2023-04-08T20:58:13.209385Z","iopub.status.idle":"2023-04-08T20:58:13.21561Z","shell.execute_reply.started":"2023-04-08T20:58:13.209326Z","shell.execute_reply":"2023-04-08T20:58:13.21415Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Helper functions\n\ndef load_img_from_dcm(path):\n    img = dicom.dcmread(path)\n    img.PhotometricInterpretation = 'YBR_FULL'\n    data = img.pixel_array\n    data = data - np.min(data)\n    if np.max(data != 0):\n        data = data/np.max(data)\n    data = (data * 255).astype(np.uint8)\n    return cv2.cvtColor(data, cv2.COLOR_GRAY2RGB)\n\ndef CT_path_to_3D_arr(folder_path, l=None): # folder_path is for folder of dcm files constituting one CT scan\n    if l == None:\n        l = os.listdir(folder_path) # list of 2D slices of the CT scan\n        l.sort()\n        l = sorted(l, key=len)\n    \n    CT_arr = [] # the full 3D CT\n    for dcm in l:\n        dcm_path = os.path.join(folder_path, dcm)\n        dcm_arr = load_img_from_dcm(dcm_path)\n        CT_arr.append(dcm_arr)\n    CT_arr = np.asarray(CT_arr)\n    return CT_arr","metadata":{"execution":{"iopub.status.busy":"2023-04-08T20:58:13.219502Z","iopub.execute_input":"2023-04-08T20:58:13.220117Z","iopub.status.idle":"2023-04-08T20:58:13.231187Z","shell.execute_reply.started":"2023-04-08T20:58:13.220065Z","shell.execute_reply":"2023-04-08T20:58:13.229786Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Bounding boxes around vertebrae\ncan be obtained from the segmentations data. First I will visualize the segmentation slices side-by-side with their corresponding CT slices.","metadata":{}},{"cell_type":"code","source":"patient_id = \"1.2.826.0.1.3680043.780\"\n\nsegm_path = os.path.join(rsna_root, 'segmentations', patient_id+'.nii')\nsegm_arr = nib.load(segm_path).get_fdata()\nsegm_arr = np.transpose(segm_arr, (2, 0, 1))\n\nCT_path = os.path.join(rsna_root, 'train_images', patient_id)\nCT_arr = CT_path_to_3D_arr(CT_path)","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2023-04-08T20:58:13.232977Z","iopub.execute_input":"2023-04-08T20:58:13.233355Z","iopub.status.idle":"2023-04-08T20:58:18.300058Z","shell.execute_reply.started":"2023-04-08T20:58:13.233322Z","shell.execute_reply":"2023-04-08T20:58:18.298937Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"fig, ax = plt.subplots(1, 2, figsize=(12, 6))\n\nslice_num = 80 # from 0 to len(array)-1\nsegm_slice = segm_arr[len(segm_arr)-1 - slice_num]\nsegm_slice = np.rot90(segm_slice)\nCT_slice = CT_arr[slice_num]\nax[0].imshow(CT_slice, cmap=plt.get_cmap('bone'))\nax[1].imshow(segm_slice, cmap=plt.get_cmap('bone'))","metadata":{"execution":{"iopub.status.busy":"2023-04-08T20:58:18.301641Z","iopub.execute_input":"2023-04-08T20:58:18.3021Z","iopub.status.idle":"2023-04-08T20:58:18.837931Z","shell.execute_reply.started":"2023-04-08T20:58:18.302054Z","shell.execute_reply":"2023-04-08T20:58:18.83665Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"We can draw a bounding box (bbox) around the vertebrae by using the segmentation slice like this:","metadata":{}},{"cell_type":"code","source":"rows = np.any(segm_slice, axis=1)\ncols = np.any(segm_slice, axis=0)\nrmin, rmax = np.where(rows)[0][[0, -1]]\ncmin, cmax = np.where(cols)[0][[0, -1]]\nwidth = cmax - cmin\nheight = rmax - rmin\n\nfig, ax = plt.subplots(1, 1, figsize=(6,6))\nax.imshow(segm_slice)\nrect = Rectangle((cmin, rmin), width, height,\n                 linewidth=1, edgecolor='r', facecolor='none')\nax.add_patch(rect)","metadata":{"execution":{"iopub.status.busy":"2023-04-08T20:58:18.840116Z","iopub.execute_input":"2023-04-08T20:58:18.84051Z","iopub.status.idle":"2023-04-08T20:58:19.134432Z","shell.execute_reply.started":"2023-04-08T20:58:18.840475Z","shell.execute_reply":"2023-04-08T20:58:19.133141Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Drawing the bbox on the original CT scan:","metadata":{}},{"cell_type":"code","source":"fig, ax = plt.subplots(1, 1, figsize=(6,6))\nax.imshow(CT_slice, cmap=plt.get_cmap('bone'))\nrect = Rectangle((cmin, rmin), width, height,\n                 linewidth=1, edgecolor='r', facecolor='none')\nax.add_patch(rect)\n\nprint(cmin, rmin, width, height)","metadata":{"execution":{"iopub.status.busy":"2023-04-08T20:58:19.136263Z","iopub.execute_input":"2023-04-08T20:58:19.136846Z","iopub.status.idle":"2023-04-08T20:58:19.456654Z","shell.execute_reply.started":"2023-04-08T20:58:19.136812Z","shell.execute_reply":"2023-04-08T20:58:19.455223Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Save labels (bboxes) in YOLOv5 format to use for training\nI used [Ultralytics YOLOv5](https://docs.ultralytics.com/quick-start/), which I found super easy to use. The specifics are well-documented on their website, but here is the gist:\n- one txt file per image\n- one row per object/bbox\n- each row in 'class_number, xcentre, ycentre, width, height' format (here, we only have one class '0: vertebra')\n- bbox coordinates are normalized, ie divided by width/height in length\n- (0,0) is top-left","metadata":{}},{"cell_type":"code","source":"def save_yolo_coord(pt_num, slice_num, slice_, destination_folder):\n    '''\n    Saves the yolov5 coord txt file for a single slice segmentation mask\n    \n    pt_num: number that follows 1.2.826.0.1.3680043.\n    slice_num: which slice\n    arr: 2D array of slice\n    destination_folder: where to save the txt files\n    '''\n    rows = np.any(slice_, axis=1)\n    cols = np.any(slice_, axis=0)\n    rmin, rmax = np.where(rows)[0][[0, -1]]\n    cmin, cmax = np.where(cols)[0][[0, -1]]\n    \n    xcentre = int((cmin+cmax)/2)\n    ycentre = int((rmin+rmax)/2)\n    width = cmax - cmin\n    height = rmax - rmin\n    img_width = slice_.shape[1]\n    img_height = slice_.shape[0]\n    \n    # yolo coordiates: class, xcentre, ycentre, width, height (normalized by width/hegith of image)\n    yolo_coord = [0, xcentre/img_width, ycentre/img_height, width/img_width, height/img_height]\n    \n    filename = os.path.join(destination_folder, str(pt_num)+\"_\"+str(slice_num)+'.txt')\n    with open(filename, 'w') as file:\n        writer = csv.writer(file, delimiter=' ')\n        writer.writerow(yolo_coord)","metadata":{"execution":{"iopub.status.busy":"2023-04-08T20:58:19.458087Z","iopub.execute_input":"2023-04-08T20:58:19.45843Z","iopub.status.idle":"2023-04-08T20:58:19.46984Z","shell.execute_reply.started":"2023-04-08T20:58:19.458398Z","shell.execute_reply":"2023-04-08T20:58:19.468909Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def save_bboxes_from_nii(pt_num, nii_folder, dest_folder_name):\n    '''\n    Saves yolov5 coord txt files for one patient's segmentation masks (ie, one nii file)\n    '''\n    destination_folder = os.path.join(os.getcwd(), dest_folder_name)\n    if not os.path.exists(destination_folder):\n        print(f'Creating destination folder in current directory: {dest_folder_name}')\n        os.mkdir(destination_folder)\n    \n    arr = nib.load(os.path.join(nii_folder, '1.2.826.0.1.3680043.'+str(pt_num)+'.nii')).get_fdata()\n    arr = np.transpose(arr, (2, 0, 1))\n    arr = np.flip(arr, axis=0)\n    \n    for slice_num, slice_ in enumerate(arr):\n        if not slice_.any():\n            continue\n        slice_ = np.rot90(slice_)\n        \n        save_yolo_coord(pt_num, slice_num, slice_, destination_folder)","metadata":{"execution":{"iopub.status.busy":"2023-04-08T20:58:19.471435Z","iopub.execute_input":"2023-04-08T20:58:19.472531Z","iopub.status.idle":"2023-04-08T20:58:19.483654Z","shell.execute_reply.started":"2023-04-08T20:58:19.472479Z","shell.execute_reply":"2023-04-08T20:58:19.482169Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# List of patient numbers in segmentations folder\nniis = os.listdir('/kaggle/input/rsna-2022-cervical-spine-fracture-detection/segmentations')\npts = [re.search('(?<=1.2.826.0.1.3680043.)([0-9]*)(?=.nii)', filename).group(0) for filename in niis]\npts = [int(s) for s in pts]\nprint(\"List of pts with segmentations: \", pts)","metadata":{"execution":{"iopub.status.busy":"2023-04-08T20:58:19.489331Z","iopub.execute_input":"2023-04-08T20:58:19.489897Z","iopub.status.idle":"2023-04-08T20:58:19.546325Z","shell.execute_reply.started":"2023-04-08T20:58:19.489855Z","shell.execute_reply":"2023-04-08T20:58:19.545086Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"dest_folder_name = 'yolo_coords'\nnii_folder = '/kaggle/input/rsna-2022-cervical-spine-fracture-detection/segmentations'\nfor pt_num in pts:\n    save_bboxes_from_nii(pt_num, nii_folder, dest_folder_name)","metadata":{"scrolled":true,"execution":{"iopub.status.busy":"2023-04-08T20:58:19.547799Z","iopub.execute_input":"2023-04-08T20:58:19.548208Z","iopub.status.idle":"2023-04-08T21:01:36.797353Z","shell.execute_reply.started":"2023-04-08T20:58:19.548169Z","shell.execute_reply":"2023-04-08T21:01:36.795995Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Test to see if saved coords are correct\n\nl = os.listdir('/kaggle/working/yolo_coords')\nl = sorted(l)\nl.sort(key=len)\n\nf = np.random.choice(l)\nprint(f)\npt_num = re.search(\"^([0-9]*)(?=_)\", f).group(0)\nslice_num = re.search(\"(?<=_)([0-9]*)(?=.txt)\", f).group(0)\n\nCT_path = os.path.join(rsna_root, 'train_images', \"1.2.826.0.1.3680043.\"+pt_num)\nCT_arr = CT_path_to_3D_arr(CT_path)\nCT_slice = CT_arr[int(slice_num)]\nimg_width = CT_slice.shape[1]\nimg_height = CT_slice.shape[0]\n\nfig, ax = plt.subplots(1,1,figsize=(6,6))\nax.imshow(CT_slice, cmap=plt.get_cmap('bone'))\n\np = '/kaggle/working/yolo_coords/'+f\nwith open(p, 'r') as txt_file:\n    reader = csv.reader(txt_file)\n    row = next(reader)\n\nrow = [float(num) for num in row[0].split()]\nbbox_xcentre = img_width * row[1]\nbbox_ycentre = img_height * row[2]\nbbox_width = img_width * row[3]\nbbox_height = img_height * row[4]\nrect = Rectangle((bbox_xcentre-int(bbox_width/2), bbox_ycentre-int(bbox_height/2)),\n                  bbox_width, bbox_height,\n                  linewidth=1, edgecolor='r', facecolor='none')\nax.add_patch(rect)","metadata":{"execution":{"iopub.status.busy":"2023-04-08T21:01:36.79892Z","iopub.execute_input":"2023-04-08T21:01:36.799482Z","iopub.status.idle":"2023-04-08T21:01:51.37851Z","shell.execute_reply.started":"2023-04-08T21:01:36.799444Z","shell.execute_reply":"2023-04-08T21:01:51.377257Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Save Images for YOLOv5 training\nI found it easier to save the CT slices as jpegs for training Ultralytics YOLO.","metadata":{}},{"cell_type":"code","source":"# yolo coord txt files for all patients\ntxt_files = os.listdir('/kaggle/working/yolo_coords')\n\n# patients with segmentation data\nniis = os.listdir('/kaggle/input/rsna-2022-cervical-spine-fracture-detection/segmentations')\npts = [re.search('(?<=1.2.826.0.1.3680043.)([0-9]*)(?=.nii)', filename).group(0) for filename in niis]\npts = [int(s) for s in pts]\n\n# Save slices corresponding to each yolo_coord txt file\ntrain_images = '/kaggle/input/rsna-2022-cervical-spine-fracture-detection/train_images'\n\nyolo_slices = os.path.join(os.getcwd(), 'yolo_slices')\nif not os.path.exists(yolo_slices):\n    os.mkdir(yolo_slices)\n    \nfor pt in pts:\n    slice_nums = []\n    for txt_file in txt_files:\n        if re.search('^([0-9]*)(?=_)', txt_file).group(0) == str(pt): # txt files for pt\n            slice_num = int(re.search(\"(?<=_)([0-9]*)(?=.txt)\", txt_file).group(0))\n            slice_nums.append(slice_num)\n    \n    pt_CT = os.path.join(train_images, f\"1.2.826.0.1.3680043.{str(pt)}\") \n\n    pt_slices_all = os.listdir(pt_CT) # list of dcms\n    pt_slices_all = sorted(pt_slices_all)\n    pt_slices_all.sort(key=len)\n    \n    pt_slices = [pt_slices_all[slice_num] for slice_num in sorted(slice_nums)]\n    \n    if min(pt_slices_all, key=len) == '2.dcm': # scans that are missing 1.dcm\n        for slice_ in pt_slices:\n            img_path = os.path.join(train_images, \"1.2.826.0.1.3680043.\"+str(pt), slice_)\n            img = load_img_from_dcm(img_path)\n\n            imgs_savepath = os.path.join(yolo_slices, f\"{str(pt)}_{int(slice_[:-4])-2}.jpg\")\n            cv2.imwrite(imgs_savepath, img)        \n    \n    else: # normal scans that start from 1.dcm\n        for slice_ in pt_slices:\n            img_path = os.path.join(train_images, \"1.2.826.0.1.3680043.\"+str(pt), slice_)\n            img = load_img_from_dcm(img_path)\n\n            imgs_savepath = os.path.join(yolo_slices, f\"{str(pt)}_{int(slice_[:-4])-1}.jpg\")\n            cv2.imwrite(imgs_savepath, img)","metadata":{"scrolled":true,"_kg_hide-input":false,"_kg_hide-output":false,"execution":{"iopub.status.busy":"2023-04-08T21:01:51.379793Z","iopub.execute_input":"2023-04-08T21:01:51.380181Z","iopub.status.idle":"2023-04-08T21:13:01.622661Z","shell.execute_reply.started":"2023-04-08T21:01:51.380146Z","shell.execute_reply":"2023-04-08T21:13:01.620456Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Organize images and labels into directories","metadata":{}},{"cell_type":"code","source":"# Split into train and valid set\nall_images = os.listdir('/kaggle/working/yolo_slices')\nall_images = [f[:-4] for f in all_images]\nrandom.shuffle(all_images)\nSPLIT_POINT = int(len(all_images) * 0.9)\ntrain_set = all_images[:SPLIT_POINT]\nvalid_set = all_images[SPLIT_POINT:]\n\n# Put each image/label into right directory\nfolders = [\"train/images\", \"train/labels\", \"valid/images\", \"valid/labels\"]\nfor folder in folders:\n    if not os.path.exists(folder):\n        os.makedirs(folder)\n        \nfor idx in train_set:\n    jpg_file = f'/kaggle/working/yolo_slices/{idx}.jpg'\n    txt_file = f'/kaggle/working/yolo_coords/{idx}.txt'\n    \n    shutil.move(jpg_file, '/kaggle/working/train/images')\n    shutil.move(txt_file, '/kaggle/working/train/labels')\n\nfor idx in valid_set:\n    jpg_file = f'/kaggle/working/yolo_slices/{idx}.jpg'\n    txt_file = f'/kaggle/working/yolo_coords/{idx}.txt'\n    \n    shutil.move(jpg_file, '/kaggle/working/valid/images')\n    shutil.move(txt_file, '/kaggle/working/valid/labels')","metadata":{"execution":{"iopub.status.busy":"2023-04-08T21:14:13.160527Z","iopub.execute_input":"2023-04-08T21:14:13.160981Z","iopub.status.idle":"2023-04-08T21:14:13.171703Z","shell.execute_reply.started":"2023-04-08T21:14:13.160943Z","shell.execute_reply":"2023-04-08T21:14:13.170161Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}