{"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":"### Setup","metadata":{}},{"cell_type":"code","source":"# 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":{"_kg_hide-input":false,"_kg_hide-output":true,"execution":{"iopub.status.busy":"2023-04-08T09:58:33.085402Z","iopub.execute_input":"2023-04-08T09:58:33.085806Z","iopub.status.idle":"2023-04-08T09:59:03.755604Z","shell.execute_reply.started":"2023-04-08T09:58:33.085771Z","shell.execute_reply":"2023-04-08T09:59:03.754162Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import 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\nsys.path.append(\"/kaggle/input/resnet3d\")\nfrom resnet3d import generate_model","metadata":{"execution":{"iopub.status.busy":"2023-04-08T09:59:20.007164Z","iopub.execute_input":"2023-04-08T09:59:20.008111Z","iopub.status.idle":"2023-04-08T09:59:22.855990Z","shell.execute_reply.started":"2023-04-08T09:59:20.008029Z","shell.execute_reply":"2023-04-08T09:59:22.854957Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Config","metadata":{}},{"cell_type":"code","source":"LR = 1e-4\nEPOCHS = 50\nBATCH_SIZE = 32\nCROP_SIZE = 256\nZ_START = 24\nZ_DIMS = 16\nTRAIN_FRAGMENTS = [\"2\", \"3\"]\nTEST_FRAGMENT = \"1\"","metadata":{"execution":{"iopub.status.busy":"2023-04-08T09:59:29.988917Z","iopub.execute_input":"2023-04-08T09:59:29.989317Z","iopub.status.idle":"2023-04-08T09:59:29.995414Z","shell.execute_reply.started":"2023-04-08T09:59:29.989282Z","shell.execute_reply":"2023-04-08T09:59:29.994289Z"},"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-04-08T09:59:34.546609Z","iopub.execute_input":"2023-04-08T09:59:34.547136Z","iopub.status.idle":"2023-04-08T09:59:34.629370Z","shell.execute_reply.started":"2023-04-08T09:59:34.547047Z","shell.execute_reply":"2023-04-08T09:59:34.628252Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Visualise input","metadata":{}},{"cell_type":"code","source":"fragment = FRAGMENTS_ZARR[\"1\"]\nx, y = 2000, 2000\n\nfragment_cropped = fragment.surface_volume[y:y+CROP_SIZE, x:x+CROP_SIZE, Z_START:Z_START+Z_DIMS]\nimageio.mimwrite(\"fragment_crop.mp4\", fragment_cropped.transpose(2, 0, 1), \"ffmpeg\")\nVideo(\"fragment_crop.mp4\", height=256, width=256)","metadata":{"execution":{"iopub.status.busy":"2023-04-08T09:59:41.601878Z","iopub.execute_input":"2023-04-08T09:59:41.602498Z","iopub.status.idle":"2023-04-08T09:59:43.852279Z","shell.execute_reply.started":"2023-04-08T09:59:41.602461Z","shell.execute_reply":"2023-04-08T09:59:43.851076Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"mask_cropped = fragment.truth[y:y+CROP_SIZE, x:x+CROP_SIZE]\nir_cropped = fragment.infrared[y:y+CROP_SIZE, x:x+CROP_SIZE]\n\nplt.figure(figsize=(6, 3))\nplt.subplot(1, 2, 1)\nplt.imshow(mask_cropped, cmap=\"gray\")\nplt.axis(\"off\")\n\nplt.subplot(1, 2, 2)\nplt.imshow(ir_cropped, cmap=\"gray\")\nplt.axis(\"off\")\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2023-04-08T09:59:49.123401Z","iopub.execute_input":"2023-04-08T09:59:49.124528Z","iopub.status.idle":"2023-04-08T09:59:49.426630Z","shell.execute_reply.started":"2023-04-08T09:59:49.124470Z","shell.execute_reply":"2023-04-08T09:59:49.425221Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"del fragment, fragment_cropped, mask_cropped, ir_cropped\ngc.collect()","metadata":{"execution":{"iopub.status.busy":"2023-04-08T09:59:56.485316Z","iopub.execute_input":"2023-04-08T09:59:56.485887Z","iopub.status.idle":"2023-04-08T09:59:56.635506Z","shell.execute_reply.started":"2023-04-08T09:59:56.485850Z","shell.execute_reply":"2023-04-08T09:59:56.634240Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Dataloaders","metadata":{}},{"cell_type":"code","source":"class VesuviusTrain(Dataset):\n    def __init__(self, fragments):\n        self.fragments = fragments\n        self.xys = []\n        \n        for fragment in fragments:\n            H, W = FRAGMENTS_SHAPE[fragment]\n            for y in range(0, H-CROP_SIZE+1, CROP_SIZE):\n                for x in range(0, W-CROP_SIZE+1, 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 = Z_START, Z_START+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 + CROP_SIZE\n        y2 = y1 + 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        \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-04-08T10:01:19.088009Z","iopub.execute_input":"2023-04-08T10:01:19.088627Z","iopub.status.idle":"2023-04-08T10:01:19.101537Z","shell.execute_reply.started":"2023-04-08T10:01:19.088590Z","shell.execute_reply":"2023-04-08T10:01:19.100535Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class VesuviusVal(Dataset):\n    def __init__(self, fragment):\n        self.fragment = FRAGMENTS_ZARR[fragment]\n        self.xys = []\n        \n        H, W = FRAGMENTS_SHAPE[fragment]\n        for y in range(0, H-CROP_SIZE+1, CROP_SIZE):\n            for x in range(0, W-CROP_SIZE+1, CROP_SIZE):\n                self.xys.append((x, y))\n                \n    def __getitem__(self, i):\n        x1, y1 = self.xys[i]\n        x2, y2 = x1+CROP_SIZE, y1+CROP_SIZE\n        z1, z2 = Z_START, Z_START+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        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-04-08T10:01:20.868158Z","iopub.execute_input":"2023-04-08T10:01:20.868521Z","iopub.status.idle":"2023-04-08T10:01:20.878551Z","shell.execute_reply.started":"2023-04-08T10:01:20.868487Z","shell.execute_reply":"2023-04-08T10:01:20.877374Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"dataset_train = VesuviusTrain(TRAIN_FRAGMENTS)\ndataloader_train = DataLoader(dataset_train, batch_size=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)\ndataloader_valid = DataLoader(dataset_valid, batch_size=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-04-08T10:01:22.845530Z","iopub.execute_input":"2023-04-08T10:01:22.846418Z","iopub.status.idle":"2023-04-08T10:01:22.853813Z","shell.execute_reply.started":"2023-04-08T10:01:22.846380Z","shell.execute_reply":"2023-04-08T10:01:22.852531Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Model\n* Encoder is a 3D ResNet model. 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-04-08T10:01:27.047618Z","iopub.execute_input":"2023-04-08T10:01:27.047998Z","iopub.status.idle":"2023-04-08T10:01:27.062805Z","shell.execute_reply.started":"2023-04-08T10:01:27.047962Z","shell.execute_reply":"2023-04-08T10:01:27.061464Z"},"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-04-08T10:01:33.225416Z","iopub.execute_input":"2023-04-08T10:01:33.225790Z","iopub.status.idle":"2023-04-08T10:01:37.166510Z","shell.execute_reply.started":"2023-04-08T10:01:33.225757Z","shell.execute_reply":"2023-04-08T10:01:37.165460Z"},"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-04-08T10:01:41.743030Z","iopub.execute_input":"2023-04-08T10:01:41.743665Z","iopub.status.idle":"2023-04-08T10:01:41.750484Z","shell.execute_reply.started":"2023-04-08T10:01:41.743627Z","shell.execute_reply":"2023-04-08T10:01:41.749225Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Training","metadata":{}},{"cell_type":"code","source":"scaler = amp.GradScaler()\ncriterion = nn.BCEWithLogitsLoss()\noptimizer = torch.optim.AdamW(model.parameters(), lr=LR)\nscheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer, max_lr=LR,\n                                                steps_per_epoch=10, epochs=EPOCHS//10,\n                                                pct_start=0.1)","metadata":{"execution":{"iopub.status.busy":"2023-04-08T10:01:44.442792Z","iopub.execute_input":"2023-04-08T10:01:44.443496Z","iopub.status.idle":"2023-04-08T10:01:44.451399Z","shell.execute_reply.started":"2023-04-08T10:01:44.443456Z","shell.execute_reply":"2023-04-08T10:01:44.448998Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"gt_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-04-08T10:01:46.362523Z","iopub.execute_input":"2023-04-08T10:01:46.362894Z","iopub.status.idle":"2023-04-08T10:01:46.855167Z","shell.execute_reply.started":"2023-04-08T10:01:46.362861Z","shell.execute_reply":"2023-04-08T10:01:46.854111Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"for epoch in range(1, 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, mloss_val, val_metric = 0.0, 0.0, 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}/{EPOCHS}\", gpu_mem, cur_lr,\n                                                              f\"Loss: {mloss_train / (i + 1):.4f}\"))\n        \n    scheduler.step()\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            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    \n    for threshold in np.arange(0.2, 0.65, 0.05):\n        fbeta = fbeta_score(final_pred_mask, gt_mask, threshold)\n        print(f\"Threshold : {threshold:.2f}\\tFBeta : {fbeta:.6f}\")\n    \n    if epoch >= 10:\n        torch.save(model.module.state_dict(), f\"./ckpts/resnet18_3d_seg_epoch_{epoch}.pt\")\n\n    if epoch == 30:\n        break\n\n    print()","metadata":{"execution":{"iopub.status.busy":"2023-04-08T10:01:50.128035Z","iopub.execute_input":"2023-04-08T10:01:50.128730Z","iopub.status.idle":"2023-04-08T10:05:03.395476Z","shell.execute_reply.started":"2023-04-08T10:01:50.128692Z","shell.execute_reply":"2023-04-08T10:05:03.392355Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}