{"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 0. Imports and Constants","metadata":{}},{"cell_type":"code","source":"# install packages\n!pip3 install -U torch torchvision --extra-index-url https://download.pytorch.org/whl/cu116\n!pip3 install python-gdcm pylibjpeg pylibjpeg-libjpeg pydicom","metadata":{"execution":{"iopub.status.busy":"2022-10-29T13:43:14.561532Z","iopub.execute_input":"2022-10-29T13:43:14.561968Z","iopub.status.idle":"2022-10-29T13:48:02.922023Z","shell.execute_reply.started":"2022-10-29T13:43:14.561923Z","shell.execute_reply":"2022-10-29T13:48:02.920771Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# import libraries\nimport os\nimport gc\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nimport pydicom\nfrom pydicom.pixel_data_handlers.util import apply_voi_lut\nfrom sklearn.model_selection import GroupKFold\nimport torch\nimport torchvision as tv\nfrom torchvision.models.feature_extraction import create_feature_extractor\nfrom tqdm import tqdm\nimport time\nimport warnings\nwarnings.filterwarnings('ignore')","metadata":{"execution":{"iopub.status.busy":"2022-10-29T13:48:02.925492Z","iopub.execute_input":"2022-10-29T13:48:02.926254Z","iopub.status.idle":"2022-10-29T13:48:05.679846Z","shell.execute_reply.started":"2022-10-29T13:48:02.926217Z","shell.execute_reply":"2022-10-29T13:48:05.678865Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# specify paths and parameters\nPATH_MAIN = '../input/rsna-2022-spine-fracture-detection-metadata'  # to be changed\nPATH_SEGMENTATION = f'{PATH_MAIN}/meta_segmentation.csv'\nPATH_TRAIN_IMAGES = '../input/rsna-2022-cervical-spine-fracture-detection/train_images'\n\nNUM_VERTEBRAE = 7\nNUM_FOLDS = 5\nBATCH_SIZE = 32\nSEED = 0","metadata":{"execution":{"iopub.status.busy":"2022-10-29T13:48:05.681283Z","iopub.execute_input":"2022-10-29T13:48:05.682207Z","iopub.status.idle":"2022-10-29T13:48:05.689341Z","shell.execute_reply.started":"2022-10-29T13:48:05.682171Z","shell.execute_reply":"2022-10-29T13:48:05.687374Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# use PyTorch CUDA for computation\nDEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'\nprint(f'Device: {DEVICE}')","metadata":{"execution":{"iopub.status.busy":"2022-10-29T13:48:05.692245Z","iopub.execute_input":"2022-10-29T13:48:05.692612Z","iopub.status.idle":"2022-10-29T13:48:05.762395Z","shell.execute_reply.started":"2022-10-29T13:48:05.692577Z","shell.execute_reply":"2022-10-29T13:48:05.761512Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def show_cuda_memory():\n    t = torch.cuda.get_device_properties(0).total_memory\n    r = torch.cuda.memory_reserved(0)\n    a = torch.cuda.memory_allocated(0)\n    f = r-a  # free inside reserved\n    print('Total:     {:0.2f} GiB'.format(t / 2**30))\n    print('Reserved:  {:0.2f} GiB'.format(r / 2**30))\n    print('Allocated: {:0.2f} GiB'.format(a / 2**30))\n    print('Free:      {:0.2f} GiB'.format(f / 2**30))\n\nshow_cuda_memory()","metadata":{"execution":{"iopub.status.busy":"2022-10-29T13:48:05.763775Z","iopub.execute_input":"2022-10-29T13:48:05.764154Z","iopub.status.idle":"2022-10-29T13:48:05.784615Z","shell.execute_reply.started":"2022-10-29T13:48:05.764116Z","shell.execute_reply":"2022-10-29T13:48:05.783607Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Step 1. Load Data","metadata":{}},{"cell_type":"code","source":"def load_image(image_path):\n    '''\n    Load image data from given path. Return the image pixel array and the file metadata.\n    \n    Parameters:\n        image_path: str\n            Path to the dicom file of the target image.\n    \n    Return:\n        image: numpy.ndarray\n            The image pixel array with shape (H, W, C) in which C = 3 (RGB image).\n            Pixel values range from 0 to 255.\n        metadata: pydicom.dataset.FileDataset\n            The metadata of the dicom file.\n    '''\n    # load metadata and alter photometric interpretation of image pixels\n    metadata = pydicom.dcmread(image_path)\n    # extract image pixel array and rescale to [0, 255]\n    image = apply_voi_lut(metadata.pixel_array, metadata)\n    image = 255 * ((image - np.min(image)) / (np.max(image) - np.min(image)))\n    # change image dtype and convert to RGB\n    image = image.astype(np.uint8)\n    image = np.stack([image] * 3, axis=-1)\n    return image, metadata","metadata":{"execution":{"iopub.status.busy":"2022-10-29T13:48:05.786181Z","iopub.execute_input":"2022-10-29T13:48:05.786861Z","iopub.status.idle":"2022-10-29T13:48:05.793687Z","shell.execute_reply.started":"2022-10-29T13:48:05.78682Z","shell.execute_reply":"2022-10-29T13:48:05.792713Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# load an image\nimage, metadata = load_image(f'{PATH_TRAIN_IMAGES}/1.2.826.0.1.3680043.17625/150.dcm')\nprint(f'image shape: {image.shape}')\nprint(f'pixel range: [{np.min(image)}, {np.max(image)}]')\nplt.imshow(image)\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2022-10-29T13:48:05.795399Z","iopub.execute_input":"2022-10-29T13:48:05.796175Z","iopub.status.idle":"2022-10-29T13:48:06.075503Z","shell.execute_reply.started":"2022-10-29T13:48:05.796139Z","shell.execute_reply":"2022-10-29T13:48:06.074516Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# load vertebrae segmentation dataset\ndf_segmentation = pd.read_csv(PATH_SEGMENTATION)\ndf_segmentation.head()","metadata":{"execution":{"iopub.status.busy":"2022-10-29T13:48:06.077203Z","iopub.execute_input":"2022-10-29T13:48:06.077589Z","iopub.status.idle":"2022-10-29T13:48:06.177845Z","shell.execute_reply.started":"2022-10-29T13:48:06.077537Z","shell.execute_reply":"2022-10-29T13:48:06.176942Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# split the dataset into folds by Group K-Fold mechanism\n# ensure that the slices from the same patient do not appear in training and test set simultaneously\nnp.random.seed(SEED)\ngroup_kfold = GroupKFold(NUM_FOLDS)\nfolds = group_kfold.split(df_segmentation, groups=df_segmentation.StudyInstanceUID)\nfor fold, (_, test_indices) in enumerate(folds):\n    df_segmentation.loc[test_indices, 'Fold'] = fold\ndf_segmentation.iloc[:,-1] = df_segmentation.iloc[:,-1].astype(np.uint8)\nprint(f'fold indices: {sorted(df_segmentation.Fold.unique())}')\ndf_segmentation.sample(5)","metadata":{"execution":{"iopub.status.busy":"2022-10-29T13:48:06.179568Z","iopub.execute_input":"2022-10-29T13:48:06.18023Z","iopub.status.idle":"2022-10-29T13:48:06.230837Z","shell.execute_reply.started":"2022-10-29T13:48:06.180192Z","shell.execute_reply":"2022-10-29T13:48:06.229832Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# add another column to dataframe: SliceRatio\n# SliceRatio specifies the ratio of the current slice \n# with respect to the total number of slices of this patient\n# i.e., SliceRatio = current slice index / max slice index\n\n# create a dictionary which maps each StudyInstanceUID (patient)\n# to the total number of slices of this patient\nuid2max_slice = df_segmentation.groupby('StudyInstanceUID')['Slice'].max().to_dict()\nmax_slices = df_segmentation['StudyInstanceUID'].map(uid2max_slice)\ndf_segmentation['SliceRatio'] = df_segmentation['Slice'] / max_slices\ndf_segmentation.sample(5)","metadata":{"execution":{"iopub.status.busy":"2022-10-29T13:48:06.234339Z","iopub.execute_input":"2022-10-29T13:48:06.234674Z","iopub.status.idle":"2022-10-29T13:48:06.263892Z","shell.execute_reply.started":"2022-10-29T13:48:06.234629Z","shell.execute_reply":"2022-10-29T13:48:06.262796Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class ImageDataset(torch.utils.data.Dataset):\n    '''\n    An image dataset to extract slice images and their target labels\n    '''\n    def __init__(self, df, image_dir):\n        '''\n        Initialize the image dataset.\n        \n        Parameters:\n            df: pandas.core.frame.DataFrame\n                The dataframe containing segmentation information of each slice.\n            image_dir: str\n                The path to the image directory.\n        '''\n        super().__init__()\n        self.df = df\n        self.image_dir = image_dir\n    \n    def __len__(self):\n        '''\n        Length of the dataset\n        \n        Return:\n            length: int\n                Total number of slices in the dataset.\n        '''\n        length = len(self.df)\n        return length\n    \n    def __getitem__(self, idx):\n        '''\n        Retrieve the idx-th slice of the dataset and its corresponding vertebrae labels.\n        \n        Parameters:\n            idx: int\n                The index of the slice in the dataset to retrieve.\n        \n        Return:\n            image: torch.Tensor\n                The tensor of slice image pixels in the shape of [C, H, W].\n            labels: torch.Tensor\n                The tensor of target vertebrae labels of the image in the shape of [NUM_VERTEBRAE].\n                The i-th entry of this tensor is either 0 or 1, \n                indicating whether the vertebrae Ci is in this slice or not.\n                Return random numbers if df does not have vertebrae label columns.\n        '''\n        # get path to the slice image\n        patient_uid = self.df.iloc[idx].StudyInstanceUID\n        slice_number = self.df.iloc[idx].Slice\n        image_path = os.path.join(self.image_dir, patient_uid, f'{slice_number}.dcm')\n        # load the slice image\n        image, _ = load_image(image_path)\n        # convert (H, W, C) to (C, H, W) in the PyTorch convention\n        image = np.transpose(image, (2, 0, 1))\n        # transform the image\n        image = torch.as_tensor(image)\n        \n        # retrieve image labels if the dataframe has label columns\n        labels = torch.randn(NUM_VERTEBRAE)\n        if 'C1' in self.df.columns:\n            # get column names of the vertebrae labels to retrieve\n            cols = [f'C{i}' for i in range(1, NUM_VERTEBRAE+1)]\n            # get the labels\n            labels = self.df.iloc[idx][cols]\n            # transform the labels\n            labels = labels.astype(np.float64)\n            labels = torch.as_tensor(labels)\n        return image, labels","metadata":{"execution":{"iopub.status.busy":"2022-10-29T13:48:06.265587Z","iopub.execute_input":"2022-10-29T13:48:06.26605Z","iopub.status.idle":"2022-10-29T13:48:06.278564Z","shell.execute_reply.started":"2022-10-29T13:48:06.266009Z","shell.execute_reply":"2022-10-29T13:48:06.277308Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# check image and labels retrieved from dataset\nds_segmentation = ImageDataset(df_segmentation, PATH_TRAIN_IMAGES)\nimage, labels = ds_segmentation[612]\nC, H, W = image.shape\nprint(f'image shape: {image.shape}')\nprint(f'labels: {labels}')","metadata":{"execution":{"iopub.status.busy":"2022-10-29T13:48:06.280292Z","iopub.execute_input":"2022-10-29T13:48:06.281226Z","iopub.status.idle":"2022-10-29T13:48:06.316363Z","shell.execute_reply.started":"2022-10-29T13:48:06.28119Z","shell.execute_reply":"2022-10-29T13:48:06.315326Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Step 2. Model Construction","metadata":{}},{"cell_type":"code","source":"class Model(torch.nn.Module):\n    '''\n    A vertebrae detection model.\n    '''\n    def __init__(self, model, weights, feature_layer, feature_dim):\n        '''\n        Model initialization.\n        \n        Parameters:\n            model: function\n                The function to create the feature extraction model to retrieve image features.\n            weights: enum\n                Initial weights of the model.\n            feature_layer: str\n                Name of the layer to extract image features.\n            feature_dim: int\n                Dimension of image features.\n        '''\n        super().__init__()\n        # image transformation procedures\n        self.transform = weights.transforms()\n        # feature extraction layer\n        self.feature_layer = feature_layer\n        # get flatten layer output, i.e., feature vector of the slice image\n        self.model = create_feature_extractor(\n            model(weights=weights),\n            [feature_layer]\n        )\n        # fully connected layer to get the final prediction vector\n        self.fc = torch.nn.Linear(feature_dim, NUM_VERTEBRAE)\n\n    def forward(self, x):\n        '''\n        Predict the logit of each vertebra being present in the image.\n        \n        Parameters:\n            x: torch.Tensor\n                The input image tensor of shape [B, C, H, W].\n        \n        Return:\n            out: torch.Tensor\n                A tensor of shape [B, NUM_VERTEBRAE] indicating the vertebrae probability logits.\n        '''\n        x = self.transform(x)\n        x = self.model(x)[self.feature_layer]  # shape [B, feature_dim]\n        out = self.fc(x)  # shape [B, NUM_VERTEBRAE]\n        return out\n    \n    def predict(self, x):\n        '''\n        Predict the probabilities of each vertebra being present in the image.\n        \n        Parameters:\n            x: torch.Tensor\n                The input image tensor of shape [B, C, H, W].\n        \n        Return:\n            probabilities: torch.Tensor\n                A tensor of shape [B, NUM_VERTEBRAE] indicating the vertebrae probabilities.\n        '''\n        logits = self.forward(x)\n        probabilities = torch.nn.Sigmoid()(logits)\n        return probabilities","metadata":{"execution":{"iopub.status.busy":"2022-10-29T13:48:06.317873Z","iopub.execute_input":"2022-10-29T13:48:06.318223Z","iopub.status.idle":"2022-10-29T13:48:06.327196Z","shell.execute_reply.started":"2022-10-29T13:48:06.318187Z","shell.execute_reply":"2022-10-29T13:48:06.325976Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Step 3. Model Training and Evaluation","metadata":{}},{"cell_type":"code","source":"def evaluate(model, ds, autocast=True):\n    '''\n    Evaluate the model performance on the test set.\n    \n    Parameters:\n        model: Model\n            The model to test performance.\n        ds: ImageDataset\n            The test image dataset.\n        autocast: bool, default=True\n            Whether to apply tensor dtype autocast or not.\n    \n    Return:\n        accuracy: float\n            The prediction accuracy of the model on the given test dataset.\n        preds: numpy.ndarray\n            The predicted probability result of all test slices, \n            with shape [N, NUM_VERTEBRAE] where N is the number of test samples.\n    '''\n    # load the test dataset\n    test_dataloader = torch.utils.data.DataLoader(\n        ds,\n        batch_size=BATCH_SIZE,\n        shuffle=False,\n        num_workers=os.cpu_count()\n    )\n    # use GPU for acceleration\n    model = model.to(DEVICE)\n    \n    # disable gradient calculation\n    with torch.no_grad():\n        # lists of predicted and true vertebrae labels\n        preds, trues = [], []\n        # evaluation mode\n        model.eval()\n        # make predictions for each batch\n        # image shape: [B, C, H, W]\n        # labels shape: [B, NUM_VERTEBRAE]\n        for image, labels in tqdm(test_dataloader, desc='Evaluation'):\n            # make predictions\n            if autocast:\n                # cast tensors to a smaller memory footprint to save memory\n                with torch.cuda.amp.autocast():\n                    pred = model.predict(image.to(DEVICE))\n            else:\n                pred = model.predict(image.to(DEVICE))\n            # store the true and predicted labels to lists\n            preds.append(pred.cpu().numpy())\n            trues.append(labels.numpy())\n        \n        # concatenate predicted and true labels of all the batches\n        preds = np.concatenate(preds)\n        trues = np.concatenate(trues)\n        # calculate the overall prediction accuracy (threshold 0.5)\n        accuracy = np.mean((preds > 0.5) == trues)  # change to f1-score/AUC\n        print(f'Evaluation accuracy: {accuracy}')\n        \n    return accuracy, preds","metadata":{"execution":{"iopub.status.busy":"2022-10-29T13:48:06.328939Z","iopub.execute_input":"2022-10-29T13:48:06.329337Z","iopub.status.idle":"2022-10-29T13:48:06.340579Z","shell.execute_reply.started":"2022-10-29T13:48:06.329302Z","shell.execute_reply":"2022-10-29T13:48:06.339419Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def train(model, ds_train, ds_test=None, num_epochs=1, autocast=True):\n    '''\n    Train the model by the training set and evaluate its performance on the test set.\n    \n    Parameters:\n        model: Model\n            The model to train and test performance.\n        ds_train: ImageDataset\n            The training dataset.\n        ds_test: ImageDataset, default=None\n            The test dataset to evaluate the model performance.\n            If not specified, do not conduct model evaluation.\n        num_epochs: int, default=1\n            Number of training epochs.\n        autocast: bool, default=True\n            Whether to apply tensor dtype autocast or not.\n            \n    Return:\n        test_accuracy: float\n            The prediction accuracy of the model on the test dataset.\n            Return None if test set not specified.\n        preds: numpy.ndarray\n            The predicted probability result of all test slices, \n            with shape [N, NUM_VERTEBRAE] where N is the number of test samples.\n            Return None if test set not specified.\n    '''\n    # load the training dataset\n    train_dataloader = torch.utils.data.DataLoader(\n        ds_train,\n        batch_size=BATCH_SIZE,\n        shuffle=True,\n        num_workers=os.cpu_count()\n    )\n    # prepare model, optimizer, and scheduler\n    model = model.to(DEVICE)\n    optimizer = torch.optim.Adam(model.parameters())\n    scheduler = torch.optim.lr_scheduler.OneCycleLR(\n        optimizer, \n        max_lr=1e-3, \n        epochs=num_epochs, \n        steps_per_epoch=len(train_dataloader)\n    )\n    \n    # train mode\n    model.train()\n    # use gradient scaler to avoid underflowing gradients\n    scaler = torch.cuda.amp.GradScaler()\n    \n    for _ in range(num_epochs):\n        # update the weights with each batch\n        # image shape: [B, C, H, W]\n        # labels shape: [B, NUM_VERTEBRAE]\n        for image, labels in tqdm(train_dataloader, desc='Train'):\n            # start from zero gradients\n            optimizer.zero_grad()\n            # make predictions and get logits\n            # compute the binary cross entropy loss between target and predicted logits\n            if autocast:\n                # cast tensors to a smaller memory footprint to save memory\n                with torch.cuda.amp.autocast():\n                    pred = model.forward(image.to(DEVICE))\n                    true = labels.to(DEVICE)\n                    loss = torch.nn.functional.binary_cross_entropy_with_logits(pred, true)\n            else:\n                pred = model.forward(image.to(DEVICE))\n                true = labels.to(DEVICE)\n                loss = torch.nn.functional.binary_cross_entropy_with_logits(pred, true)\n            # update weights by backward propagation\n            scaler.scale(loss).backward()\n            scaler.step(optimizer)\n            scaler.update()\n            # update the scheduler\n            scheduler.step()\n\n        # evaluate the trained model, get test accuracy and prediction\n        test_accuracy, preds = None, None\n        if ds_test:\n            test_accuracy, preds = evaluate(model, ds_test, autocast)\n        \n    return test_accuracy, preds","metadata":{"execution":{"iopub.status.busy":"2022-10-29T13:48:06.342463Z","iopub.execute_input":"2022-10-29T13:48:06.343145Z","iopub.status.idle":"2022-10-29T13:48:06.35648Z","shell.execute_reply.started":"2022-10-29T13:48:06.343109Z","shell.execute_reply":"2022-10-29T13:48:06.355325Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# A dictionary that stores all the model templates.\n# Key: model name.\n# Value: a tuple of (model function, model initial weights, feature extraction layer, extracted feature dimension)\nmodels = {\n    'AlexNet':            (tv.models.alexnet,            tv.models.AlexNet_Weights.DEFAULT,                 'flatten',      9216),\n    'VGG16_BN':           (tv.models.vgg16_bn,           tv.models.VGG16_BN_Weights.DEFAULT,                'flatten',      25088),\n    'ResNet50':           (tv.models.resnet50,           tv.models.ResNet50_Weights.DEFAULT,                'flatten',      2048),\n    'ResNeXt50_32X4D':    (tv.models.resnext50_32x4d,    tv.models.ResNeXt50_32X4D_Weights.DEFAULT,         'flatten',      2048),\n    'GoogLeNet':          (tv.models.googlenet,          tv.models.GoogLeNet_Weights.DEFAULT,               'flatten',      1024),\n    'Inception_V3':       (tv.models.inception_v3,       tv.models.Inception_V3_Weights.DEFAULT,            'flatten',      2048),\n    'DenseNet121':        (tv.models.densenet121,        tv.models.DenseNet121_Weights.DEFAULT,             'flatten',      1024),\n    'EfficientNet_V2_S':  (tv.models.efficientnet_v2_s,  tv.models.EfficientNet_V2_S_Weights.DEFAULT,       'flatten',      1280),\n    'MobileNet_V3_Large': (tv.models.mobilenet_v3_large, tv.models.MobileNet_V3_Large_Weights.DEFAULT,      'flatten',      960),\n    'ConvNeXt_Small':     (tv.models.convnext_small,     tv.models.ConvNeXt_Small_Weights.DEFAULT,          'classifier.1', 768),\n    'MNASNet1_3':         (tv.models.mnasnet1_3,         tv.models.MNASNet1_3_Weights.DEFAULT,              'classifier.0', 1280),\n    'ShuffleNet_V2_X2_0': (tv.models.shufflenet_v2_x2_0, tv.models.ShuffleNet_V2_X2_0_Weights.DEFAULT,      'mean',         2048),\n    'Swin_S':             (tv.models.swin_s,             tv.models.Swin_S_Weights.DEFAULT,                  'flatten',      768),\n    'RegNet_Y_8GF':       (tv.models.regnet_y_8gf,       tv.models.RegNet_Y_8GF_Weights.DEFAULT,            'flatten',      2016),\n    # 'ViT_B_16':           (tv.models.vit_b_16,           tv.models.ViT_B_16_Weights.IMAGENET1K_SWAG_E2E_V1, 'getitem_5',    768)\n}\n\ndef create_model(model_name):\n    '''\n    Create a model instance by the specified model name.\n    \n    Parameters:\n        model_name: str\n            Name of the model.\n            Has to be in the model list.\n            \n    Return:\n        model: Model\n            An initialized model instance.\n    '''\n    model = Model(*models[model_name])\n    return model","metadata":{"execution":{"iopub.status.busy":"2022-10-29T13:48:06.358233Z","iopub.execute_input":"2022-10-29T13:48:06.358672Z","iopub.status.idle":"2022-10-29T13:48:06.371234Z","shell.execute_reply.started":"2022-10-29T13:48:06.358633Z","shell.execute_reply":"2022-10-29T13:48:06.369913Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# model selection\ntest_accs = dict()  # model evaluation accuracy\ntrain_time = dict()  # model training time\n\n# load the training set: train on folds 1 ~ NUM_FOLDS-1\nds_train = ImageDataset(df_segmentation[df_segmentation.Fold != 0], PATH_TRAIN_IMAGES)\n# load the test set: test on fold 0\nds_test = ImageDataset(df_segmentation[df_segmentation.Fold == 0], PATH_TRAIN_IMAGES)\n\n# for each model in the dictionary\nfor model_name in models:\n    print(f'Model: {model_name}')\n    # initialize the model\n    model = create_model(model_name)\n    # train and evaluate the model; record training time\n    start = time.time()\n    test_acc, _ = train(model, ds_train, ds_test)\n    end = time.time()\n    time_elapsed = end - start\n    # store evaluation accuracy and training time\n    test_accs[model_name] = test_acc\n    train_time[model_name] = time_elapsed\n    # clean cache\n    del model\n    gc.collect()\n    torch.cuda.empty_cache()\n\n# clean cache\ndel ds_train, ds_test\ngc.collect()\ntorch.cuda.empty_cache()","metadata":{"execution":{"iopub.status.busy":"2022-10-27T18:40:54.456892Z","iopub.execute_input":"2022-10-27T18:40:54.45728Z","iopub.status.idle":"2022-10-27T20:28:08.946178Z","shell.execute_reply.started":"2022-10-27T18:40:54.457241Z","shell.execute_reply":"2022-10-27T20:28:08.945061Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# model selection\ncv_accs = dict()  # model cross validation accuracy\n    \n# for each model with test accuracy in the previous stage > 0.96\nfor model_name in ['RegNet_Y_8GF', 'DenseNet121', 'EfficientNet_V2_S', 'ResNeXt50_32X4D', 'Inception_V3']:\n    print(f'Model: {model_name}')\n    # perform K-fold cross validation\n    test_accs = []\n    for fold in range(NUM_FOLDS):\n        # initialize the model\n        model = create_model(model_name)\n        # load the training set\n        ds_train = ImageDataset(df_segmentation[df_segmentation.Fold != fold], PATH_TRAIN_IMAGES)\n        # load the test set\n        ds_test = ImageDataset(df_segmentation[df_segmentation.Fold == fold], PATH_TRAIN_IMAGES)\n        # train and evaluate the model\n        test_acc, _ = train(model, ds_train, ds_test)\n        test_accs.append(test_acc)\n        # clean cache\n        del model, ds_train, ds_test\n        gc.collect()\n        torch.cuda.empty_cache()\n    cv_acc = sum(test_accs) / NUM_FOLDS\n    print(f'Cross validation accuracy: {cv_acc}')\n    cv_accs[model_name] = cv_acc","metadata":{"execution":{"iopub.status.busy":"2022-10-29T13:53:30.637116Z","iopub.execute_input":"2022-10-29T13:53:30.637503Z","iopub.status.idle":"2022-10-29T17:06:41.215928Z","shell.execute_reply.started":"2022-10-29T13:53:30.637468Z","shell.execute_reply":"2022-10-29T17:06:41.213389Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Step 4. Inference","metadata":{}},{"cell_type":"code","source":"# initialize the model\nweights = tv.models.EfficientNet_V2_S_Weights.DEFAULT\nmodel = tv.models.efficientnet_v2_s(weights=weights)\nmodel = Model(model, transform=weights.transforms())\n# load the training set (the entire segmentation dataset)\nds_train = ImageDataset(df_segmentation, PATH_TRAIN_IMAGES)\n# train the model by the entire segmentation dataset\ntrain(model, ds_train)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def infer(model, df, image_dir):\n    '''\n    Infer the vertebrae probabilities for all the slice samples in the given dataframe.\n    \n    Parameters:\n        model: Model\n            The model used for inference.\n        df: pandas.core.frame.DataFrame\n            The dataframe containing patient and slice indices to infer vertebrae.\n        image_dir: str\n            The path to the slice image directory.\n    \n    Return:\n        preds: numpy.ndarray\n            The predicted probability result of all test slices, \n            with shape [N, NUM_VERTEBRAE] where N is the total number of slices in df.\n    '''\n    ds = ImageDataset(df, PATH_TRAIN_IMAGES)\n    dataloader = torch.utils.data.DataLoader(\n        ds,\n        batch_size=BATCH_SIZE,\n        shuffle=False,\n        num_workers=os.cpu_count()\n    )\n    # the final prediction results\n    preds = []\n    # use GPU for acceleration\n    model = model.to(DEVICE)\n    # disable gradient calculation\n    with torch.no_grad():\n        # evaluation mode\n        model.eval()\n        # make predictions for each image\n        # image shape: [B, C, H, W]\n        for image, _ in tqdm(dataloader, desc='Inference'):\n            # cast tensors to a smaller memory footprint to save memory\n            if True:#with torch.cuda.amp.autocast():\n                # make predictions\n                pred = model.predict(image.to(DEVICE))\n            # store the true and predicted labels to lists\n            preds.append(pred.cpu().numpy())\n        # concatenate predicted and true labels of all the batches\n        preds = np.concatenate(preds)\n    return preds","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# make inference for all the training sample slices\n# to check the inference function is working properly\npreds = infer(model, df_segmentation, image_dir=PATH_TRAIN_IMAGES)\n# calculate all the predicted vertebrae probabilities\npred_cols = [f'C{i}_pred' for i in range(1, NUM_VERTEBRAE+1)]\ndf_segmentation[pred_cols] = preds\ndf_segmentation.head()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# check the accuracy on the training set \n# to ensure the inference function is working properly\ncols = [f'C{i}' for i in range(1, NUM_VERTEBRAE+1)]\ntrues = df_segmentation[cols].values\npreds = df_segmentation[pred_cols].values\naccuracy = np.mean((preds > 0.5) == trues)\nprint(f'Training set accuracy: {accuracy}')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# infer the vertebrae probability for all the samples\ndf_all = pd.read_csv(os.path.join(PATH_MAIN, 'meta_train_clean.csv'))\ndf_all = df_all.head(100)  # extract 100 samples for testing, delete this line in the future\ndf_all.head()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# make inference for all the sample slices\npreds = infer(model, df_all, image_dir=PATH_TRAIN_IMAGES)\n# save all the predicted vertebrae probabilities\ndf_all[cols] = preds\ndf_all.to_csv('train_segmented.csv', index=False)\ndf_all.head()","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}