{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.14","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":29653,"databundleVersionId":2420395,"sourceType":"competition"},{"sourceId":2542390,"sourceType":"datasetVersion","datasetId":1541666},{"sourceId":2723611,"sourceType":"datasetVersion","datasetId":1659960},{"sourceId":9638426,"sourceType":"datasetVersion","datasetId":5885258},{"sourceId":9649374,"sourceType":"datasetVersion","datasetId":5893578},{"sourceId":9656364,"sourceType":"datasetVersion","datasetId":5898939},{"sourceId":138532,"sourceType":"modelInstanceVersion","isSourceIdPinned":true,"modelInstanceId":117310,"modelId":140539}],"dockerImageVersionId":30786,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\n%pip install monai===0.7.0\n%pip install segmentation-models-pytorch=0.2.0\n%pip install gdown==3.6.4\n%pip install segmentation_models_pytorch==0.2.0\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.\nimport os\n","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2024-10-18T06:17:18.557051Z","iopub.execute_input":"2024-10-18T06:17:18.557353Z","iopub.status.idle":"2024-10-18T06:18:05.998205Z","shell.execute_reply.started":"2024-10-18T06:17:18.557305Z","shell.execute_reply":"2024-10-18T06:18:05.997052Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# import tarfile\n# file = tarfile.open('../input/brats-2021-task1/BraTS2021_Training_Data.tar')\n\n# file.extractall('./TrainingData')\n# file.close()\nimport shutil\ntry:\n    shutil.rmtree(\"/kaggle/working/TrainingData\")\nexcept FileNotFoundError:\n    print(\"The directory does not exist.\")\nexcept OSError as e:\n    print(f\"Error: {e}\")","metadata":{"execution":{"iopub.status.busy":"2024-10-18T05:34:53.189544Z","iopub.execute_input":"2024-10-18T05:34:53.190347Z","iopub.status.idle":"2024-10-18T05:34:53.195972Z","shell.execute_reply.started":"2024-10-18T05:34:53.190310Z","shell.execute_reply":"2024-10-18T05:34:53.194987Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Segementation","metadata":{}},{"cell_type":"code","source":"import os\nimport json\n\nimport logging\nimport pandas as pd \nimport numpy as np\nimport cv2\nimport matplotlib.pyplot as plt\nimport seaborn as sns\nfrom tqdm import tqdm\n\n\nfrom monai.data import DataLoader, ImageDataset\n\nfrom monai.transforms import (\n    AddChannel,\n    Compose,\n    Resize,\n    Transform\n)\n\n\nfrom sklearn.model_selection import GroupKFold\n\nimport torch.nn.functional as F\n\nfrom multiprocessing import Pool\n\n\nSETTINGS = {\n    \"DICOM_DATA_DIR\":\"/kaggle/input/rsna-miccai-brain-tumor-radiogenomic-classification\",\n    \"TASK1_DIR\":\"/kaggle/working/TrainingData\", \n    \"CLASSIFICATION_RAW_JPG\":\"/kaggle/input/miccaibraintumorjpgdata\",\n    \"SEGMENT_DATA_DIR\":\"data/processed_segmentation_data\", \n    \"CLASSIFICATION_DATA_DIR\":\"data/processed_classification_data\",\n    \"KFOLD_PATH\":\"/kaggle/input/sub-file-for-tumor/data/train_stratifiedgroupkfold.csv\", \n    \"SEGMENT_MODEL_DIR\":\"/kaggle/working/models/densenet121_2d_segment\",\n    \"CLASSIFICATION_MODEL_DIR\":\"models/eca_nfnet_l0_2d_classification\",\n    \"TEMP_DATA_DIR\":\"temp\",\n    \"TEST_PREDICTION_FILE\":\"data/test_prediction.csv\"\n}\n\nIM_FOLDER_TASK1 = SETTINGS['TASK1_DIR']\n\nRUN_FOLDS = [0]\nKFOLD_PATH = SETTINGS['KFOLD_PATH']\n\nSEED = 67\nN_PROCESSES = 4\n\nOUT_FOLDER = SETTINGS['SEGMENT_DATA_DIR']\n\nPLANES = ['sagital', 'coronal', 'axial']\nMRI_TYPES = ['t1', 't1ce', 't2', 'flair']\n\nimport tarfile\n\n\n\n\n# ============ Helper functions ===========\nclass ScaleRange(Transform):\n    def __init__(self, new_max = 255.0):\n        super(ScaleRange, self).__init__()\n        self.new_max = new_max\n        \n    def __call__(self, data):\n        dmin, dmax = data.min(), data.max()\n        return (data - dmin) / (dmax-dmin) * self.new_max\n\nclass ConvertToMultiChannelBasedOnBratsClasses(Transform):\n    \"\"\"\n    Convert labels to multi channels based on brats classes:\n    label 2 is the peritumoral edema\n    label 4 is the GD-enhancing tumor\n    label 1 is the necrotic and non-enhancing tumor core\n    The possible classes are TC (Tumor core), WT (Whole tumor)\n    and ET (Enhancing tumor).\n    Ehancing Tumor (ET) = enhancing tumor\n    Tumor Core (TC) = enhancing tumor + necrotic\n    Whole Tumor (WT) = enhancing tumor + necrotic + edema    \n    \"\"\"\n\n    def __call__(self, masks):\n        '''This time we only use 2 label: 0 - WT and 1 - ET'''\n        result = []\n\n        # merge labels 1, 2 and 4 to construct WT\n        result.append(\n            np.logical_or(\n                np.logical_or(masks == 1, masks == 2), masks == 4\n            )\n        )\n        # label 4 is ET\n        result.append(masks == 4)\n        \n        return np.stack(result, axis=0).astype(np.float32)\n\ndef get_non_0_voxels_and_masks(voxels, masks_2channels, ax=0, min_avg=0.01):\n    '''Get non-empty slices from the 3D mask\n        A 2D slice is considered to be empty if its mean pixel value < min_avg'''\n    masks = np.logical_or(masks_2channels[0], masks_2channels[1])\n    remain_axes = tuple([i for i in range(len(voxels.shape)) if i != ax])\n    ax_mean = masks.mean(axis=remain_axes)\n    ax_non_0_inds = ax_mean > min_avg\n    if(ax==0):\n        return voxels[ax_non_0_inds], masks_2channels[:, ax_non_0_inds, :, :]\n    if(ax==1):\n        return voxels[:,ax_non_0_inds,:], masks_2channels[:, :, ax_non_0_inds,:]\n    if(ax==2):\n        return voxels[:,:,ax_non_0_inds], masks_2channels[:,:,:,ax_non_0_inds]\n    \ndef sampling_slices(non_0_voxels, non_0_masks, ax=0, keep_rate=0.1):\n    '''Nearby slices are similar to each other, we use sample to only get the different ones'''\n    total_slices = non_0_voxels.shape[ax]\n    T = max(round(total_slices * keep_rate), 1)\n    sampling_inds = np.arange(0, total_slices, T)\n    \n    if(ax==0):\n        return non_0_voxels[sampling_inds], non_0_masks[:, sampling_inds, :, :]\n    if(ax==1):\n        return non_0_voxels[:, sampling_inds, :], non_0_masks[:, :, sampling_inds, :]\n    if(ax==2):\n        return non_0_voxels[:, :, sampling_inds], non_0_masks[:, :, :, sampling_inds]\n    \n    \ndef process_one_patient(voxels, masks, patient_id):\n    '''Perform slicing 2D images and tumor masks for this patient'''\n    current_list_patient_id = []\n    current_list_plane = []\n    current_list_mri_type = []\n    current_list_slice_index = []\n    current_list_file_path = []\n    current_list_segfile_path = []\n    \n    for ax, plane in enumerate(PLANES):\n        non_0_voxels, non_0_masks = get_non_0_voxels_and_masks(voxels, masks, ax=ax)\n        if(non_0_voxels.shape[ax]==0):\n            print(f'Cannot get any slice in patient: {patient_id}, plane: {plane} due to the masks are too small')\n            continue\n        sampled_non_0_voxels, sampled_non_0_masks = sampling_slices(non_0_voxels, non_0_masks, ax=ax)\n\n        for j in range(sampled_non_0_voxels.shape[ax]):\n            file_path = os.path.join(OUT_FOLDER + '/2D_slice_data/', \n                                     f'BraTS2021_{patient_id:05d}',\n                                     f'BraTS2021_{patient_id:05d}_{mri_type}',\n                                    f'BraTS2021_{patient_id:05d}_{mri_type}_{plane}_{j:03d}')\n            seg_file_path = os.path.join(OUT_FOLDER + '/2D_slice_data/', \n                                     f'BraTS2021_{patient_id:05d}',\n                                    f'BraTS2021_{patient_id:05d}_segmask',\n                                    f'BraTS2021_{patient_id:05d}_segmask_{plane}_{j:03d}')\n\n            os.makedirs(os.path.dirname(file_path), exist_ok=True)\n            os.makedirs(os.path.dirname(seg_file_path), exist_ok=True)\n            \n            if(ax==0):\n                np.save(file_path, sampled_non_0_voxels[j])\n                np.save(seg_file_path, sampled_non_0_masks[:,j])\n            elif(ax==1):\n                np.save(file_path, sampled_non_0_voxels[:,j,:])\n                np.save(seg_file_path, sampled_non_0_masks[:,:,j,:])\n            elif(ax==2):\n                np.save(file_path, sampled_non_0_voxels[:,:,j])\n                np.save(seg_file_path, sampled_non_0_masks[:,:,:,j])\n            else:\n                raise ValueError('No such ax')\n\n            current_list_patient_id.append(patient_id)\n            current_list_plane.append(plane)\n            current_list_mri_type.append(mri_type)\n            current_list_slice_index.append(j)\n            current_list_file_path.append(file_path)\n            current_list_segfile_path.append(seg_file_path)\n\n    return current_list_patient_id, current_list_plane, current_list_mri_type,  \\\n            current_list_slice_index, current_list_file_path, current_list_segfile_path\n\n\ndef update(args):\n    global list_patient_id, list_plane, list_mri_type, list_slice_index, list_file_path, list_segfile_path\n    pbar.update()\n    current_list_patient_id, current_list_plane, current_list_mri_type,  \\\n            current_list_slice_index, current_list_file_path, current_list_segfile_path = args\n    \n    list_patient_id += current_list_patient_id\n    list_plane += current_list_plane\n    list_mri_type += current_list_mri_type\n    list_slice_index += current_list_slice_index\n    list_file_path += current_list_file_path\n    list_segfile_path += current_list_segfile_path\n\n\ndef error(e):\n    print(e)\n        \n# =========================================\n\n# ============ Read meta data =============\nfold_df = pd.read_csv(KFOLD_PATH)\nfold_df['pfolder'] = fold_df.BraTS21ID.map(lambda x: f'BraTS2021_{x:05d}')\n\nPATIENT_DIRS = []\nfor p in os.listdir(IM_FOLDER_TASK1):\n    try:\n       \n        int(p.split('_')[-1])\n        PATIENT_DIRS.append(p)\n    except:\n        print('Non patient dir:', p)\n\ndf = pd.DataFrame(PATIENT_DIRS, columns=['pfolder'])\n\n\ndf['BraTS21ID'] = df['pfolder'].map(lambda x: int(x.split('_')[-1]))\ndf = df.dropna()\n\ndf = df[~df.BraTS21ID.isin(fold_df.BraTS21ID.tolist())]\n\nfor t in MRI_TYPES:\n    df[f'{t}_data_path'] = df.pfolder.map(lambda x: os.path.join(IM_FOLDER_TASK1, x, x+f'_{t}.nii.gz'))\ndf['seg_label_path'] = df.pfolder.map(lambda x: os.path.join(IM_FOLDER_TASK1, x, x+f'_seg.nii.gz'))\n\n# =========================================\n\n\n# ============ Create a nii gz file loader ==========\ntransforms = Compose([ScaleRange()])\n\nseg_transforms = Compose([ConvertToMultiChannelBasedOnBratsClasses(),\n                         ])\n\nmri_type = MRI_TYPES[0]\n# Define nifti dataset, data loader\ndataset = ImageDataset(image_files=df[f'{mri_type}_data_path'].tolist(),\n                             seg_files = df.seg_label_path.tolist(),\n                             seg_transform=seg_transforms,\n                            transform=transforms\n                      )\n# =====================================================\n\n\n\n# ========== Perform slicing data and mask ============\n\nfor mri_type in MRI_TYPES:\n    dataset = ImageDataset(image_files=df[f'{mri_type}_data_path'].tolist(),\n                                 seg_files = df.seg_label_path.tolist(),\n                                   labels = df['BraTS21ID'].tolist(),\n                                 seg_transform=seg_transforms,\n                                transform=transforms\n                          )\n    \n    os.makedirs(OUT_FOLDER + '/2D_slice_data/', exist_ok=True)\n\n    list_patient_id = []\n    list_plane = []\n    list_mri_type = []\n    list_slice_index = []\n    list_file_path = []\n    list_segfile_path = []\n\n    pool = Pool(processes=N_PROCESSES)   \n\n    iterations = range(len(dataset))\n    pbar = tqdm(iterations)\n\n    for i in iterations:\n        voxels, masks, patient_id = dataset[i]\n        pool.apply_async(\n            process_one_patient,\n            args=(voxels, masks, patient_id),\n            callback=update,\n            error_callback=error,\n        )\n\n    pool.close()\n    pool.join()\n    pbar.close()\n    \nout_df = pd.DataFrame({\n    'BraTS21ID':list_patient_id,\n    'mri_type':list_mri_type,\n    'plane':list_plane,\n    'slice_index':list_slice_index,\n    'file_path':list_file_path,\n    'segfile_path':list_segfile_path\n})\n\nout_df.to_csv(os.path.join(OUT_FOLDER, 'segment_meta.csv'))\n# =====================================================\n\n\n# ================= Kfold split ====================\nkfold = GroupKFold(n_splits=5)\n\ni = 0\nfor train_ind, valid_ind in kfold.split(out_df,out_df,out_df['BraTS21ID']):\n    out_df.loc[valid_ind, 'fold'] = i\n    i+=1\nout_df.to_csv(f'{OUT_FOLDER}/segment_meta_groupkfold.csv', index=False)\n# ==================================================\n    \n    \n  ","metadata":{"scrolled":true,"execution":{"iopub.status.busy":"2024-10-18T05:34:53.197172Z","iopub.execute_input":"2024-10-18T05:34:53.197695Z","iopub.status.idle":"2024-10-18T05:35:13.862192Z","shell.execute_reply.started":"2024-10-18T05:34:53.197642Z","shell.execute_reply":"2024-10-18T05:35:13.860668Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import numpy as np\nimport matplotlib.pyplot as plt\n\nimg=np.load(\"/kaggle/working/data/processed_segmentation_data/2D_slice_data/BraTS2021_01295/BraTS2021_01295_t1/BraTS2021_01295_t1_sagital_004.npy\")\nmask=np.load(\"/kaggle/working/data/processed_segmentation_data/2D_slice_data/BraTS2021_01295/BraTS2021_01295_segmask/BraTS2021_01295_segmask_sagital_007.npy\")\nprint(img.shape)\n\n\nplt.figure()\nplt.imshow(mask[1, :, :])\nplt.imshow(img, cmap=\"gray\", alpha=0.5)\n\n","metadata":{"execution":{"iopub.status.busy":"2024-10-18T05:37:54.047160Z","iopub.execute_input":"2024-10-18T05:37:54.048044Z","iopub.status.idle":"2024-10-18T05:37:54.266637Z","shell.execute_reply.started":"2024-10-18T05:37:54.048000Z","shell.execute_reply":"2024-10-18T05:37:54.265333Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Models","metadata":{}},{"cell_type":"code","source":"import gdown\nimport os\n\n# os.makedirs('pretrained_models', exist_ok=True)\n# url = 'https://drive.google.com/uc?id=14dmJ2HEEsdDzcsbJC_qTAQPxOGSSkDY3'\n# output = 'pretrained_models/eca_nfnet_l0.pth'\n# gdown.download(url, output, quiet=False)\n\n# os.makedirs('pretrained_models', exist_ok=True)\n# url = 'https://drive.google.com/uc?id=1hV4HECWeiHpFkBTES2nYPpnlYM0dApRJ'\n# output = 'pretrained_models/unet_pp_densenet121_2channels_out.pth'\n# gdown.download(url, output, quiet=False)\n\n# os.makedirs('models/densenet121_2d_segment', exist_ok=True)\n# url = 'https://drive.google.com/uc?id=12EVeyHI_kQlryAp6554Au4S1pt1ektnY'\n# output = 'models/densenet121_2d_segment/Fold0_densenet121_2d_segment.pth'\n# gdown.download(url, output, quiet=False)\n\n# os.makedirs('models/eca_nfnet_l0_2d_classification/T1w', exist_ok=True)\n# url = 'https://drive.google.com/uc?id=1wFJAurtdm8_G-nV2DVUQlWy7X4h4lBjo'\n# output = 'models/eca_nfnet_l0_2d_classification/T1w/T1w_Fold0_eca_nfnet_l0_2d_classification.pth'\n# gdown.download(url, output, quiet=False)\n\n# os.makedirs('models/eca_nfnet_l0_2d_classification/T1wCE', exist_ok=True)\n# url = 'https://drive.google.com/uc?id=1fmMScOnUyDWNDEYewtX7CpDMz6EFr3Vd'\n# output = 'models/eca_nfnet_l0_2d_classification/T1wCE/T1wCE_Fold0_eca_nfnet_l0_2d_classification.pth'\n# gdown.download(url, output, quiet=False)\n\n# os.makedirs('models/eca_nfnet_l0_2d_classification/T2w', exist_ok=True)\n# url = 'https://drive.google.com/uc?id=1JPAcR2vCDtvblDqPljx_poZNUTC00OET'\n# output = 'models/eca_nfnet_l0_2d_classification/T2w/T2w_Fold0_eca_nfnet_l0_2d_classification.pth'\n# gdown.download(url, output, quiet=False)\n\n# os.makedirs('models/eca_nfnet_l0_2d_classification/FLAIR', exist_ok=True)\n# url = 'https://drive.google.com/uc?id=12UBvr4ewswPWzf-8OyM4GjqYugonlTjY'\n# output = 'models/eca_nfnet_l0_2d_classification/FLAIR/FLAIR_Fold0_eca_nfnet_l0_2d_classification.pth'\n# gdown.download(url, output, quiet=False)\n","metadata":{"execution":{"iopub.status.busy":"2024-10-18T05:35:13.865267Z","iopub.status.idle":"2024-10-18T05:35:13.865664Z","shell.execute_reply.started":"2024-10-18T05:35:13.865451Z","shell.execute_reply":"2024-10-18T05:35:13.865470Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Utility","metadata":{}},{"cell_type":"code","source":"import pandas as pd\nimport numpy as np\nimport json\nimport re\nimport torch\nimport random\nimport sys\nimport os\nimport matplotlib.pyplot as plt\nimport logging\n\nfrom tqdm import tqdm\nfrom copy import deepcopy\nfrom typing import List\n\ndef seed_torch(seed=42):\n    random.seed(seed)\n    os.environ['PYTHONHASHSEED'] = str(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    torch.backends.cudnn.deterministic = True\n\n\ndef get_logger(name, path, mode='a'):\n    logger = logging.getLogger(name)  \n\n    if not logger.hasHandlers():\n        # set log level\n        logger.setLevel(logging.INFO)\n\n        # define file handler and set formatter\n        file_handler = logging.FileHandler(path, mode=mode)\n        formatter    = logging.Formatter('%(asctime)s : %(levelname)s : %(name)s : %(message)s')\n        file_handler.setFormatter(formatter)\n\n        # add file handler to logger\n        logger.addHandler(file_handler)\n    \n    return logger\n    \n    \ndef log_and_print(logger, obj):\n    print(obj)\n    logger.info(obj)    \n    \ndef init_progress_dict(metrics):\n    progress_dict = dict()\n    for metric in metrics:\n        progress_dict[f'train_{metric}'] = []\n        progress_dict[f'valid_{metric}'] = []\n    return progress_dict\n\ndef log_to_progress_dict(progress_dict, metric_dict):\n    for k, v in metric_dict.items():\n        progress_dict[k].append(v)\n       \n    return progress_dict\n\ndef save_progress(progress_dict, out_folder, out_folder_name, fold, show=False):\n    metric_names = list(progress_dict.keys())\n    epochs = len(progress_dict[metric_names[0]])+1\n    \n    # plot figure and save the progress chart\n    n_cols = 4\n    n_rows = int(np.ceil(len(metric_names) / 2 / n_cols))\n    \n    plt.figure(figsize=(7*n_cols, 7*n_rows))\n    \n    for i in range(0, len(metric_names), 2):\n        plt.subplot(n_rows,n_cols,int(i/2+1))\n\n        plt.plot(range(1, epochs), progress_dict[metric_names[i]])\n        plt.plot(range(1, epochs), progress_dict[metric_names[i+1]])\n        plt.legend([metric_names[i], metric_names[i+1]])\n        plt.xlabel('Epoch')\n        plt.title(f'{metric_names[i]} and {metric_names[i+1]}')\n\n    save_name = f'training_progress_{out_folder_name}_fold{fold}'\n    plt.savefig(os.path.join(out_folder, save_name+'.jpg'))\n\n    if(show):\n        plt.show()\n\n    pd.DataFrame({'epoch':range(1, epochs), **progress_dict}).to_csv(os.path.join(out_folder, save_name+'.csv'), index=False)\n\n    \ndef check_mem(cuda_device):\n    devices_info = os.popen('\"/usr/bin/nvidia-smi\" --query-gpu=memory.total,memory.used --format=csv,nounits,noheader').read().strip().split(\"\\n\")\n    total, used = devices_info[int(cuda_device)].split(',')\n    return total,used\n\ndef occumpy_mem(cuda_device):\n    total, used = check_mem(cuda_device)\n    total = int(total)\n    used = int(used)\n    max_mem = int(total * 0.9)\n    block_mem = max_mem - used\n    x = torch.cuda.FloatTensor(256,1024,block_mem)\n    del x","metadata":{"execution":{"iopub.status.busy":"2024-10-18T06:35:19.890598Z","iopub.execute_input":"2024-10-18T06:35:19.891403Z","iopub.status.idle":"2024-10-18T06:35:19.912138Z","shell.execute_reply.started":"2024-10-18T06:35:19.891359Z","shell.execute_reply":"2024-10-18T06:35:19.911138Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Training","metadata":{}},{"cell_type":"code","source":"\nimport os\nimport logging\nimport pandas as pd \nimport numpy as np\nimport cv2\nimport matplotlib.pyplot as plt\nimport seaborn as sns\nfrom tqdm import tqdm\nimport torch\nfrom torch import nn\n\nfrom torch.optim import AdamW\nfrom torch.optim.lr_scheduler import ReduceLROnPlateau, CosineAnnealingWarmRestarts, CosineAnnealingLR\n\nimport torch.nn.functional as F\n\nimport json\nimport gc\nfrom sklearn.metrics import roc_auc_score, accuracy_score\n\nfrom segmentation_models_pytorch.unetplusplus.model import UnetPlusPlus\nfrom segmentation_models_pytorch.losses import DiceLoss\nfrom segmentation_models_pytorch.utils.metrics import IoU\n\nimport albumentations as A\nfrom albumentations.pytorch.transforms import ToTensorV2\n\nimport argparse\n\n# parser = argparse.ArgumentParser(description='Insert some arguments')\n# parser.add_argument('--gpu', type=int,\n#                     help='GPU ID', default=0)\n# parser.add_argument('--batch_size', type=int,\n#                     help='Batch size', default=128)\n# parser.add_argument('--n_workers', type=int,\n#                     help='Number of parrallel workers', default=8)\n# args = parser.parse_args()\nargv = sys.argv[1:]  # Exclude the first argument (script name)\nif '-f' in argv:\n    argv = argv[:argv.index('-f')]  # Ignore everything after '-f'\n\n# Set up the parser\nparser = argparse.ArgumentParser(description='Insert some arguments')\nparser.add_argument('--gpu', type=int, help='GPU ID', default=0)\nparser.add_argument('--batch_size', type=int, help='Batch size', default=128)\nparser.add_argument('--n_workers', type=int, help='Number of parallel workers', default=8)\n\n# Parse the cleaned arguments\nargs = parser.parse_args(argv)\n\n\n# with open('SETTINGS.json', 'r') as f:\nSETTINGS = {\n    \"DICOM_DATA_DIR\":\"/kaggle/input/rsna-miccai-brain-tumor-radiogenomic-classification\",\n    \"TASK1_DIR\":\"/kaggle/working/TrainingData\", \n    \"CLASSIFICATION_RAW_JPG\":\"/kaggle/input/miccaibraintumorjpgdata\",\n    \"SEGMENT_DATA_DIR\":\"/kaggle/input/segmentaed-csv\", \n    \"CLASSIFICATION_DATA_DIR\":\"data/processed_classification_data\",\n    \"KFOLD_PATH\":\"/kaggle/input/sub-file-for-tumor/data/train_stratifiedgroupkfold.csv\", \n    \"SEGMENT_MODEL_DIR\":\"/kaggle/working/models/densenet121_2d_segment\",\n    \"CLASSIFICATION_MODEL_DIR\":\"models/eca_nfnet_l0_2d_classification\",\n    \"TEMP_DATA_DIR\":\"temp\",\n    \"TEST_PREDICTION_FILE\":\"data/test_prediction.csv\"\n        }\n\nFOLDER = SETTINGS['SEGMENT_DATA_DIR']\nMETA_FILE_PATH = os.path.join(SETTINGS['SEGMENT_DATA_DIR'], 'segment_meta_groupkfold.csv')\n\nRUN_FOLDS = [0]\nSEED = 67\nDIM = (224, 224, 3)\n# N_WORKERS = args.n_workers\nBATCH_SIZE = 32\nBASE_LR = 1e-4\nNUM_EPOCHS = 50\nPATIENT = 10\nSAMPLE = None\nDEVICE = torch.device(f'cuda:{args.gpu}')\n\nPARENT_OUT_FOLDER = f'models/'    \n\nCANDIDATES = [\n    {\n        'backbone_name':'densenet121',\n        'ver_note':'2d_segment',\n        'backbone_pretrained':f'/kaggle/input/pretranied_models/tensorflow2/default/1/unet_pp_densenet121_2channels_out (1).pth',\n        'batch_size':BATCH_SIZE,\n        'warm_up_epochs':5,\n    },\n]\n\nimport sys\n# from utils.general import seed_torch, init_progress_dict, log_to_progress_dict, save_progress, log_and_print, get_logger\n\n# seed every thing\nseed_torch(SEED)\n\n# ================= Some helper functions ====================\nclass BrainSegment2DDataset(torch.utils.data.Dataset):\n    \n    def __init__(self, csv, transforms=None):\n        self.csv = csv.reset_index(drop=True)\n        self.augmentations = transforms\n\n    def __len__(self):\n        return self.csv.shape[0]\n\n    def __getitem__(self, index):\n        row = self.csv.iloc[index]\n        image = np.load(row['file_path']+'.npy')\n        image = np.stack([image]*3, axis=-1)\n        mask = np.load(row['segfile_path']+'.npy')\n        mask = np.stack([mask[0], mask[1]], axis=-1)\n        \n        if self.augmentations:\n            augmented = self.augmentations(image=image, mask=mask)\n            image = augmented['image']\n            mask = augmented['mask']\n            mask = mask.permute(2,0,1)\n        \n        return image, mask\n    \n\ndef get_train_transforms(candidate):\n    dim = candidate.get('dim', DIM)\n    return A.Compose(\n        [\n            A.HorizontalFlip(p=0.5),\n            A.VerticalFlip(p=0.5),\n            A.ShiftScaleRotate(p=0.5),\n     \n            A.Resize(width=dim[1], height=dim[0], always_apply=True),\n            A.Normalize(),\n            ToTensorV2(p=1.0)\n        ]\n    )\n\ndef get_valid_transforms(candidate):\n    dim = candidate.get('dim', DIM)\n    return A.Compose(\n        [\n            A.Resize(width=dim[1], height=dim[0], always_apply=True),\n            A.Normalize(),\n            ToTensorV2(p=1.0)\n        ]\n    )\n\ndef get_model(candidate):\n    model = UnetPlusPlus(\n        encoder_name = candidate['backbone_name'],\n        encoder_depth = 5,\n        encoder_weights = None,\n        classes = 2,\n        activation = 'sigmoid',\n    )\n\n    weight_path = candidate.get('backbone_pretrained')\n    if(weight_path is not None):\n        print('Load pretrained:', weight_path)\n        model.load_state_dict(torch.load(weight_path, map_location='cpu'))\n        \n    return model\n\n\nclass AverageMeter(object):\n    def __init__(self):\n        self.reset()\n    \n    def reset(self):\n        self.val = 0\n        self.avg = 0\n        self.sum = 0\n        self.count = 0\n    \n    def update(self, val, n=1):\n        self.val = val\n        self.sum += val * n\n        self.count += n\n        self.avg = self.sum / self.count\n        \ndef train_valid_fn(dataloader,model,criterion,iou_metric,optimizer=None,device='cuda:0',\n                            scheduler=None,epoch=0, mode='train', scaler=None):\n    '''Perform model training'''\n    if(mode=='train'):\n        model.train()\n    elif(mode=='valid'):\n        model.eval()\n    else:\n        raise ValueError('No such mode')\n        \n    loss_score = AverageMeter()\n    iou_score = AverageMeter()\n    \n    tk0 = tqdm(enumerate(dataloader), total=len(dataloader))\n    for i, batch in tk0:\n        if(mode=='train'):\n            optimizer.zero_grad()\n            \n        # input, gt\n        images, gt_masks = batch\n        images = images.to(DEVICE)\n        gt_masks = gt_masks.to(DEVICE)\n\n        with torch.cuda.amp.autocast():\n            # prediction\n            pred_masks = model(images)\n\n            # compute loss\n            loss = criterion(y_true=gt_masks, y_pred=pred_masks)\n            \n            # compute metric\n            iou = iou_metric(pred_masks, gt_masks)\n        \n        if(mode=='train'):\n            scaler.scale(loss).backward()\n            scaler.step(optimizer)\n            scaler.update()\n        \n        loss_score.update(loss.detach().cpu().item(), dataloader.batch_size)\n        iou_score.update(iou.detach().cpu().item(), dataloader.batch_size)\n        \n        if(mode=='train'):\n            tk0.set_postfix(Loss_Train=loss_score.avg, IOU_Train=iou_score.avg, \n                            Epoch=epoch, LR=optimizer.param_groups[0]['lr'])\n        elif(mode=='valid'):\n            tk0.set_postfix(Loss_Valid=loss_score.avg, IOU_Valid=iou_score.avg, Epoch=epoch)\n        \n        del batch, images, gt_masks, pred_masks, loss, iou\n        torch.cuda.empty_cache()\n        \n    if(mode=='train'):\n        if(scheduler.__class__.__name__ == 'CosineAnnealingWarmRestarts'):\n            scheduler.step(epoch=epoch)\n        elif(scheduler.__class__.__name__ == 'ReduceLROnPlateau'):\n            scheduler.step(loss_score.avg)\n    \n    return loss_score.avg, iou_score.avg\n\ndef dfs_freeze(module):\n    for name, child in module.named_children():\n        for param in child.parameters():\n            param.requires_grad = False\n        dfs_freeze(child)\n        \ndef dfs_unfreeze(module):\n    for name, child in module.named_children():\n        for param in child.parameters():\n            param.requires_grad = True\n        dfs_unfreeze(child)\n# ===========================================================\n  \n# ===========================================================\n        \n    \n# ================ Read metadata =================\ndf = pd.read_csv(META_FILE_PATH)\n# ================================================\n\n\n# ============================ Training ==============================\nfor candidate in CANDIDATES:\n    print(f\"######################### Candidate: {candidate['backbone_name']} ############################\")\n    run_folds = candidate.get('run_folds', RUN_FOLDS)\n    \n    parent_out_folder = candidate.get('parent_out_folder', PARENT_OUT_FOLDER)\n    ver_note = candidate['ver_note']\n    out_folder_name = f\"{candidate['backbone_name']}_{ver_note}\"\n    out_folder = os.path.join(parent_out_folder, out_folder_name)\n\n    os.makedirs(out_folder, exist_ok=True)\n    \n    for valid_fold in run_folds:\n        # Read data\n        if(SAMPLE):\n            df = df.sample(SAMPLE, random_state=SEED)\n\n        train_df = df[df.fold!=valid_fold]\n        valid_df = df[df.fold==valid_fold]\n\n        print(f'\\n\\n================= Fold {valid_fold} ==================')\n        print(f'Number of training images: {len(train_df)}. Number of valid images: {len(valid_df)}')\n#         print(\"filepath\",traindf)\n        \n        train_dataset = BrainSegment2DDataset(train_df, get_train_transforms(candidate))\n        valid_dataset = BrainSegment2DDataset(train_df, get_valid_transforms(candidate))\n        \n        batch_size = candidate.get('batch_size', BATCH_SIZE)\n        train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=4)\n        valid_loader = torch.utils.data.DataLoader(valid_dataset, batch_size=batch_size, shuffle=False, num_workers=4)\n        \n        # model\n        model = get_model(candidate)\n        # freeze layer\n        dfs_freeze(model.encoder)\n        print(' -------- Start warm up process ----------')\n        print('Freeze encoder')\n        model.to(DEVICE)\n        print()\n        \n        # Optimizer and scheduler\n        base_lr = candidate.get('base_lr', BASE_LR)\n        optim = AdamW(filter(lambda p: p.requires_grad, model.parameters()), lr=BASE_LR)\n\n        num_training_steps = NUM_EPOCHS * len(train_loader)\n        lr_scheduler = ReduceLROnPlateau(optimizer=optim)\n\n        # loss\n        criterion = DiceLoss(mode='binary', from_logits=False)\n        iou_metric = IoU()\n\n        # use amp to accelerate training\n        scaler = torch.cuda.amp.GradScaler()\n\n        # Logging\n        logger = get_logger(\n            name = f'training_log_fold{valid_fold}.txt',\n            path=os.path.join(out_folder, f'training_log_fold{valid_fold}.txt')\n        )\n\n        best_valid_loss = 9999\n        best_valid_ep = 0\n        patient = PATIENT\n\n        progress_dict = init_progress_dict(['loss', 'IOU'])\n\n        start_ep = candidate.get('warm_start_ep', 1)\n        print('Start ep:', start_ep)\n\n        # warm up epochs\n        warm_up_epochs = candidate.get('warm_up_epochs', 1)\n\n        \n        for epoch in range(start_ep, NUM_EPOCHS+1):\n            if (epoch==warm_up_epochs+1):\n                print(' -------- Finish warm up process ----------')\n                print('Unfreeze encoder')\n                dfs_unfreeze(model.encoder)\n                optim = AdamW(filter(lambda p: p.requires_grad, model.parameters()), lr=BASE_LR)\n                lr_scheduler = ReduceLROnPlateau(optimizer=optim)\n                \n            # =============== Training ==============\n            train_loss, train_iou = train_valid_fn(train_loader,model, criterion, iou_metric,\n                                                    optimizer=optim,device=DEVICE,\n                                                    scheduler=lr_scheduler,epoch=epoch,mode='train',\n                                                  scaler=scaler)\n            valid_loss, valid_iou = train_valid_fn(valid_loader,model, criterion, iou_metric,\n                                                     device=DEVICE, \n                                                     epoch=epoch,mode='valid',\n                                                  scaler=scaler)\n\n            current_lr = optim.param_groups[0]['lr']\n            log_line = f'Model: {out_folder_name}. Epoch: {epoch}. '\n            log_line += f'Train loss:{train_loss} - Valid loss: {valid_loss}. '\n            log_line += f'Train IOU:{train_iou} - Valid IOU: {valid_iou}. '\n            log_line += f'Lr: {current_lr}.'\n\n            log_and_print(logger, log_line)\n\n            metric_dict = {'train_loss':train_loss,'valid_loss':valid_loss,\n                           'train_IOU':train_iou, 'valid_IOU':valid_iou,\n                       }\n\n            progress_dict = log_to_progress_dict(progress_dict, metric_dict)\n\n            # plot figure and save the progress chart\n            save_progress(progress_dict, out_folder, out_folder_name, valid_fold, show=False)\n\n            if(valid_loss < best_valid_loss):\n                best_valid_loss = valid_loss\n                best_valid_ep = epoch\n                patient = PATIENT # reset patient\n\n                # save model\n                name = os.path.join(out_folder, 'Fold%d_%s.pth'%(valid_fold, \n                                                                 out_folder_name, \n                                                                ))\n                log_and_print(logger, 'Saving model to: ' + name)\n                torch.save(model.state_dict(), name)\n            else:\n                patient -= 1\n                log_and_print(logger, 'Decrease early-stopping patient by 1 due valid loss not decreasing. Patient='+ str(patient))\n\n            if(patient == 0):\n                log_and_print(logger, 'Early stopping patient = 0. Early stop')\n                break\n\n# ======================================================================\n\n\n","metadata":{"execution":{"iopub.status.busy":"2024-10-18T06:35:24.658798Z","iopub.execute_input":"2024-10-18T06:35:24.659166Z"},"trusted":true},"execution_count":null,"outputs":[]}]}