{"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":"**Step 1: Import packages**","metadata":{}},{"cell_type":"code","source":"import os\nimport pickle\nimport time\nimport SimpleITK as sitk\nimport numpy as np\nimport torch\nimport path\nimport sys\nimport pandas as pd\nimport skimage\nimport skimage.morphology\nimport threading\nimport math\nimport gc\n\n# ------------------------------------ #\nsrc_path = \"../input/source-code/src\"\nif src_path not in sys.path:\n    sys.path.insert(0, src_path)\n\nprint(os.listdir(src_path))\nprint(sys.path)\n# ------------------------------------ #\n\nfrom Utils.CommonTools import sitk_base\nfrom Utils.CommonTools.NiiIO import read_from_DICOM_dir\nfrom Utils.PreProcessing.resampling import sitk_dummy_3D_resample\nfrom Utils.CommonTools.bbox import get_bbox, extend_bbox\nfrom Utils.post_processing import keep_largest_cervical_cc\nfrom Utils.Inference.nnunet_inference import NNUnetCTPredictor\nfrom Utils.CommonTools.dir import try_recursive_mkdir\nfrom Utils.CommonTools.NiiIO import read_from_DICOM_dir\nfrom Utils.CommonTools.sitk_base import resample, copy_nii_info, get_nii_info\n\nfrom Training.Task_301_PostProcessing_Overall.model import Model as ModelStage3\nprint(\"==> Import success\")","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2022-11-07T13:12:41.32831Z","iopub.execute_input":"2022-11-07T13:12:41.328648Z","iopub.status.idle":"2022-11-07T13:12:44.721771Z","shell.execute_reply.started":"2022-11-07T13:12:41.328575Z","shell.execute_reply":"2022-11-07T13:12:44.720428Z"}},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**Step 2: Define predictors**\n\n* PredictorStage1: Segment C1-C7\n* PredictorStage2: Segment bone fracture region\n* PredictorStage2: Predict final score using outputs from PredictorStage1 and PredictorStage2","metadata":{}},{"cell_type":"code","source":"class PredictorStage2(NNUnetCTPredictor):\n    def __init__(self, *args, **kwargs):\n        super(PredictorStage2, self).__init__(*args, **kwargs)\n\n    def resampling(self, ct_nii):\n\n        ori_spacing = ct_nii.GetSpacing()[::-1]  # to z,y,x\n        ori_size = ct_nii.GetSize()[::-1]\n        new_spacing = self.plan['plans_per_stage'][self.plan_stage]['current_spacing']\n\n        # For faster inference, but may reduce accuracy\n        new_size = [int(math.ceil(ori_size[0] * ori_spacing[0] / 0.8)), 224, 224]\n\n        new_spacing[0] = 0.8\n        new_spacing[1] = ori_size[1] * ori_spacing[1] / 224.\n        new_spacing[2] = ori_size[2] * ori_spacing[2] / 224.\n\n        do_resampling = np.any(np.abs(np.array(ori_spacing) - np.array(new_spacing)) > self.resampling_tolerance)\n        if do_resampling:\n            ct_nii = sitk_dummy_3D_resample(\n                ct_nii,\n                new_spacing=new_spacing[::-1],\n                new_size=new_size[::-1],\n                interp_xy=self.resampling_mode,\n                interp_z=sitk.sitkNearestNeighbor,\n                out_dtype=self.resampling_dtype,\n                constant_value=self.resampling_constance_value\n            )\n        else:\n            print(f'==> No necessary to do resampling ori {ori_spacing}, new: {new_spacing}')\n\n        return ct_nii\n\nclass PredictorStage3:\n    def __init__(self, list_model_pth, device, tta=False, tta_flip_axis=(4, )):\n        self.list_model_pth = list_model_pth\n        self.device = device\n        self.tta = tta\n        self.tta_flip_axis = tta_flip_axis\n\n        self.list_model = None\n\n        self.in_ch = 2\n        self.out_ch = 1\n        self.list_ch = [-1, 16, 32, 64, 128]\n\n        self.init_model()\n\n    def init_model(self):\n        with torch.no_grad():\n            self.list_model = []\n            for i in range(len(self.list_model_pth)):\n                model = ModelStage3(in_ch=self.in_ch, out_ch=self.out_ch, list_ch=self.list_ch, random_init=False)\n\n                ckpt = torch.load(self.list_model_pth[i], map_location='cpu')\n\n                model.load_state_dict(ckpt)\n                model.eval()\n                model = model.to(self.device)\n                self.list_model.append(model)\n\n                print(f'==> Init model from {self.list_model_pth[i]} to device {self.device}')\n\n    def predict(self, image):\n        with torch.no_grad():\n\n            list_pred = []\n            input_ori = image.copy()\n\n            # Flip TTA\n            if self.tta:\n                p_flip_z = (0, 1) if 2 in self.tta_flip_axis else (0,)\n                p_flip_y = (0, 1) if 3 in self.tta_flip_axis else (0,)\n                p_flip_x = (0, 1) if 4 in self.tta_flip_axis else (0,)\n            else:\n                p_flip_z = (0,)\n                p_flip_y = (0,)\n                p_flip_x = (0,)\n\n            for flip_z in p_flip_z:\n                for flip_y in p_flip_y:\n                    for flip_x in p_flip_x:\n                        patch_input = torch.from_numpy(input_ori).to(self.device).unsqueeze(0)\n\n                        # Get flip axis\n                        flip_axis = []\n                        if flip_z == 1:\n                            flip_axis.append(2)\n                        if flip_y == 1:\n                            flip_axis.append(3)\n                        if flip_x == 1:\n                            flip_axis.append(4)\n\n                        # Flip aug\n                        do_flip = (flip_z == 1) or (flip_y == 1) or (flip_x == 1)\n                        if do_flip:\n                            patch_input = torch.flip(patch_input, dims=flip_axis)\n\n                        for model in self.list_model:\n                            pred = model(patch_input)[0]\n                            pred = torch.sigmoid(pred[0, 0])\n\n                            list_pred.append(pred.cpu().numpy())\n            return np.mean(list_pred)\n\nclass DICOMReader(threading.Thread):\n    def __init__(self, func=read_from_DICOM_dir, args=()):\n        super(DICOMReader, self).__init__()\n\n        self.func = func\n        self.args = args\n\n        self.result = None\n\n    def run(self):\n        self.result = self.func(*self.args)\n\n    def get_result(self):\n        threading.Thread.join(self)\n        return self.result\n\n\nclass FractureDetector:\n    def __init__(self, predictor_stage1, predictor_stage2, predictor_stage3,  extend_roi=(5.0, 5.0, 5.0)):\n        self.predictor_stage1 = predictor_stage1\n        self.predictor_stage2 = predictor_stage2\n        self.predictor_stage3 = predictor_stage3\n        self.extend_roi = extend_roi\n\n        self.params = {\n            'alpha': [0.055, 0.044, 0.052, 0.05, 0.07, 0.077, 0.09, 0.024],\n            'beta': [0.475, 0.34, 0.37, 0.38, 0.31, 0.35, 0.38, 0.36],\n            'min_score': [0.116, 0.01, 0.015, 0.015, 0.01, 0.02, 0.032, 0.048],\n            'max_score': [0.99, 0.999, 0.993, 0.99, 1.0, 0.943, 0.997, 0.999]\n        }\n\n        self.results = {}\n\n    def get_c1_c7_bbox(self, pred, image_spacing):\n        c1_c7_bbox = get_bbox(pred > 0)\n        if c1_c7_bbox is None:\n            return None\n        c1_c7_bbox = extend_bbox(\n            c1_c7_bbox,\n            max_shape=pred.shape,\n            list_extend_length=self.extend_roi,\n            spacing=image_spacing,\n            approximate_method=np.ceil\n        )\n        return c1_c7_bbox\n\n    def predict_stage1(self, ct_nii):\n        # Resampling\n        time_start = time.time()\n        ori_nii_info = get_nii_info(ct_nii)\n        ct_nii = self.predictor_stage1.resampling(ct_nii)\n        print(f\"                  Resampling use: {time.time() - time_start}\")\n\n        # Pre_processing\n        time_start = time.time()\n        image = sitk.GetArrayFromImage(ct_nii)[np.newaxis]\n        image = self.predictor_stage1.pre_processing(image)\n        print(f\"                  Pre_processing use: {time.time() - time_start}\")\n\n        # Model forward\n        time_start = time.time()\n        pred = self.predictor_stage1.sliding_window_inference(image)\n        print(f\"                  Model forward use: {time.time() - time_start}\")\n\n        # Post processing\n        time_start = time.time()\n        pred = np.argmax(pred, axis=0)\n        pred = keep_largest_cervical_cc(pred, ct_nii.GetSpacing()[::-1])\n        pred[pred > 7] = 0\n        print(f\"                  Post processing use: {time.time() - time_start}\")\n\n        # Resampling back\n        time_start = time.time()\n        pred_nii = sitk.GetImageFromArray(np.uint8(pred))\n        pred_nii = copy_nii_info(ct_nii, pred_nii)\n        pred_nii = resample(\n            pred_nii,\n            new_spacing=ori_nii_info['spacing'],\n            new_origin=ori_nii_info['origin'],\n            new_size=ori_nii_info['size'],\n            new_direction=ori_nii_info['direction'],\n            center_origin=None,\n            interp=sitk.sitkNearestNeighbor,\n            dtype=sitk.sitkUInt8,\n            constant_value=0\n        )\n\n        pred = sitk.GetArrayFromImage(pred_nii)\n        print(f\"                  Resampling back use: {time.time() - time_start}\")\n        return pred\n\n    def predict_stage2(self, ct_nii):\n\n        # Resampling\n        time_start = time.time()\n        ori_nii_info = get_nii_info(ct_nii)\n        ct_nii = self.predictor_stage2.resampling(ct_nii)\n        print(f\"                  Resampling use: {time.time() - time_start}\")\n\n        # Pre_processing\n        time_start = time.time()\n        image = sitk.GetArrayFromImage(ct_nii)[np.newaxis]\n        image = self.predictor_stage2.pre_processing(image)\n        print(f\"                  Pre_processing use: {time.time() - time_start}\")\n        print(f\"!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!{image.shape}\")\n        # Model forward\n        time_start = time.time()\n        pred = self.predictor_stage2.sliding_window_inference(image)\n        print(f\"                  Model forward use: {time.time() - time_start}\")\n\n        # Post processing\n        time_start = time.time()\n        pred = pred[1]  # 0 for background, 1 for foreground\n        print(f\"                  Post processing use: {time.time() - time_start}\")\n\n        # Resampling back\n        time_start = time.time()\n        pred_nii = sitk.GetImageFromArray(pred)\n        pred_nii = copy_nii_info(ct_nii, pred_nii)\n        pred_nii = resample(\n            pred_nii,\n            new_spacing=ori_nii_info['spacing'],\n            new_origin=ori_nii_info['origin'],\n            new_size=ori_nii_info['size'],\n            new_direction=ori_nii_info['direction'],\n            center_origin=None,\n            interp=sitk.sitkLinear,\n            dtype=sitk.sitkFloat32,\n            constant_value=0.\n        )\n        pred = sitk.GetArrayFromImage(pred_nii)\n        print(f\"                  Resampling back use: {time.time() - time_start}\")\n\n        return pred\n    def predict_stage3(self, pred_c1_c7, pred_fracture):\n        # Resampling\n        pred_c1_c7_nii = sitk.GetImageFromArray(pred_c1_c7)\n        pred_fracture_nii = sitk.GetImageFromArray(pred_fracture)\n\n        new_size = (96, 96, 96)\n        new_spacing = list(np.array(pred_fracture_nii.GetSize()[::-1]) / np.array(new_size))\n\n        pred_fracture_nii = sitk_base.resample(\n            pred_fracture_nii,\n            new_spacing[::-1],\n            new_origin=None,\n            new_size=new_size[::-1],\n            new_direction=None,\n            center_origin=None,\n            interp=sitk.sitkLinear,\n            dtype=sitk.sitkFloat32,\n            constant_value=0\n        )\n        pred_c1_c7_nii = sitk_base.resample(\n            pred_c1_c7_nii,\n            new_spacing[::-1],\n            new_origin=None,\n            new_size=new_size[::-1],\n            new_direction=None,\n            center_origin=None,\n            interp=sitk.sitkNearestNeighbor,\n            dtype=sitk.sitkUInt8,\n            constant_value=0\n        )\n\n        pred_c1_c7 = sitk.GetArrayFromImage(pred_c1_c7_nii)\n        pred_fracture = sitk.GetArrayFromImage(pred_fracture_nii)\n\n        input_ = np.concatenate((pred_fracture[np.newaxis], pred_c1_c7[np.newaxis]), axis=0)\n        score = self.predictor_stage3.predict(input_)\n\n        return score\n    def get_score(self, pred_c1_c7, pred_fracture):\n        output = np.zeros(8, np.float32)  # Overall, C1-C7\n\n        if (pred_c1_c7 is not None) and (pred_fracture is not None):\n            # Overall, C1-C7\n            for C_i in range(8):\n                if C_i == 0:\n                    roi_fracture = pred_fracture[np.logical_and(pred_fracture >= self.params['alpha'][C_i], pred_c1_c7 > 0)]\n                else:\n                    roi_fracture = pred_fracture[np.logical_and(pred_fracture >= self.params['alpha'][C_i], pred_c1_c7 == C_i)]\n\n                if len(roi_fracture) == 0:\n                    output[C_i] = self.params['min_score'][C_i]\n                else:\n                    output[C_i] = max(self.params['min_score'][C_i],\n                                      min(self.params['max_score'][C_i],\n                                          np.percentile(roi_fracture, 100 * self.params['beta'][C_i])\n                                          )\n                                      )\n        else:\n            for C_i in range(8):\n                output[C_i] = self.params['min_score'][C_i]\n        output[0] = max(self.params['min_score'][0],  np.max(output[1:]))\n        return output\n\n    @staticmethod\n    def read_DICOM_multi_thread(list_DICOM_dirs):\n        list_thread = []\n        list_outputs = []\n\n        for DICOM_dir in list_DICOM_dirs:\n            cur_thread = DICOMReader(func=read_from_DICOM_dir, args=(DICOM_dir, ))\n            cur_thread.start()\n            list_thread.append(cur_thread)\n\n        for cur_thread in list_thread:\n            list_outputs.append(cur_thread.get_result())\n        list_thread.clear()\n\n        return list_outputs\n\n    def predict(self,\n                list_test_files,\n                num_thread=4,\n                output_dir=None\n                ):\n        \n        with torch.no_grad():\n            #try_recursive_mkdir(output_dir)\n\n            overall_time_start = time.time()\n            num_split = math.ceil(len(list_test_files) / num_thread)\n            for split_i in range(num_split):\n                cur_test_files = list_test_files[num_thread * split_i:num_thread * (split_i + 1)]\n                cur_case_ids = [test_file.split('/')[-1] for test_file in cur_test_files]\n\n                print(f\"==> Predicting {split_i}: {cur_case_ids}\")\n\n                # --------------------------------- Inference one split ------------------------------ #\n                # Step 1, Read all images\n                time_start = time.time()\n                cur_ct_niis = self.read_DICOM_multi_thread(cur_test_files)\n                print(f\"    Finish Reading use : {time.time() - time_start} seconds\")\n\n                time_start = time.time()\n                for case_i in range(len(cur_ct_niis)):\n                    case_id = cur_case_ids[case_i]\n                    ct_nii = cur_ct_niis[case_i]\n\n                    ori_nii_info = get_nii_info(ct_nii)  # Record original nii info before processing, eg. spacing, size\n\n                    # Step 2, predictor_stage1, segment C1-C7\n                    pred_1 = self.predict_stage1(ct_nii)\n\n                    #  Step 3, get c1-c7 bounding bbox\n                    c1_c7_bbox = self.get_c1_c7_bbox(pred_1, ori_nii_info['spacing'][::-1])\n\n                    # Step 4, predictor_stage2,  segment fracture\n                    if c1_c7_bbox is not None:\n                        bz, ez, by, ey, bx, ex = c1_c7_bbox\n                        roi_ct_nii = ct_nii[bx:ex + 1, by:ey + 1, bz:ez + 1]\n                        roi_pred_1 = pred_1[bz:ez + 1, by:ey + 1, bx:ex + 1]\n\n                        # Stage 2 inference\n                        roi_pred_2 = self.predict_stage2(roi_ct_nii)\n                    else:\n                        roi_pred_1 = None\n                        roi_pred_2 = None\n\n                    # Saving 5, get score\n                    if roi_pred_1 is None:\n                        roi_pred_1 = np.zeros((2, 2, 2), np.uint8)\n                        roi_pred_2 = np.zeros((2, 2, 2), np.float32)\n\n                    #sitk.WriteImage(sitk.GetImageFromArray(np.uint8(roi_pred_1)), f\"{output_dir}/{case_id}_C1_C7.nii.gz\")\n                    #sitk.WriteImage(sitk.GetImageFromArray(np.float32(roi_pred_2)), f\"{output_dir}/{case_id}_fracture.nii.gz\")\n\n                    score = self.get_score(roi_pred_1, roi_pred_2)\n                    score[0] = self.predict_stage3(roi_pred_1, roi_pred_2)\n                    \n                    self.results[case_id] = score\n                    print(f\"!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!     {score}\")\n                print(f\"    Finish this split use : {time.time() - time_start} seconds\")\n                print(f\"    Overall use : {time.time() - overall_time_start} seconds\")\n                \n                gc.collect()","metadata":{"execution":{"iopub.status.busy":"2022-11-07T13:05:40.976739Z","iopub.execute_input":"2022-11-07T13:05:40.977268Z","iopub.status.idle":"2022-11-07T13:05:41.14546Z","shell.execute_reply.started":"2022-11-07T13:05:40.977232Z","shell.execute_reply":"2022-11-07T13:05:41.143203Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**Step3: Init all models, get final predictions of all DICOM dirs in test dir**","metadata":{}},{"cell_type":"code","source":"# ----------------------- main ---------------------------- #\ntime_start = time.time()\n\n\nDATA_DIR = \"../input/rsna-2022-cervical-spine-fracture-detection/test_images\"\nSAVE_CSV = \"submission.csv\"\n\n# Predictor\nwith torch.no_grad():\n    # --------------------------- Init models --------------------------------- #\n    list_model_C1_C7_segmentation = [\n        '../input/models-final/models/stage1_0.model',\n        '../input/models-final/models/stage1_1.model',\n        '../input/models-final/models/stage1_2.model'\n    ]\n    plan_C1_C7_segmentation = '../input/plans-nnunet/stage1.pkl'\n    \n    list_model_fracture_detection = [\n        '../input/models-final/models/stage2_0.model',\n        '../input/models-final/models/stage2_1.model',\n        '../input/models-final/models/stage2_2.model',\n        '../input/models-final/models/stage2_3.model'\n    ]\n    list_model_post_processing = [\n        '../input/models-final/models/stage3_111.model',\n        '../input/models-final/models/stage3_222.model',\n        '../input/models-final/models/stage3_333.model',\n        '../input/models-final/models/stage3_444.model',\n        '../input/models-final/models/stage3_555.model'\n    ]\n    plan_fracture_detection = '../input/plans-nnunet/stage2.pkl'\n    \n    predictor_1 = NNUnetCTPredictor(\n        list_model_pth=list_model_C1_C7_segmentation,\n        plan_file=plan_C1_C7_segmentation,\n        plan_stage=-1,\n        device=torch.device('cuda:0'),\n        use_gaussian_for_sliding_window=True,\n\n        patch_size=None,\n        stride=None,\n        tta=False,\n        tta_flip_axis=(4,),\n\n        resampling_tolerance=0.01,\n        resampling_mode=sitk.sitkNearestNeighbor,\n        resampling_dtype=sitk.sitkInt16,\n        resampling_constance_value=-1024,\n\n        remove_air_CT=True\n    )\n    predictor_2 = PredictorStage2(\n        list_model_pth=list_model_fracture_detection,\n        plan_file=plan_fracture_detection,\n        plan_stage=-1,\n        device=torch.device('cuda:0'),\n        use_gaussian_for_sliding_window=True,\n\n        patch_size=(96, 224, 224),\n        stride=(96, 224, 224),\n        tta=True,\n        tta_flip_axis=(4,),\n\n        resampling_tolerance=0.01,\n        resampling_mode=sitk.sitkNearestNeighbor,\n        resampling_dtype=sitk.sitkInt16,\n        resampling_constance_value=-1024,\n\n        remove_air_CT=False,\n\n        save_dtype=np.float32\n    )\n\n    predictor_3 = PredictorStage3(\n            list_model_pth=list_model_post_processing,\n            device=torch.device('cuda:0'),\n            tta=True,\n            tta_flip_axis=(4,)\n        )\n\n    c2f_predictor = FractureDetector(\n        predictor_stage1=predictor_1,\n        predictor_stage2=predictor_2,\n        predictor_stage3=predictor_3\n    )\n\n    # ------------------------- Get all DICOM dirs ----------------------------- #\n    list_DICOM_dirs = os.listdir(DATA_DIR)\n    list_DICOM_dirs = [f\"{DATA_DIR}/{sub_dir}\" for sub_dir in list_DICOM_dirs]\n    print(f\"==> Total {len(list_DICOM_dirs)} cases\")\n    \n    list_size_DICOM_dirs = []\n    for case_i in range(len(list_DICOM_dirs)):\n        size_ = 0\n        for file in os.listdir(list_DICOM_dirs[case_i]):\n            size_ += os.path.getsize(f\"{list_DICOM_dirs[case_i]}/{file}\")\n        list_size_DICOM_dirs.append(size_)\n\n    list_DICOM_dirs = list(np.array(list_DICOM_dirs)[np.argsort(list_size_DICOM_dirs)[::-1]])\n    print(f\"==> Sort DICOM dirs by size\")\n    print(list_DICOM_dirs[0])\n    \n    c2f_predictor.predict(\n        list_test_files=list_DICOM_dirs\n    )\n\n    results = c2f_predictor.results\n\n    # To CSV\n    results_csv = {'row_id':[],\n                   'fractured':[]\n                   }\n    for case_id in results.keys():\n        for C_i in range(1, 8):\n            results_csv['row_id'].append(f\"{case_id}_C{C_i}\")\n            results_csv['fractured'].append(f\"{results[case_id][C_i]}\")\n\n        results_csv['row_id'].append(f\"{case_id}_patient_overall\")\n        results_csv['fractured'].append(f\"{results[case_id][0]}\")\n\n    results_csv = pd.DataFrame(results_csv)\n    results_csv.to_csv(SAVE_CSV, index=False)\n\n    print(f'==> Finish using time: {time.time() - time_start}')\n\n\n","metadata":{"execution":{"iopub.status.busy":"2022-11-07T13:05:41.151893Z","iopub.execute_input":"2022-11-07T13:05:41.152576Z"},"trusted":true},"execution_count":null,"outputs":[]}]}