{"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":"## Vesuvius Challenge - Ink Detection Training Notebook","metadata":{}},{"cell_type":"markdown","source":"Summary:\n- Model training uses pre-trained weights \n- Training on fragment 2 & 3, validation on fragment 1","metadata":{}},{"cell_type":"markdown","source":"### Setup","metadata":{}},{"cell_type":"code","source":"%%capture\n!pip install segmentation_models_pytorch\n\n# Pretrained weights\n# ref - https://github.com/kenshohara/3D-ResNets-PyTorch\n!pip install gdown\n!gdown 1Nb4abvIkkp_ydPFA9sNPT1WakoVKA8Fa\n\n# Utility packages for reading and visualizing volumes\n!pip install zarr imageio-ffmpeg\n\n# save model checkpoints\n!mkdir ./ckpts","metadata":{"execution":{"iopub.status.busy":"2023-06-27T13:16:49.713604Z","iopub.execute_input":"2023-06-27T13:16:49.714050Z","iopub.status.idle":"2023-06-27T13:17:58.820261Z","shell.execute_reply.started":"2023-06-27T13:16:49.714003Z","shell.execute_reply":"2023-06-27T13:17:58.818797Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import glob\nimport os\nimport gc\nimport sys\nimport zarr\nimport random\nimport imageio\nimport numpy as np\nimport pandas as pd\nfrom tqdm import tqdm\nimport matplotlib.pyplot as plt\nfrom IPython.display import Video\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.cuda import amp\nfrom torch.utils.data import Dataset, DataLoader\n\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\nfrom albumentations import ImageOnlyTransform\n\nimport segmentation_models_pytorch as smp\n\nsys.path.append(\"/kaggle/input/resnet3d\")\nfrom resnet3d import generate_model\n\nnp.random.seed(42)","metadata":{"_kg_hide-input":false,"_kg_hide-output":true,"execution":{"iopub.status.busy":"2023-06-27T13:17:58.822658Z","iopub.execute_input":"2023-06-27T13:17:58.823898Z","iopub.status.idle":"2023-06-27T13:18:04.551120Z","shell.execute_reply.started":"2023-06-27T13:17:58.823847Z","shell.execute_reply":"2023-06-27T13:18:04.549736Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Config","metadata":{}},{"cell_type":"code","source":"TRAIN_FRAGMENTS = [\"2\", \"3\"]\nTEST_FRAGMENT = \"1\"\n    \nclass ModelConfig:\n    # model\n    crop_size_scaling=2\n    crop_size = 256*crop_size_scaling\n    z_start = 24\n    z_dims = 16\n    \n    # training\n    init_lr = 1e-4\n    batch_size = int(32/(crop_size_scaling**2))\n    epochs = 50\n    \n    # augmentation\n    train_aug_list = [\n        #A.PadIfNeeded(min_height=crop_size, min_width=crop_size, position=\"top_left\"),\n        A.ToFloat(),\n        A.HorizontalFlip(p=0.5),\n        A.VerticalFlip(p=0.5),\n        A.ShiftScaleRotate(p=0.75),\n        A.CoarseDropout(max_holes=1, max_width=int(crop_size * 0.1), max_height=int(crop_size * 0.1)),\n        A.ChannelDropout(p=0.2),\n        ToTensorV2()\n    ]\n\n    valid_aug_list = [\n        #A.PadIfNeeded(min_height=crop_size, min_width=crop_size, position=\"top_left\"),\n        A.ToFloat(),\n        ToTensorV2(),\n    ]\n    \n    #inference\n    tta = None # rotate, flip\n\ndef get_transforms(aug_list):\naug = A.Compose(aug_list)\nreturn aug","metadata":{"execution":{"iopub.status.busy":"2023-06-27T13:18:04.553174Z","iopub.execute_input":"2023-06-27T13:18:04.553539Z","iopub.status.idle":"2023-06-27T13:18:04.564433Z","shell.execute_reply.started":"2023-06-27T13:18:04.553500Z","shell.execute_reply":"2023-06-27T13:18:04.562896Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Load data","metadata":{}},{"cell_type":"code","source":"FRAGMENTS_ZARR = {\n    \"1\" : zarr.open(\"/kaggle/input/vesuvius-zarr-files/train-1.zarr\", mode=\"r\"),\n    \"2\" : zarr.open(\"/kaggle/input/vesuvius-zarr-files/train-2.zarr\", mode=\"r\"),\n    \"3\" : zarr.open(\"/kaggle/input/vesuvius-zarr-files/train-3.zarr\", mode=\"r\")\n}\n\nFRAGMENTS_SHAPE = {k : v.mask.shape for k, v in FRAGMENTS_ZARR.items()}","metadata":{"execution":{"iopub.status.busy":"2023-06-26T14:04:35.342429Z","iopub.execute_input":"2023-06-26T14:04:35.342786Z","iopub.status.idle":"2023-06-26T14:04:35.664415Z","shell.execute_reply.started":"2023-06-26T14:04:35.342726Z","shell.execute_reply":"2023-06-26T14:04:35.663001Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Dataloaders","metadata":{}},{"cell_type":"code","source":"class VesuviusTrain(Dataset):\n    def __init__(self, fragments, cfg):\n        self.fragments = fragments\n        self.xys = []\n        self.transform = get_transforms(cfg.train_aug_list)\n        self.cfg = cfg\n        \n        for fragment in fragments:\n            H, W = FRAGMENTS_SHAPE[fragment]\n            for y in range(0, H-cfg.crop_size+1, cfg.crop_size):\n                for x in range(0, W-cfg.crop_size+1, cfg.crop_size):\n                    self.xys.append((fragment, x, y, W, H))\n        \n    def __getitem__(self, i):\n        fragment, x1, y1, W, H = self.xys[i]\n        z1, z2 = self.cfg.z_start, self.cfg.z_start+self.cfg.z_dims\n        \n        x_offset = random.randint(-32 if x1 != 0 else 0, 32)\n        y_offset = random.randint(-32 if y1 != 0 else 0, 32)\n        \n        x1 += x_offset\n        y1 += y_offset\n        \n        x2 = x1 + self.cfg.crop_size\n        y2 = y1 + self.cfg.crop_size\n        \n        if x2 > W:\n            x1 -= x_offset\n            x2 -= x_offset\n            \n        if y2 > H:\n            y1 -= y_offset\n            y2 -= y_offset\n        \n        frag_crop = FRAGMENTS_ZARR[fragment].surface_volume[y1:y2, x1:x2, z1:z2]\n        mask_crop = FRAGMENTS_ZARR[fragment].truth[y1:y2, x1:x2]\n        if self.transform is not None:\n            data = self.transform(image=frag_crop.astype(np.float32), mask=mask_crop.astype(np.float32))\n            image = data['image']\n            label = data['mask']\n            image = image/65535.0\n            image = (image - 0.45)/0.225\n            image, label = torch.unsqueeze(image, 0), torch.unsqueeze(label, 0)\n            return image, label\n        else:\n            if random.random() > 0.5:\n                frag_crop = np.flip(frag_crop, axis=1).copy()\n                mask_crop = np.flip(mask_crop, axis=1).copy()\n\n            frag_crop = torch.from_numpy(frag_crop.astype(np.float32)).unsqueeze(0).permute(0, 3, 1, 2)\n            frag_crop = frag_crop/65535.0\n            frag_crop = (frag_crop - 0.45)/0.225\n\n            mask_crop = torch.from_numpy(mask_crop.astype(np.float32)).unsqueeze(0)\n            return frag_crop, mask_crop\n\n    def __len__(self):\n        return len(self.xys)","metadata":{"execution":{"iopub.status.busy":"2023-06-26T14:04:35.669372Z","iopub.execute_input":"2023-06-26T14:04:35.669723Z","iopub.status.idle":"2023-06-26T14:04:35.690981Z","shell.execute_reply.started":"2023-06-26T14:04:35.669687Z","shell.execute_reply":"2023-06-26T14:04:35.689769Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class VesuviusVal(Dataset):\n    def __init__(self, fragment, cfg):\n        self.fragment = FRAGMENTS_ZARR[fragment]\n        self.xys = []\n        self.transform = None #get_transforms(cfg.valid_aug_list)\n        self.cfg=cfg\n        \n        H, W = FRAGMENTS_SHAPE[fragment]\n        for y in range(0, H-cfg.crop_size+1, cfg.crop_size):\n            for x in range(0, W-cfg.crop_size+1, cfg.crop_size):\n                self.xys.append((x, y))\n                \n    def __getitem__(self, i):\n        x1, y1 = self.xys[i]\n        x2, y2 = x1+self.cfg.crop_size, y1+self.cfg.crop_size\n        z1, z2 = self.cfg.z_start, self.cfg.z_start+self.cfg.z_dims\n        \n        frag_crop = self.fragment.surface_volume[y1:y2, x1:x2, z1:z2]\n        mask_crop = self.fragment.truth[y1:y2, x1:x2]\n        \n        if self.transform is not None:\n            data = self.transform(image=frag_crop, mask=mask_crop.astype(np.float32))\n            image = data['image']\n            label = data['mask']\n            image = image/65535.0\n            image = (image - 0.45)/0.225\n            image, label = torch.unsqueeze(image, 0), torch.unsqueeze(label, 0)\n            return image, label, torch.tensor([x1, y1, x2, y2], dtype=torch.int32)\n        else:\n            frag_crop = torch.from_numpy(frag_crop.astype(np.float32)).unsqueeze(0).permute(0, 3, 1, 2)\n            frag_crop = frag_crop/65535.0\n            frag_crop = (frag_crop - 0.45)/0.225\n\n            mask_crop = torch.from_numpy(mask_crop.astype(np.float32)).unsqueeze(0)\n            return frag_crop, mask_crop, torch.tensor([x1, y1, x2, y2], dtype=torch.int32)\n\n    def __len__(self):\n        return len(self.xys)","metadata":{"execution":{"iopub.status.busy":"2023-06-26T14:04:35.692692Z","iopub.execute_input":"2023-06-26T14:04:35.693393Z","iopub.status.idle":"2023-06-26T14:04:35.711081Z","shell.execute_reply.started":"2023-06-26T14:04:35.693348Z","shell.execute_reply":"2023-06-26T14:04:35.709909Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Create data loaders\ndataset_train = VesuviusTrain(TRAIN_FRAGMENTS, ModelConfig)\ndataloader_train = DataLoader(dataset_train, batch_size=ModelConfig.batch_size, num_workers=2,\n                              shuffle=True, pin_memory=True, drop_last=True)\nn_train = len(dataloader_train)\n\ndataset_valid = VesuviusVal(TEST_FRAGMENT, ModelConfig)\ndataloader_valid = DataLoader(dataset_valid, batch_size=ModelConfig.batch_size, num_workers=2,\n                              shuffle=False, pin_memory=True, drop_last=False)\nn_valid = len(dataloader_valid)","metadata":{"execution":{"iopub.status.busy":"2023-06-26T14:04:35.713172Z","iopub.execute_input":"2023-06-26T14:04:35.713642Z","iopub.status.idle":"2023-06-26T14:04:35.725034Z","shell.execute_reply.started":"2023-06-26T14:04:35.713598Z","shell.execute_reply":"2023-06-26T14:04:35.723813Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Model\n* Encoder is a 3D ResNet model with 18 layers. The architecture has been modified to remove temporal downsampling between blocks.\n* A 2D decoder is used for predicting the segmentation map.\n* The encoder feature maps are average pooled over the Z dimension before passing it to the decoder.","metadata":{}},{"cell_type":"code","source":"class Decoder(nn.Module):\n    def __init__(self, encoder_dims, upscale):\n        super().__init__()\n        self.convs = nn.ModuleList([\n            nn.Sequential(\n                nn.Conv2d(encoder_dims[i]+encoder_dims[i-1], encoder_dims[i-1], 3, 1, 1, bias=False),\n                nn.BatchNorm2d(encoder_dims[i-1]),\n                nn.ReLU(inplace=True)\n            ) for i in range(1, len(encoder_dims))])\n\n        self.logit = nn.Conv2d(encoder_dims[0], 1, 1, 1, 0)\n        self.up = nn.Upsample(scale_factor=upscale, mode=\"bilinear\")\n\n    def forward(self, feature_maps):\n        for i in range(len(feature_maps)-1, 0, -1):\n            f_up = F.interpolate(feature_maps[i], scale_factor=2, mode=\"bilinear\")\n            f = torch.cat([feature_maps[i-1], f_up], dim=1)\n            f_down = self.convs[i-1](f)\n            feature_maps[i-1] = f_down\n\n        x = self.logit(feature_maps[0])\n        mask = self.up(x)\n        return mask\n\n\nclass SegModel(nn.Module):\n    def __init__(self):\n        super().__init__()\n        self.encoder = generate_model(model_depth=18, n_input_channels=1)\n        self.decoder = Decoder(encoder_dims=[64, 128, 256, 512], upscale=4)\n        \n    def forward(self, x):\n        feat_maps = self.encoder(x)\n        feat_maps_pooled = [torch.mean(f, dim=2) for f in feat_maps]\n        pred_mask = self.decoder(feat_maps_pooled)\n        return pred_mask\n    \n    def load_pretrained_weights(self, state_dict):\n        # Convert 3 channel weights to single channel\n        # ref - https://timm.fast.ai/models#Case-1:-When-the-number-of-input-channels-is-1\n        conv1_weight = state_dict['conv1.weight']\n        state_dict['conv1.weight'] = conv1_weight.sum(dim=1, keepdim=True)\n        print(self.encoder.load_state_dict(state_dict, strict=False))","metadata":{"execution":{"iopub.status.busy":"2023-06-27T13:18:04.567800Z","iopub.execute_input":"2023-06-27T13:18:04.568297Z","iopub.status.idle":"2023-06-27T13:18:04.583921Z","shell.execute_reply.started":"2023-06-27T13:18:04.568245Z","shell.execute_reply":"2023-06-27T13:18:04.582502Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model = SegModel()\nmodel.load_pretrained_weights(torch.load(\"r3d18_K_200ep.pth\")[\"state_dict\"])\nmodel = nn.DataParallel(model, device_ids=[0, 1])\nmodel = model.cuda()","metadata":{"execution":{"iopub.status.busy":"2023-06-27T13:18:19.619432Z","iopub.execute_input":"2023-06-27T13:18:19.620600Z","iopub.status.idle":"2023-06-27T13:18:20.474427Z","shell.execute_reply.started":"2023-06-27T13:18:19.620552Z","shell.execute_reply":"2023-06-27T13:18:20.472936Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Competition metric (F0.5 Score)","metadata":{}},{"cell_type":"code","source":"# ref - https://www.kaggle.com/competitions/vesuvius-challenge-ink-detection/discussion/397288\ndef fbeta_score(preds, targets, threshold, beta=0.5, smooth=1e-5):\n    preds_t = torch.where(preds > threshold, 1.0, 0.0).float()\n    y_true_count = targets.sum()\n    \n    ctp = preds_t[targets==1].sum()\n    cfp = preds_t[targets==0].sum()\n    beta_squared = beta * beta\n\n    c_precision = ctp / (ctp + cfp + smooth)\n    c_recall = ctp / (y_true_count + smooth)\n    dice = (1 + beta_squared) * (c_precision * c_recall) / (beta_squared * c_precision + c_recall + smooth)\n\n    return dice","metadata":{"execution":{"iopub.status.busy":"2023-06-26T14:04:39.906661Z","iopub.execute_input":"2023-06-26T14:04:39.907495Z","iopub.status.idle":"2023-06-26T14:04:39.917030Z","shell.execute_reply.started":"2023-06-26T14:04:39.907449Z","shell.execute_reply":"2023-06-26T14:04:39.915937Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Training","metadata":{}},{"cell_type":"code","source":"# Define loss, optimize and scheduler\nDiceLoss = smp.losses.DiceLoss(mode='binary')\nBCELoss = smp.losses.SoftBCEWithLogitsLoss()\n\nalpha = 0.5\nbeta = 1 - alpha\nTverskyLoss = smp.losses.TverskyLoss(\n    mode='binary', log_loss=False, alpha=alpha, beta=beta)\n\ndef criterion(y_pred, y_true):\n    #return BCELoss(y_pred, y_true)\n    return 0.5 * BCELoss(y_pred, y_true) + 0.5 * DiceLoss(y_pred, y_true)\n    #return DiceLoss(y_pred, y_true)\n\nscaler = amp.GradScaler()\noptimizer = torch.optim.AdamW(model.parameters(), lr=ModelConfig.init_lr)\nscheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer, max_lr=ModelConfig.init_lr,\n                                                steps_per_epoch=10, epochs=ModelConfig.epochs//10,\n                                                pct_start=0.1)","metadata":{"execution":{"iopub.status.busy":"2023-06-26T14:04:39.918844Z","iopub.execute_input":"2023-06-26T14:04:39.919244Z","iopub.status.idle":"2023-06-26T14:04:39.935002Z","shell.execute_reply.started":"2023-06-26T14:04:39.919203Z","shell.execute_reply":"2023-06-26T14:04:39.933965Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#transform = get_transforms(ModelConfig.valid_aug_list)\n#mask = np.array(FRAGMENTS_ZARR[TEST_FRAGMENT].truth.astype(np.float32))\n#data = transform(image = np.zeros_like(mask), mask=mask)\n#gt_mask = data['mask']\n\ngt_mask = torch.from_numpy(np.asarray(FRAGMENTS_ZARR[TEST_FRAGMENT].truth)).float().cuda()\ngt_shape = FRAGMENTS_SHAPE[TEST_FRAGMENT]","metadata":{"execution":{"iopub.status.busy":"2023-06-26T14:04:39.936288Z","iopub.execute_input":"2023-06-26T14:04:39.937149Z","iopub.status.idle":"2023-06-26T14:04:40.544774Z","shell.execute_reply.started":"2023-06-26T14:04:39.937100Z","shell.execute_reply":"2023-06-26T14:04:40.543455Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def validate(model, dataloader_valid):\n    mloss_val = 0.0\n    model.eval()\n    pbar_val = enumerate(dataloader_valid)\n    pbar_val = tqdm(pbar_val, total=n_valid, bar_format=\"{l_bar}{bar:10}{r_bar}{bar:-10b}\")\n    final_pred_mask = torch.zeros(gt_shape, dtype=torch.float32, device='cuda')\n\n    for i, (fragments, masks, xys) in pbar_val:\n        fragments, masks = fragments.cuda(), masks.cuda()\n        with torch.no_grad():\n            if tta:\n                preds = []\n                for fragment in fragments:\n                    x = fragment\n                    if tta == \"flip\":\n                        x = [x, torch.flip(x, (-1,)), torch.flip(x, (-2,), torch.flip(x, (-2,-1)))]\n                    elif tta == \"rotate\":\n                        x=[torch.rot90(x, k=i, dims=(-2,-1)) for i in range(4)]\n                    x=torch.stack(x,dim=0)\n                    x=model(x)\n                    x=torch.sigmoid(x)\n                    if tta == \"flip\":\n                        x = [x[0], torch.flip(x[1], (-1,)), torch.flip(x[2], (-2,)), torch.flip(x[2], (-2,-1))]\n                    elif tta == \"rotate\":\n                        x=[torch.rot90(x[i],k=-i,dims=(-2,-1)) for i in range(4)]\n                    x=torch.stack(x,dim=0)\n                    preds.append(x.mean(0))\n                pred_masks = torch.stack(preds,dim=0)\n            else:\n                pred_masks = model(fragments)\n                mloss_val += criterion(pred_masks, masks).item()\n                pred_masks = torch.sigmoid(pred_masks)\n\n        for j, xy in enumerate(xys):\n            final_pred_mask[xy[1]:xy[3], xy[0]:xy[2]] = pred_masks[j, 0]\n\n        pbar_val.set_description((\"%10s\") % (f\"Val Loss: {mloss_val / (i+1):.4f}\"))\n    return final_pred_mask, mloss_val","metadata":{"execution":{"iopub.status.busy":"2023-06-26T14:04:40.546505Z","iopub.execute_input":"2023-06-26T14:04:40.546911Z","iopub.status.idle":"2023-06-26T14:04:40.557669Z","shell.execute_reply.started":"2023-06-26T14:04:40.546868Z","shell.execute_reply":"2023-06-26T14:04:40.556498Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"best_fbeta = 0\ntrain_losses = []\nval_losses = []\nfbetas = []\n\nfor epoch in range(1, ModelConfig.epochs+1):\n    model.train()\n    cur_lr = f\"LR : {scheduler.get_last_lr()[0]:.2E}\"\n    pbar_train = enumerate(dataloader_train)\n    pbar_train = tqdm(pbar_train, total=n_train, bar_format=\"{l_bar}{bar:10}{r_bar}{bar:-10b}\")\n    mloss_train = 0.0\n\n    for i, (fragments, masks) in pbar_train:\n        fragments, masks = fragments.cuda(), masks.cuda()\n        optimizer.zero_grad()\n        with amp.autocast():\n            pred_masks = model(fragments)\n            loss = criterion(pred_masks, masks)\n            scaler.scale(loss).backward()\n            scaler.step(optimizer)\n            scaler.update()\n            mloss_train += loss.detach().item()\n\n        gpu_mem = f\"Mem : {torch.cuda.memory_reserved() / 1E9:.3g}GB\"\n        pbar_train.set_description((\"%10s  \" * 3 + \"%10s\") % (f\"Epoch {epoch}/{ModelConfig.epochs}\", gpu_mem, cur_lr,\n                                                              f\"Loss: {mloss_train / (i + 1):.4f}\"))\n        \n    scheduler.step()\n    final_pred_mask, mloss_val = validate(model, dataloader_valid)\n    plt.imshow(final_pred_mask.cpu().numpy())\n    plt.show()\n    \n    fbeta_ = 0\n    for threshold in np.arange(0.2, 0.85, 0.05):\n        fbeta = fbeta_score(final_pred_mask, gt_mask, threshold)\n        fbeta_ = max(fbeta, fbeta_)\n        print(f\"Threshold : {threshold:.2f}\\tFBeta : {fbeta:.6f}\")\n        \n    # save losses and metrics\n    train_losses.append(mloss_train)\n    val_losses.append(mloss_val)\n    fbetas.append(fbeta_.item())\n    \n    if fbeta_ > best_fbeta:\n        best_fbeta = fbeta_\n        torch.save(model.module.state_dict(), f\"./ckpts/resnet18_3d_seg_{epoch}_{best_fbeta:.2f}.pt\")","metadata":{"execution":{"iopub.status.busy":"2023-06-26T14:04:40.563044Z","iopub.execute_input":"2023-06-26T14:04:40.564117Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Plot training curve\nplt.plot(np.array(train_losses)/len(dataloader_train), color=\"blue\", label=\"train_loss\")\nplt.plot(np.array(val_losses)/len(dataloader_valid), color=\"orange\", label=\"val_loss\")\nplt.plot(fbetas, color=\"cyan\", label=\"fbeta\")\nplt.legend()\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2023-06-26T16:01:25.269640Z","iopub.execute_input":"2023-06-26T16:01:25.270330Z","iopub.status.idle":"2023-06-26T16:01:25.374772Z","shell.execute_reply.started":"2023-06-26T16:01:25.270290Z","shell.execute_reply":"2023-06-26T16:01:25.373080Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Load best model\ncheckpoints = glob.glob(\"/kaggle/working/ckpts/*.pt\")\ncheckpoints.sort(key=os.path.getmtime)\nmodel_ckpt = checkpoints[-1]\nprint(model_ckpt)\n#model_ckpt  = \"/kaggle/working/ckpts/resnet18_3d_seg_12_0.56.pt\"\ncheckpoint = torch.load(model_ckpt)\nmodel.module.load_state_dict(checkpoint)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Get optimal theshold\nfinal_pred_mask, _ = validate(model, dataloader_valid, ModelConfig.tta)\nopt_f, opt_t = 0, 0\nfor threshold in np.arange(0.2, 1.0, 0.05):\n    fbeta = fbeta_score(final_pred_mask, gt_mask, threshold)\n    if fbeta > opt_f:\n        opt_f, opt_t = fbeta, threshold\n    print(f\"{threshold:.2f}: {fbeta.item():.2f}\")\nfinal_pred_mask = final_pred_mask.cpu().numpy()\nthresholded = np.zeros_like(final_pred_mask)\nthresholded[final_pred_mask >= opt_t] = 1","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"np_gt_mask = gt_mask.cpu().numpy()\n\n# Plot final model predictions\nfig, (ax0, ax1, ax2) = plt.subplots(1, 3)\nax0.imshow(final_pred_mask)\nax1.imshow(thresholded)\nax2.imshow(np_gt_mask)\nfig.suptitle(f'T: {opt_t:.2f} F0.5: {opt_f.item():.2f}')\n\nplt.show()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Uncomment to remove checkpoints if file persistence is on.\n# !rm /kaggle/working/ckpts/*","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}