{"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-05-19T01:11:53.133957Z","iopub.execute_input":"2023-05-19T01:11:53.134339Z","iopub.status.idle":"2023-05-19T01:12:32.505726Z","shell.execute_reply.started":"2023-05-19T01:11:53.134303Z","shell.execute_reply":"2023-05-19T01:12:32.504251Z"},"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-05-19T01:12:32.508412Z","iopub.execute_input":"2023-05-19T01:12:32.508766Z","iopub.status.idle":"2023-05-19T01:12:35.317259Z","shell.execute_reply.started":"2023-05-19T01:12:32.508731Z","shell.execute_reply":"2023-05-19T01:12:35.316195Z"},"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-05-19T01:12:35.318661Z","iopub.execute_input":"2023-05-19T01:12:35.319565Z","iopub.status.idle":"2023-05-19T01:12:35.326741Z","shell.execute_reply.started":"2023-05-19T01:12:35.319523Z","shell.execute_reply":"2023-05-19T01:12:35.325312Z"},"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-05-19T01:12:35.330266Z","iopub.execute_input":"2023-05-19T01:12:35.330999Z","iopub.status.idle":"2023-05-19T01:12:35.528869Z","shell.execute_reply.started":"2023-05-19T01:12:35.330958Z","shell.execute_reply":"2023-05-19T01:12:35.527671Z"},"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-05-19T01:12:35.533031Z","iopub.execute_input":"2023-05-19T01:12:35.533376Z","iopub.status.idle":"2023-05-19T01:12:38.617625Z","shell.execute_reply.started":"2023-05-19T01:12:35.533344Z","shell.execute_reply":"2023-05-19T01:12:38.616568Z"},"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-05-19T01:12:38.619765Z","iopub.execute_input":"2023-05-19T01:12:38.620165Z","iopub.status.idle":"2023-05-19T01:12:38.880977Z","shell.execute_reply.started":"2023-05-19T01:12:38.620122Z","shell.execute_reply":"2023-05-19T01:12:38.879711Z"},"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-05-19T01:12:38.882750Z","iopub.execute_input":"2023-05-19T01:12:38.883881Z","iopub.status.idle":"2023-05-19T01:12:39.043927Z","shell.execute_reply.started":"2023-05-19T01:12:38.883839Z","shell.execute_reply":"2023-05-19T01:12:39.042617Z"},"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-05-19T01:12:39.045984Z","iopub.execute_input":"2023-05-19T01:12:39.047913Z","iopub.status.idle":"2023-05-19T01:12:39.063597Z","shell.execute_reply.started":"2023-05-19T01:12:39.047869Z","shell.execute_reply":"2023-05-19T01:12:39.062270Z"},"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-05-19T01:12:39.065215Z","iopub.execute_input":"2023-05-19T01:12:39.065742Z","iopub.status.idle":"2023-05-19T01:12:39.080097Z","shell.execute_reply.started":"2023-05-19T01:12:39.065683Z","shell.execute_reply":"2023-05-19T01:12:39.078878Z"},"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-05-19T01:12:39.084790Z","iopub.execute_input":"2023-05-19T01:12:39.085120Z","iopub.status.idle":"2023-05-19T01:12:39.095864Z","shell.execute_reply.started":"2023-05-19T01:12:39.085090Z","shell.execute_reply":"2023-05-19T01:12:39.094902Z"},"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-05-19T01:12:39.097111Z","iopub.execute_input":"2023-05-19T01:12:39.098122Z","iopub.status.idle":"2023-05-19T01:12:39.114796Z","shell.execute_reply.started":"2023-05-19T01:12:39.098084Z","shell.execute_reply":"2023-05-19T01:12:39.113683Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model = SegModel()\nmodel.load_pretrained_weights(torch.load(\"r3d18_K_200ep.pth\")[\"state_dict\"])\n# model = nn.DataParallel(model, device_ids=[0, 1])\nmodel = model.cuda()","metadata":{"execution":{"iopub.status.busy":"2023-05-19T01:12:39.116492Z","iopub.execute_input":"2023-05-19T01:12:39.116938Z","iopub.status.idle":"2023-05-19T01:12:42.842737Z","shell.execute_reply.started":"2023-05-19T01:12:39.116899Z","shell.execute_reply":"2023-05-19T01:12:42.841618Z"},"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-05-19T01:12:42.844356Z","iopub.execute_input":"2023-05-19T01:12:42.844877Z","iopub.status.idle":"2023-05-19T01:12:42.854555Z","shell.execute_reply.started":"2023-05-19T01:12:42.844834Z","shell.execute_reply":"2023-05-19T01:12:42.853241Z"},"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-05-19T01:12:42.856590Z","iopub.execute_input":"2023-05-19T01:12:42.857383Z","iopub.status.idle":"2023-05-19T01:12:42.870348Z","shell.execute_reply.started":"2023-05-19T01:12:42.857340Z","shell.execute_reply":"2023-05-19T01:12:42.869176Z"},"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-05-19T01:12:42.871890Z","iopub.execute_input":"2023-05-19T01:12:42.872344Z","iopub.status.idle":"2023-05-19T01:12:43.360831Z","shell.execute_reply.started":"2023-05-19T01:12:42.872304Z","shell.execute_reply":"2023-05-19T01:12:43.359785Z"},"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.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-05-19T01:12:43.362520Z","iopub.execute_input":"2023-05-19T01:12:43.363417Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}