{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":94147,"databundleVersionId":11390004,"sourceType":"competition"}],"dockerImageVersionId":30919,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install -qq monai\n!pip install -qq segmentation_models_pytorch","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-01T09:01:57.861137Z","iopub.execute_input":"2025-04-01T09:01:57.861363Z","iopub.status.idle":"2025-04-01T09:02:11.707258Z","shell.execute_reply.started":"2025-04-01T09:01:57.861341Z","shell.execute_reply":"2025-04-01T09:02:11.706209Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import cv2\nimport h5py\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\n\nimport torch\nfrom torch import nn\nfrom torch import optim\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\n\nfrom scipy.stats import spearmanr\nfrom scipy.interpolate import griddata\n\nimport albumentations as albu\nfrom albumentations.pytorch import ToTensorV2\n\nimport pytorch_lightning as pl\nfrom pytorch_lightning import Trainer\nfrom pytorch_lightning import seed_everything\nfrom pytorch_lightning.callbacks import LearningRateMonitor\nfrom pytorch_lightning.callbacks import ModelCheckpoint\nfrom pytorch_lightning.loggers import WandbLogger\n\nimport segmentation_models_pytorch as smp\n\nfrom monai.inferers import sliding_window_inference\n\nimport timm\nfrom timm.optim import create_optimizer_v2\nfrom timm.scheduler import create_scheduler_v2","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-04-01T09:02:14.344764Z","iopub.execute_input":"2025-04-01T09:02:14.345111Z","iopub.status.idle":"2025-04-01T09:02:46.179629Z","shell.execute_reply.started":"2025-04-01T09:02:14.345081Z","shell.execute_reply":"2025-04-01T09:02:46.178749Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import wandb\n\nwandb_api = '######'\nwandb.login(key=wandb_api)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"H5PY_PATH = '/kaggle/input/elucidata-test-data/elucidata_ai_challenge_data.h5'\nIMAGE_SIZE = 256\nSTRIDE = 128\nPADDING = 128\nNUM_EPOCHS = 100\nBATCH_SIZE = 4\nPOINT_RADIUS = 12\nLEARNING_RATE = 2e-4\n\nTRAIN_SLIDE_IDS = ['S_2', 'S_3', 'S_4', 'S_5', 'S_6']\nVAL_SLIDE_IDS = ['S_1']\nTEST_SLIDE_IDS = ['S_7']\nSEED = 42\n\nTRAIN_AUG = albu.Compose([\n    albu.ColorJitter(p=1.0),\n    albu.RandomRotate90(p=1.0),\n    albu.HorizontalFlip(p=0.5),\n    albu.VerticalFlip(p=0.5),\n    albu.Transpose(p=0.5),\n    albu.Normalize(max_pixel_value=1),\n    ToTensorV2(transpose_mask=True),\n])\n\nEVAL_AUG = albu.Compose([\n    albu.Normalize(max_pixel_value=1),\n    ToTensorV2(transpose_mask=True),\n])\n\nseed_everything(SEED)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-01T09:02:46.180772Z","iopub.execute_input":"2025-04-01T09:02:46.181691Z","iopub.status.idle":"2025-04-01T09:02:46.199208Z","shell.execute_reply.started":"2025-04-01T09:02:46.181651Z","shell.execute_reply":"2025-04-01T09:02:46.198298Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def calculate_loss(mask, pred_mask, loss_fn):\n    loss = loss_fn(mask, (mask > 0) * pred_mask)\n    return loss","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-01T09:02:46.201015Z","iopub.execute_input":"2025-04-01T09:02:46.201371Z","iopub.status.idle":"2025-04-01T09:02:46.2106Z","shell.execute_reply.started":"2025-04-01T09:02:46.201346Z","shell.execute_reply":"2025-04-01T09:02:46.209783Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def split_image_with_overlap(image, patch_size=512, stride=256):\n    h, w, c = image.shape\n    patches = []\n\n    i = 0\n    while i + patch_size <= h:\n        j = 0\n        while j + patch_size <= w:\n            patch = image[i:i+patch_size, j:j+patch_size, :]\n            patches.append(patch)\n            j += stride\n        i += stride\n\n    return np.stack(patches)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-01T09:02:46.212012Z","iopub.execute_input":"2025-04-01T09:02:46.21236Z","iopub.status.idle":"2025-04-01T09:02:46.224749Z","shell.execute_reply.started":"2025-04-01T09:02:46.212338Z","shell.execute_reply":"2025-04-01T09:02:46.223946Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def crop_image_with_padding(image, xs, ys):\n    height, width, _ = image.shape\n\n    xmin = max(int(xs.min() - PADDING), 0)\n    xmax = min(int(xs.max() + PADDING), width)\n    ymin = max(int(ys.min() - PADDING), 0)\n    ymax = min(int(ys.max() + PADDING), height)\n\n    image = image[ymin:ymax, xmin:xmax]\n    xs = xs - xmin\n    ys = ys - ymin\n\n    return image, xs, ys","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-01T09:02:46.225621Z","iopub.execute_input":"2025-04-01T09:02:46.225933Z","iopub.status.idle":"2025-04-01T09:02:46.244057Z","shell.execute_reply.started":"2025-04-01T09:02:46.225903Z","shell.execute_reply":"2025-04-01T09:02:46.243249Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def create_mask(image, xs, ys, ch):\n    height, width, _ = image.shape\n    ch = (1 + np.argsort(np.argsort(ch, axis=1), axis=1)) / 36\n    mask = []\n    for i in range(35):\n        mask_slice = np.zeros((height, width), dtype=np.float32)\n        for x, y, c in zip(xs, ys, ch[:, i]):\n            mask_slice = cv2.circle(mask_slice, (int(x), int(y)), radius=POINT_RADIUS, color=c, thickness=-1)\n        mask.append(mask_slice)\n\n    return np.stack(mask, axis=-1)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-01T09:02:46.244817Z","iopub.execute_input":"2025-04-01T09:02:46.245198Z","iopub.status.idle":"2025-04-01T09:02:46.260546Z","shell.execute_reply.started":"2025-04-01T09:02:46.245162Z","shell.execute_reply":"2025-04-01T09:02:46.259726Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def calculate_average_spearman(gt, pred):\n    if gt.shape != pred.shape:\n        raise ValueError(\"Ground truth and prediction must have the same shape\")\n\n    n, target = gt.shape\n    correlations = []\n\n    for i in range(n):\n        gt_sample = gt[i, :]\n        pred_sample = pred[i, :]\n\n        corr, _ = spearmanr(gt_sample, pred_sample)\n        correlations.append(corr)\n\n    valid_correlations = [c for c in correlations if not np.isnan(c)]\n    if not valid_correlations:\n        return np.nan\n\n    return np.mean(valid_correlations)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-01T09:02:46.261428Z","iopub.execute_input":"2025-04-01T09:02:46.261649Z","iopub.status.idle":"2025-04-01T09:02:46.276644Z","shell.execute_reply.started":"2025-04-01T09:02:46.26163Z","shell.execute_reply":"2025-04-01T09:02:46.275799Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"@torch.no_grad()\ndef create_submission(model, loader, gt_exist=False):\n    batch = next(iter(loader))\n    image = batch['image'].cuda()\n    spots = batch['spot'][0]\n\n    pred_masks = []\n    for rot in range(4):\n        rotated_image = torch.rot90(image, k=rot, dims=(2, 3))\n        pred_mask = sliding_window_inference(\n            rotated_image,\n            roi_size=(IMAGE_SIZE, IMAGE_SIZE),\n            overlap=0.5,\n            sw_batch_size=BATCH_SIZE,\n            predictor=model,\n            mode='constant',\n        )\n\n        pred_mask = torch.rot90(pred_mask, k=-rot, dims=(2, 3))\n        pred_masks.append(pred_mask)\n\n    pred_mask = torch.stack(pred_masks).mean(dim=0)\n    pred_mask = pred_mask.detach().cpu()\n    pred_mask = pred_mask.numpy()\n    pred_mask = pred_mask.clip(0)\n    pred_mask = pred_mask.squeeze()\n\n    predictions = []\n    if gt_exist:\n        for idx, (x, y, *gt) in enumerate(spots):\n            x, y, gt = int(x), int(y), torch.tensor(gt).numpy()\n            row = [idx, *pred_mask[:, y, x], *gt]\n            predictions.append(row)\n        predictions = pd.DataFrame(predictions, columns=['ID', *([f'CP{i}' for i in range(1, 36)] + [f'CT{i}' for i in range(1, 36)])])\n    else:\n        for idx, (x, y, *gt) in enumerate(spots):\n            x, y = int(x), int(y)\n            row = [idx, *pred_mask[:, y, x]]\n            predictions.append(row)\n        predictions = pd.DataFrame(predictions, columns=['ID', *[f'C{i}' for i in range(1, 36)]])\n\n    return predictions, pred_mask","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-01T09:02:46.278873Z","iopub.execute_input":"2025-04-01T09:02:46.279161Z","iopub.status.idle":"2025-04-01T09:02:46.292103Z","shell.execute_reply.started":"2025-04-01T09:02:46.279135Z","shell.execute_reply":"2025-04-01T09:02:46.291228Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class TrainDataset(Dataset):\n    def __init__(self, h5py_path, slide_names, image_size, stride, transforms):\n        super().__init__()\n        self.images = []\n        self.masks = []\n        self.tfs = transforms\n\n        with h5py.File(h5py_path, 'r') as h5file:\n            all_images = h5file['images/Train']\n            all_spots = h5file['spots/Train']\n\n            image_list, mask_list = [], []\n            for slide_name in slide_names:\n                image = np.array(all_images[slide_name])\n                spots = np.array([list(spot) for spot in list(all_spots[slide_name])])\n\n                xs, ys, ch = spots[:, 0], spots[:, 1], spots[:, 2:]\n\n                if slide_name in ['S_1', 'S_2']:\n                    xs, ys = xs - 64, ys - 64\n\n                image, xs, ys = crop_image_with_padding(image, xs, ys)\n                mask = create_mask(image, xs, ys, ch)\n\n                images = split_image_with_overlap(image, image_size, stride)\n                masks = split_image_with_overlap(mask, image_size, stride)\n\n                image_list.append(images)\n                mask_list.append(masks)\n\n            self.images = np.concatenate(image_list)\n            self.masks = np.concatenate(mask_list)\n\n    def __getitem__(self, idx):\n        return self.tfs(\n            image=self.images[idx],\n            mask=self.masks[idx],\n        )\n\n    def __len__(self):\n        return len(self.images)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-01T09:02:46.293259Z","iopub.execute_input":"2025-04-01T09:02:46.293518Z","iopub.status.idle":"2025-04-01T09:02:46.310332Z","shell.execute_reply.started":"2025-04-01T09:02:46.293491Z","shell.execute_reply":"2025-04-01T09:02:46.309454Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class EvalDataset(Dataset):\n    def __init__(self, h5py_path, slide_names, transforms):\n        super().__init__()\n        self.images = []\n        self.spots = []\n        self.tfs = transforms\n\n        with h5py.File(h5py_path, 'r') as h5file:\n            all_images = h5file['images/Train']\n            all_spots = h5file['spots/Train']\n\n            image_list, spot_list = [], []\n            for slide_name in slide_names:\n                image = np.array(all_images[slide_name])\n                spots = np.array([list(spot) for spot in list(all_spots[slide_name])])\n\n                image_list.append(image)\n                spot_list.append(spots)\n\n            self.images = np.stack(image_list)\n            self.spots = np.stack(spot_list)\n\n    def __getitem__(self, idx):\n        return {\n            'image': self.tfs(image=self.images[idx])['image'],\n            'spot': self.spots[idx],\n        }\n\n    def __len__(self):\n        return len(self.images)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-01T09:02:46.311153Z","iopub.execute_input":"2025-04-01T09:02:46.311475Z","iopub.status.idle":"2025-04-01T09:02:46.324992Z","shell.execute_reply.started":"2025-04-01T09:02:46.311445Z","shell.execute_reply":"2025-04-01T09:02:46.324173Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class TestDataset(Dataset):\n    def __init__(self, h5py_path, slide_names, transforms):\n        super().__init__()\n        self.images = []\n        self.spots = []\n        self.tfs = transforms\n\n        with h5py.File(h5py_path, 'r') as h5file:\n            all_images = h5file['images/Test']\n            all_spots = h5file['spots/Test']\n\n            image_list, spot_list = [], []\n            for slide_name in slide_names:\n                image = np.array(all_images[slide_name])\n                spots = np.array([list(spot) for spot in list(all_spots[slide_name])])\n\n                image_list.append(image)\n                spot_list.append(spots)\n\n            self.images = np.stack(image_list)\n            self.spots = np.stack(spot_list)\n\n    def __getitem__(self, idx):\n        return {\n            'image': self.tfs(image=self.images[idx])['image'],\n            'spot': self.spots[idx],\n        }\n\n    def __len__(self):\n        return len(self.images)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-01T09:02:46.325758Z","iopub.execute_input":"2025-04-01T09:02:46.326011Z","iopub.status.idle":"2025-04-01T09:02:46.345395Z","shell.execute_reply.started":"2025-04-01T09:02:46.325991Z","shell.execute_reply":"2025-04-01T09:02:46.344576Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class Model(pl.LightningModule):\n    def __init__(self):\n        super().__init__()\n        self.model = smp.create_model(\n            arch='Unet',\n            encoder_name='tu-tf_efficientnetv2_s',\n            encoder_weights='imagenet',\n            in_channels=3,\n            classes=35,\n            activation='sigmoid',\n        )\n\n        self.loss_fn = nn.MSELoss()\n        self.validation_step_outputs = []\n\n    def forward(self, image):\n        return self.model(image)\n\n    def training_step(self, batch, batch_idx):\n        image, mask = batch['image'], batch['mask']\n        logits_mask = self.forward(image)\n\n        loss = calculate_loss(mask, logits_mask, self.loss_fn)\n        return loss\n\n    def validation_step(self, batch, batch_idx):\n        image = batch['image']\n        spots = batch['spot'][0]\n\n        pred_mask = sliding_window_inference(\n            image,\n            roi_size=(IMAGE_SIZE, IMAGE_SIZE),\n            overlap=0.25,\n            sw_batch_size=BATCH_SIZE,\n            predictor=self.forward,\n            mode='constant',\n        )\n\n        pred_mask = pred_mask.detach().cpu()\n        pred_mask = pred_mask.numpy()\n        pred_mask = pred_mask.clip(0)\n        pred_mask = pred_mask.squeeze()\n\n        preds, gts = [], []\n        for idx, (x, y, *gt) in enumerate(spots):\n            x = int(x.detach().cpu().numpy())\n            y = int(y.detach().cpu().numpy())\n            gt_spot = torch.tensor(gt).detach().cpu().numpy()\n\n            pred_spot = pred_mask[:, y, x]\n            preds.append(pred_spot)\n            gts.append(gt_spot)\n\n        preds = np.stack(preds)\n        gts = np.stack(gts)\n\n        score = calculate_average_spearman(gts, preds)\n        self.validation_step_outputs.append(score)\n\n    def on_validation_epoch_end(self):\n        score = torch.tensor(self.validation_step_outputs)\n        score_avg = score.mean()\n\n        self.log(f'val_score', score_avg, on_step=False, on_epoch=True, prog_bar=True)\n        self.validation_step_outputs.clear()\n\n    def configure_optimizers(self):\n        optimizer = create_optimizer_v2(\n            self.parameters(),\n            opt='adamw',\n            lr=LEARNING_RATE,\n            weight_decay=1e-2,\n            filter_bias_and_bn=True,\n        )\n\n        scheduler, _ = create_scheduler_v2(\n            optimizer,\n            sched='cosine',\n            num_epochs=NUM_EPOCHS,\n            min_lr=0.0,\n            warmup_lr=1e-5,\n            warmup_epochs=0,\n            warmup_prefix=False,\n            step_on_epochs=True,\n        )\n\n        return {\n            'optimizer': optimizer,\n            'lr_scheduler': {\n                'scheduler': scheduler,\n                'interval': 'epoch',\n            },\n        }\n\n    def lr_scheduler_step(self, scheduler, metric):\n        scheduler.step(epoch=self.current_epoch)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-01T09:02:46.34644Z","iopub.execute_input":"2025-04-01T09:02:46.346729Z","iopub.status.idle":"2025-04-01T09:02:46.362034Z","shell.execute_reply.started":"2025-04-01T09:02:46.346697Z","shell.execute_reply":"2025-04-01T09:02:46.361148Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_loader = DataLoader(\n    TrainDataset(\n        h5py_path=H5PY_PATH,\n        slide_names=TRAIN_SLIDE_IDS,\n        image_size=IMAGE_SIZE,\n        stride=STRIDE,\n        transforms=TRAIN_AUG,\n    ),\n    batch_size=BATCH_SIZE,\n    shuffle=True,\n    drop_last=True,\n    num_workers=2,\n)\n\nvalid_loader = DataLoader(\n    EvalDataset(\n        h5py_path=H5PY_PATH,\n        slide_names=VAL_SLIDE_IDS,\n        transforms=EVAL_AUG,\n    ),\n    batch_size=1,\n    shuffle=False,\n    drop_last=False,\n    num_workers=2,\n)\n\ntest_loader = DataLoader(\n    TestDataset(\n        h5py_path=H5PY_PATH,\n        slide_names=TEST_SLIDE_IDS,\n        transforms=EVAL_AUG,\n    ),\n    batch_size=1,\n    shuffle=False,\n    drop_last=False,\n    num_workers=2,\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-01T09:02:46.362957Z","iopub.execute_input":"2025-04-01T09:02:46.363239Z","iopub.status.idle":"2025-04-01T09:03:00.323331Z","shell.execute_reply.started":"2025-04-01T09:02:46.363212Z","shell.execute_reply":"2025-04-01T09:03:00.32223Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"trainer = Trainer(\n    max_epochs=NUM_EPOCHS,\n    callbacks = [\n        ModelCheckpoint(\n            dirpath='checkpoint',\n            filename='best_model',\n            save_weights_only=True,\n            save_top_k=1,\n            monitor='val_score',\n            mode='max',\n            save_last=False,\n        ),\n        LearningRateMonitor(logging_interval='step'),\n    ],\n    logger=[WandbLogger(project='elucidata_final_segmentation_project', log_model=False)],\n    precision='16-mixed',\n    deterministic=True,\n    benchmark=False,\n    check_val_every_n_epoch=5,\n    accelerator='gpu',\n    devices=[0,],\n    log_every_n_steps=1,\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-01T09:03:05.622134Z","iopub.execute_input":"2025-04-01T09:03:05.622446Z","iopub.status.idle":"2025-04-01T09:03:05.701786Z","shell.execute_reply.started":"2025-04-01T09:03:05.622423Z","shell.execute_reply":"2025-04-01T09:03:05.700833Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"trainer.fit(\n    Model(),\n    train_dataloaders=train_loader,\n    val_dataloaders=valid_loader,\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-01T09:03:05.781744Z","iopub.execute_input":"2025-04-01T09:03:05.782105Z","iopub.status.idle":"2025-04-01T09:03:28.136788Z","shell.execute_reply.started":"2025-04-01T09:03:05.782077Z","shell.execute_reply":"2025-04-01T09:03:28.135086Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"model_path = 'checkpoint/best_model.ckpt'\n\nmodel = Model()\nmodel.load_state_dict(torch.load(model_path, weights_only=True)['state_dict'])\nmodel = model.cuda()\nmodel = model.eval()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df_subm, mask = create_submission(model, test_loader, gt_exist=False)\ndf_subm.to_csv(f'submission.csv', index=False)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df_subm, mask = create_submission(model, valid_loader, gt_exist=True)\ndf_subm.to_csv(f'oof_submission.csv', index=False)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}