{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":20561,"databundleVersionId":1187889,"sourceType":"competition"}],"dockerImageVersionId":31194,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-11-28T06:21:14.684531Z","iopub.execute_input":"2025-11-28T06:21:14.685378Z","iopub.status.idle":"2025-11-28T06:21:14.953727Z","shell.execute_reply.started":"2025-11-28T06:21:14.685343Z","shell.execute_reply":"2025-11-28T06:21:14.953036Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport torchvision\nimport albumentations as A\nimport transformers\nfrom albumentations.pytorch import ToTensorV2\nimport os\nfrom torchmetrics.functional.classification import multiclass_f1_score\nfrom tqdm.auto import tqdm\nimport pandas as pd","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-28T06:21:15.904934Z","iopub.execute_input":"2025-11-28T06:21:15.905358Z","iopub.status.idle":"2025-11-28T06:21:25.593159Z","shell.execute_reply.started":"2025-11-28T06:21:15.905327Z","shell.execute_reply":"2025-11-28T06:21:25.592368Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"torch.cuda.is_available()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-28T06:21:25.594515Z","iopub.execute_input":"2025-11-28T06:21:25.594919Z","iopub.status.idle":"2025-11-28T06:21:25.656502Z","shell.execute_reply.started":"2025-11-28T06:21:25.594899Z","shell.execute_reply":"2025-11-28T06:21:25.655525Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class Dataset(torchvision.datasets.VisionDataset):\n    def __init__(self, transform, images, masks):\n        super().__init__(transform)\n        self.images = images\n        self.masks = masks\n        self.transform = transform\n\n    def __getitem__(self, item):\n        image = self.images[item]\n        mask = self.masks[item]\n        transformed = self.transform(image=image, mask=mask)\n        image = transformed[\"image\"]\n        mask = transformed[\"mask\"]\n        mask = mask.permute((2, 0, 1))\n        return image, mask\n\n    def __len__(self):\n        return len(self.images)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-28T06:21:25.657619Z","iopub.execute_input":"2025-11-28T06:21:25.657893Z","iopub.status.idle":"2025-11-28T06:21:25.688571Z","shell.execute_reply.started":"2025-11-28T06:21:25.65787Z","shell.execute_reply":"2025-11-28T06:21:25.687702Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"prefix = \"/kaggle/input/covid-segmentation/\"\n\nimages_medseg = np.load(os.path.join(prefix, \"images_medseg.npy\")).astype(np.float32)\nimages_radiopedia = np.load(os.path.join(prefix, \"images_radiopedia.npy\")).astype(np.float32)\nmasks_medseg = np.load(os.path.join(prefix, \"masks_medseg.npy\")).astype(np.int8)\nmasks_radiopedia = np.load(os.path.join(prefix, \"masks_radiopedia.npy\")).astype(np.int8)\ntest_images_medseg = np.load(os.path.join(prefix, \"test_images_medseg.npy\")).astype(np.float32)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-28T06:21:25.68998Z","iopub.execute_input":"2025-11-28T06:21:25.690198Z","iopub.status.idle":"2025-11-28T06:21:42.088088Z","shell.execute_reply.started":"2025-11-28T06:21:25.690181Z","shell.execute_reply":"2025-11-28T06:21:42.087227Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# visualization of NumPy arrays\ndef visualize(image_batch, mask_batch, num_samples=8):\n    fig, ax = plt.subplots(nrows=1 + mask_batch.shape[-1], ncols=num_samples,\n                           figsize=(num_samples * 2, mask_batch.shape[-1] * 2))\n    for i in range(num_samples):\n        ax[0, i].imshow(image_batch[i, :, :, 0])\n        for channel in range(1, mask_batch.shape[-1] + 1):\n            ax[channel, i].imshow(mask_batch[i, :, :, channel - 1])\n    plt.tight_layout()\n    plt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-28T06:21:42.088839Z","iopub.execute_input":"2025-11-28T06:21:42.089087Z","iopub.status.idle":"2025-11-28T06:21:42.09474Z","shell.execute_reply.started":"2025-11-28T06:21:42.089065Z","shell.execute_reply":"2025-11-28T06:21:42.093721Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# visualization of Torch tensors\ndef visualize_transposed(image_batch, mask_batch, num_samples=8):\n    num_samples = min(num_samples, len(image_batch))\n    fig, ax = plt.subplots(nrows=1 + mask_batch.shape[1], ncols=num_samples,\n                           figsize=(num_samples * 2, mask_batch.shape[1] * 2))\n    for i in range(num_samples):\n        ax[0, i].imshow(image_batch[i, 0, :, :])\n        for channel in range(1, mask_batch.shape[1] + 1):\n            ax[channel, i].imshow(mask_batch[i, channel - 1, :, :])\n    plt.tight_layout()\n    plt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-28T06:21:42.095662Z","iopub.execute_input":"2025-11-28T06:21:42.095912Z","iopub.status.idle":"2025-11-28T06:21:42.120061Z","shell.execute_reply.started":"2025-11-28T06:21:42.095889Z","shell.execute_reply":"2025-11-28T06:21:42.119425Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# define transforms for datasets\nTARGET_SIZE = 512\n\ntrain_transforms = A.Compose([\n    A.RandomResizedCrop(size=(TARGET_SIZE, TARGET_SIZE),\n                        scale=(0.7, 1.0), ratio=(0.9, 1.1), p=0.5),\n    A.HorizontalFlip(p=0.3),\n    A.RandomBrightnessContrast(p=0.1, brightness_limit=0.1, contrast_limit=0.1),\n    A.ElasticTransform(alpha=40, sigma=4, p=0.15),\n    A.GridDistortion(num_steps=5, distort_limit=0.2, p=0.15),\n    A.OpticalDistortion(distort_limit=0.5, p=0.15),\n    ToTensorV2()\n])\n\ntest_transforms = A.Compose([\n    ToTensorV2()\n])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-28T06:21:42.120713Z","iopub.execute_input":"2025-11-28T06:21:42.120914Z","iopub.status.idle":"2025-11-28T06:21:42.141408Z","shell.execute_reply.started":"2025-11-28T06:21:42.120897Z","shell.execute_reply":"2025-11-28T06:21:42.140617Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# split raw data, clip it, normalize it\ntrain_size = 0.8","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-27T20:27:28.797008Z","iopub.execute_input":"2025-11-27T20:27:28.797549Z","iopub.status.idle":"2025-11-27T20:27:28.800893Z","shell.execute_reply.started":"2025-11-27T20:27:28.797524Z","shell.execute_reply":"2025-11-27T20:27:28.80018Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# prepare train\ntrain_images = np.concatenate((images_medseg[0:int(train_size * len(images_medseg))], images_radiopedia))\n\n# clip and normalize train\ntrain_images = np.clip(train_images, -1500, 500)\ntrain_images = (train_images + 1500) / 2000\n\ntrain_masks_c = np.concatenate((masks_medseg[0:int(train_size * len(images_medseg))], masks_radiopedia))\n\n# unite class 3 and 4, treat it as background\ntrain_masks = np.concatenate(\n    (train_masks_c[:, :, :, 0:2], np.logical_or(train_masks_c[:, :, :, 2], train_masks_c[:, :, :, 3])[..., np.newaxis]),\n    axis=-1)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-27T20:27:29.420769Z","iopub.execute_input":"2025-11-27T20:27:29.421274Z","iopub.status.idle":"2025-11-27T20:27:33.173969Z","shell.execute_reply.started":"2025-11-27T20:27:29.42125Z","shell.execute_reply":"2025-11-27T20:27:33.173036Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# prepare val\nval_images = images_medseg[int(train_size * len(images_medseg)):]\n\n# clip and normalize val\nval_images = np.clip(val_images, -1500, 500)\nval_images = (val_images + 1500) / 2000\n\nval_masks_c = masks_medseg[int(train_size * len(images_medseg)):]\n\n# unite class 3 and 4, treat it as background\nval_masks = np.concatenate(\n    (val_masks_c[:, :, :, 0:2], np.logical_or(val_masks_c[:, :, :, 2], val_masks_c[:, :, :, 3])[..., np.newaxis]),\n    axis=-1)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-27T20:27:33.175306Z","iopub.execute_input":"2025-11-27T20:27:33.175609Z","iopub.status.idle":"2025-11-27T20:27:33.250867Z","shell.execute_reply.started":"2025-11-27T20:27:33.17559Z","shell.execute_reply":"2025-11-27T20:27:33.249938Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# prepare test\ntest_images_medseg = np.clip(test_images_medseg, -1500, 500)\ntest_images_medseg = (test_images_medseg + 1500) / 2000","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-28T06:21:42.143391Z","iopub.execute_input":"2025-11-28T06:21:42.143707Z","iopub.status.idle":"2025-11-28T06:21:42.175104Z","shell.execute_reply.started":"2025-11-28T06:21:42.143685Z","shell.execute_reply":"2025-11-28T06:21:42.174385Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# prepare datasets\ntrain_dataset = Dataset(train_transforms, train_images, train_masks)\nval_dataset = Dataset(test_transforms, val_images, val_masks)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-27T20:27:33.265978Z","iopub.execute_input":"2025-11-27T20:27:33.266291Z","iopub.status.idle":"2025-11-27T20:27:33.274748Z","shell.execute_reply.started":"2025-11-27T20:27:33.266273Z","shell.execute_reply":"2025-11-27T20:27:33.273926Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# prepare loaders\ntrain_loader = torch.utils.data.dataloader.DataLoader(train_dataset, batch_size=12, shuffle=True, pin_memory=True, num_workers=3, persistent_workers=True)\nval_loader = torch.utils.data.dataloader.DataLoader(val_dataset, batch_size=12, shuffle=False, pin_memory=True)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-27T20:27:33.275566Z","iopub.execute_input":"2025-11-27T20:27:33.275782Z","iopub.status.idle":"2025-11-27T20:27:33.289911Z","shell.execute_reply.started":"2025-11-27T20:27:33.275765Z","shell.execute_reply":"2025-11-27T20:27:33.289161Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"device = \"cuda\"","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-28T06:21:42.176066Z","iopub.execute_input":"2025-11-28T06:21:42.176348Z","iopub.status.idle":"2025-11-28T06:21:42.18045Z","shell.execute_reply.started":"2025-11-28T06:21:42.176324Z","shell.execute_reply":"2025-11-28T06:21:42.179833Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# define config and processor for HF Mask2Former\nconfig = transformers.Mask2FormerConfig.from_pretrained(\n    \"facebook/mask2former-swin-tiny-coco-instance\",\n    num_labels=3\n)\n\nprocessor = transformers.Mask2FormerImageProcessor.from_pretrained(\n    \"facebook/mask2former-swin-tiny-coco-instance\",\n    do_normalize=True,\n    do_rescale=False\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-28T06:21:42.181383Z","iopub.execute_input":"2025-11-28T06:21:42.181679Z","iopub.status.idle":"2025-11-28T06:22:00.95371Z","shell.execute_reply.started":"2025-11-28T06:21:42.181659Z","shell.execute_reply":"2025-11-28T06:22:00.952838Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# THIS CELL SHOULD BE RUN ONLY FOR THE FIRST TRAIN ITERATION\n\nmodel = transformers.Mask2FormerForUniversalSegmentation.from_pretrained(\n    \"facebook/mask2former-swin-tiny-coco-instance\",\n    config=config,\n    ignore_mismatched_sizes=True\n)\n\nmodel = model.to(device)\n\n# Stage-1:\n# We freeze the entire backbone (Swin), the pixel decoder and almost all\n# transformer decoder layers. Only the prediction heads remain trainable:\n#\n# - queries_embedder: learnable query embeddings (100 vectors)\n# - queries_features: initial feature vectors for each query\n# - mask_predictor: MLP that turns query features into mask embeddings\n# - class_predictor: linear layer predicting class labels for each query\n#\n# This stage only trains the top segmentation heads, while keeping all\n# pretrained dense feature extractors frozen.\nfor name, p in model.model.pixel_level_module.encoder.named_parameters():\n    p.requires_grad = False\n\nfor name, p in model.model.pixel_level_module.named_parameters():\n    p.requires_grad = False\n\nfor name, p in model.model.transformer_module.decoder.named_parameters():\n    if \"mask_predictor\" not in name:\n        p.requires_grad = False\n\nfor p in model.model.transformer_module.queries_embedder.parameters():\n    p.requires_grad = True\n\nfor p in model.model.transformer_module.queries_features.parameters():\n    p.requires_grad = True\n\nfor p in model.class_predictor.parameters():\n    p.requires_grad = True\n\nfor name, p in model.named_parameters():\n    if p.requires_grad:\n        print(name)\n\nn_epochs = 40\npatience = 10\noptimizer = torch.optim.AdamW(model.parameters(), lr=0.0001, weight_decay=0.00001)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-28T06:22:00.95517Z","iopub.execute_input":"2025-11-28T06:22:00.955717Z","iopub.status.idle":"2025-11-28T06:22:04.802476Z","shell.execute_reply.started":"2025-11-28T06:22:00.955694Z","shell.execute_reply":"2025-11-28T06:22:04.801641Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# helper for the processor\ndef inputs_to_device(_inputs):\n    return {\n        \"pixel_values\": _inputs[\"pixel_values\"].to(device).float(),\n        \"pixel_mask\": _inputs[\"pixel_mask\"].to(device),\n        \"mask_labels\": [m_l.to(device) for m_l in _inputs[\"mask_labels\"]],\n        \"class_labels\": [c_l.to(device) for c_l in _inputs[\"class_labels\"]],\n    }","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-28T06:22:04.803582Z","iopub.execute_input":"2025-11-28T06:22:04.803869Z","iopub.status.idle":"2025-11-28T06:22:04.808328Z","shell.execute_reply.started":"2025-11-28T06:22:04.80384Z","shell.execute_reply":"2025-11-28T06:22:04.807564Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def test_model_f(imgs, model, processor, device):\n    model.eval()\n\n    with torch.no_grad():\n        inputs = processor(images=[i for i in imgs], return_tensors=\"pt\")\n        inputs = {k: v.to(device) for k, v in inputs.items()}\n\n        outputs = model(**inputs)\n\n        target_sizes = [(imgs.shape[2], imgs.shape[3])] * imgs.size(0)\n        preds = processor.post_process_semantic_segmentation(outputs, target_sizes)\n\n        preds = torch.stack([p.to(device) for p in preds], dim=0)\n\n        return preds\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-28T06:22:04.809166Z","iopub.execute_input":"2025-11-28T06:22:04.809442Z","iopub.status.idle":"2025-11-28T06:22:04.831857Z","shell.execute_reply.started":"2025-11-28T06:22:04.809419Z","shell.execute_reply":"2025-11-28T06:22:04.831078Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# train\ndef train_one_epoch(train_loader, optimizer, model):\n    model.train()\n\n    running_loss, running_correct, total = 0.0, 0, 0\n    running_f1_macro = 0.0\n\n    autocast_ctx = torch.amp.autocast(device_type=\"cuda\")\n\n    scaler = torch.amp.GradScaler(\"cuda\")\n\n    pbar = tqdm(train_loader, desc=\"Train\", unit=\"batch\")\n\n    for imgs, masks in pbar:\n        imgs, masks = imgs.to(device), masks.to(device)\n\n        imgs = imgs.repeat(1, 3, 1, 1)\n\n        seg_maps = torch.argmax(masks, dim=1)\n        inputs = processor(\n            images=[i for i in imgs],\n            segmentation_maps=[s for s in seg_maps.cpu().numpy()],\n            return_tensors=\"pt\"\n        )\n\n        inputs = inputs_to_device(inputs)\n\n        # inputs = {k: v.to(device) for k, v in inputs.items()}\n\n        optimizer.zero_grad()\n\n        with autocast_ctx:\n            outputs = model(**inputs)\n            loss_total_per_batch = outputs.loss\n\n        if scaler is not None:\n            scaler.scale(loss_total_per_batch).backward()\n            scaler.step(optimizer)\n            scaler.update()\n\n        running_loss += loss_total_per_batch * imgs.size(0)\n        total += imgs.size(0)\n\n        with torch.no_grad():\n            target_sizes = [seg_maps.shape[1:]] * imgs.size(0)\n            preds = processor.post_process_semantic_segmentation(\n                outputs, target_sizes)\n            preds = torch.stack([p.to(device) for p in preds], dim=0)\n\n            preds_flattened = preds.view(-1)\n            targets_flattened = seg_maps.view(-1)\n\n            f1_batch_macro = multiclass_f1_score(preds_flattened, targets_flattened, num_classes=3, average=None)\n\n        running_f1_macro += f1_batch_macro * imgs.size(0)\n\n        avg_loss = running_loss / max(1, total)\n        avg_f1m = running_f1_macro / max(1, total)\n\n        pbar.set_postfix(\n            loss=f\"{avg_loss:.4f}\",\n            f1_per_class=\", \".join(f\"{num:.4f}\" for num in avg_f1m)\n        )\n    return model","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-27T20:27:58.785376Z","iopub.execute_input":"2025-11-27T20:27:58.785706Z","iopub.status.idle":"2025-11-27T20:27:58.804861Z","shell.execute_reply.started":"2025-11-27T20:27:58.785679Z","shell.execute_reply":"2025-11-27T20:27:58.804133Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# validation\ndef validate(val_loader, model):\n    model.eval()\n    running_loss, running_correct, running_2_classes, total = 0.0, 0, 0.0, 0\n    running_f1_macro = 0.0\n    autocast_ctx = torch.amp.autocast(device_type=\"cuda\")\n\n    pbar = tqdm(val_loader, desc=\"Val\", unit=\"batch\")\n\n    with torch.no_grad():\n        for imgs, masks in pbar:\n            visualize_transposed(imgs, masks)\n\n            imgs, masks = imgs.to(device), masks.to(device)\n\n\n            imgs = imgs.repeat(1, 3, 1, 1)\n\n\n            seg_maps = torch.argmax(masks, dim=1)\n\n            inputs = processor(\n                images=[i for i in imgs],\n                segmentation_maps=[s for s in seg_maps.cpu().numpy()],\n                return_tensors=\"pt\"\n            )\n            inputs = inputs_to_device(inputs)\n\n            with autocast_ctx:\n                outputs = model(**inputs)\n                loss_total_per_batch = outputs.loss\n\n            running_loss += loss_total_per_batch * imgs.size(0)\n            total += imgs.size(0)\n\n            with torch.no_grad():\n                target_sizes = [seg_maps.shape[1:]] * imgs.size(0)\n\n                preds = processor.post_process_semantic_segmentation(\n                outputs, target_sizes)\n                preds = torch.stack([p.to(device) for p in preds], dim=0)\n\n                preds_one_hot = torch.nn.functional.one_hot(preds, num_classes=3)\n                preds_one_hot = preds_one_hot.permute(0, 3, 1, 2).float()\n\n                visualize_transposed(imgs.cpu(), preds_one_hot.cpu())\n\n                preds_flattened = preds.view(-1)\n                targets_flattened = seg_maps.view(-1)\n\n                f1_batch_macro = multiclass_f1_score(preds_flattened, targets_flattened, num_classes=3, average=None)\n                f1_batch_2_classes = multiclass_f1_score(preds_flattened, targets_flattened, num_classes=3, average=\"macro\")\n            running_f1_macro += f1_batch_macro * imgs.size(0)\n\n            avg_loss = running_loss / max(1, total)\n            avg_f1m = running_f1_macro / max(1, total)\n\n            pbar.set_postfix(\n                loss=f\"{avg_loss:.4f}\",\n                f1_per_class=\", \".join(f\"{num:.4f}\" for num in avg_f1m),\n                f1_macro=f\"{f1_batch_2_classes:.4f}\"\n            )\n    return avg_loss","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-27T20:27:58.805757Z","iopub.execute_input":"2025-11-27T20:27:58.806578Z","iopub.status.idle":"2025-11-27T20:27:58.828484Z","shell.execute_reply.started":"2025-11-27T20:27:58.806557Z","shell.execute_reply":"2025-11-27T20:27:58.827745Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# training pipeline\ndef fit(_model, num_epochs, _patience):\n    best_loss = float(\"inf\")\n    fit_patience = 0\n    _best_model = None\n    for epoch in range(num_epochs):\n        model_trained = train_one_epoch(train_loader, optimizer, _model)\n        torch.cuda.empty_cache()\n        val_loss = validate(val_loader, model_trained)\n        if val_loss < best_loss:\n            best_loss = val_loss\n            fit_patience = 0\n            torch.save(model_trained, f\"best_model.pth\")\n        else:\n            fit_patience += 1\n        if fit_patience > _patience:\n            break\n        print(fit_patience)\n    return\n\nfit(model, n_epochs, patience)\ndel model","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-27T20:27:58.829238Z","iopub.execute_input":"2025-11-27T20:27:58.829489Z","iopub.status.idle":"2025-11-27T21:36:34.58622Z","shell.execute_reply.started":"2025-11-27T20:27:58.829469Z","shell.execute_reply":"2025-11-27T21:36:34.585368Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# THIS CELL SHOULD BE RUN ONLY FOR THE SECOND TRAIN ITERATION\nmodel = torch.load(f\"best_model.pth\", weights_only=False)\n\nmodel = model.to(device)\n\n# Stage-2:\n# We train everything except the backbone\nfor name, p in model.named_parameters():\n    p.requires_grad = True\n\nfor name, p in model.model.pixel_level_module.encoder.named_parameters():\n    p.requires_grad = False\n\nfor name, p in model.named_parameters():\n    if p.requires_grad:\n        print(name)\n\nn_epochs = 60\npatience = 10\noptimizer = torch.optim.AdamW(model.parameters(), lr=0.00001, weight_decay=0.0000001)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-27T21:36:34.588866Z","iopub.execute_input":"2025-11-27T21:36:34.589116Z","iopub.status.idle":"2025-11-27T21:36:34.824801Z","shell.execute_reply.started":"2025-11-27T21:36:34.589098Z","shell.execute_reply":"2025-11-27T21:36:34.824032Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# training pipeline\ndef fit(_model, num_epochs, _patience):\n    best_loss = float(\"inf\")\n    fit_patience = 0\n    _best_model = None\n    for epoch in range(num_epochs):\n        model_trained = train_one_epoch(train_loader, optimizer, _model)\n        torch.cuda.empty_cache()\n        val_loss = validate(val_loader, model_trained)\n        if val_loss < best_loss:\n            best_loss = val_loss\n            fit_patience = 0\n            torch.save(model_trained, f\"best_model.pth\")\n        else:\n            fit_patience += 1\n        if fit_patience > _patience:\n            break\n        print(fit_patience)\n    return\n\nfit(model, n_epochs, patience)\ndel model","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-27T21:36:34.825874Z","iopub.execute_input":"2025-11-27T21:36:34.826252Z","iopub.status.idle":"2025-11-27T23:02:13.516203Z","shell.execute_reply.started":"2025-11-27T21:36:34.826231Z","shell.execute_reply":"2025-11-27T23:02:13.5154Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# THIS CELL SHOULD BE RUN ONLY FOR THE LAST TRAIN ITERATION\nmodel = torch.load(f\"best_model.pth\", weights_only=False)\n\nmodel = model.to(device)\n\n# Stage-3:\n# We train everything\nfor name, p in model.named_parameters():\n    p.requires_grad = True\n\nfor name, p in model.named_parameters():\n    if p.requires_grad:\n        print(name)\n\nn_epochs = 60\npatience = 20\noptimizer = torch.optim.AdamW(model.parameters(), lr=0.000001, weight_decay=0.000000001)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-27T23:02:13.517188Z","iopub.execute_input":"2025-11-27T23:02:13.517465Z","iopub.status.idle":"2025-11-27T23:02:13.746067Z","shell.execute_reply.started":"2025-11-27T23:02:13.51744Z","shell.execute_reply":"2025-11-27T23:02:13.745169Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# training pipeline\ndef fit(_model, num_epochs, _patience):\n    best_loss = float(\"inf\")\n    fit_patience = 0\n    _best_model = None\n    for epoch in range(num_epochs):\n        model_trained = train_one_epoch(train_loader, optimizer, _model)\n        torch.cuda.empty_cache()\n        val_loss = validate(val_loader, model_trained)\n        if val_loss < best_loss:\n            best_loss = val_loss\n            fit_patience = 0\n            torch.save(model_trained, f\"best_model.pth\")\n        else:\n            fit_patience += 1\n        if fit_patience > _patience:\n            break\n        print(fit_patience)\n    return\n\nfit(model, n_epochs, patience)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-27T23:02:13.748166Z","iopub.execute_input":"2025-11-27T23:02:13.748489Z","iopub.status.idle":"2025-11-27T23:02:16.576755Z","shell.execute_reply.started":"2025-11-27T23:02:13.74846Z","shell.execute_reply":"2025-11-27T23:02:16.575481Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def test_model():\n    # model = torch.load(f\"best_model.pth\", weights_only=False).to(device)\n    test = torch.tensor(test_images_medseg).permute((0, 3, 1, 2)).to(device)\n    model.eval()\n    test = test.repeat(1, 3, 1, 1)\n    test_preds = test_model_f(test, model, processor, device)\n    test_preds_one_hot = torch.nn.functional.one_hot(test_preds, num_classes=3)\n    test_preds_one_hot = test_preds_one_hot.permute((0, 3, 1, 2))\n    visualize_transposed(test.cpu(), test_preds_one_hot.cpu(), num_samples=10)\n    test_preds_one_hot = test_preds_one_hot[:, 0:2, :, :]\n    test_preds_mask = test_preds_one_hot.permute((0, 2, 3, 1))\n    pd.DataFrame(test_preds_mask.reshape(-1).int().cpu().numpy(), columns=['Predicted']).reset_index(names=\"Id\").to_csv(\"submission.csv\", index=False)\n    return\n\n\ntest_model()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-28T06:22:15.457453Z","iopub.execute_input":"2025-11-28T06:22:15.457736Z","iopub.status.idle":"2025-11-28T06:22:25.453344Z","shell.execute_reply.started":"2025-11-28T06:22:15.457717Z","shell.execute_reply":"2025-11-28T06:22:25.45252Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}