{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.12.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[],"dockerImageVersionId":28755,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Imports","metadata":{}},{"cell_type":"code","source":"!pip uninstall -y albumentations -y\n!pip install --no-deps --upgrade albumentationsx albucore numkong\n!pip uninstall torch torchvision torch-tensorrt -y\n!pip install torch torchvision torch-tensorrt ","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-10T06:50:22.289092Z","iopub.execute_input":"2026-10-10T06:50:22.289457Z","iopub.status.idle":"2026-10-10T06:56:37.59294Z","shell.execute_reply.started":"2026-10-10T06:50:22.289421Z","shell.execute_reply":"2026-10-10T06:56:37.591949Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import zipfile\nimport shutil\nimport os\nfrom pathlib import Path\nfrom IPython.display import FileLink\n\nfrom PIL import Image\nimport cv2\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\n\nimport numpy as np\nimport matplotlib.pyplot as plt\n\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nimport torchvision\nfrom torchvision.transforms.v2 import functional as TF\nfrom torchvision.utils import draw_segmentation_masks\nimport torchvision.models as models\nfrom torchinfo import summary\nimport torch.multiprocessing as mp\n\nfrom tqdm import tqdm\n\nimport random\n\nimport matplotlib.pyplot as plt\nimport seaborn as sns\nimport pandas as pd\nimport time\nimport optuna\nimport torch_tensorrt\nfrom torch_tensorrt import MutableTorchTensorRTModule\nimport logging\n\ntorch_tensorrt.logging.set_level(logging.INFO)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-10T06:56:37.595177Z","iopub.execute_input":"2026-10-10T06:56:37.595579Z","iopub.status.idle":"2026-10-10T06:57:03.874835Z","shell.execute_reply.started":"2026-10-10T06:56:37.595547Z","shell.execute_reply":"2026-10-10T06:57:03.873913Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nprint(torch.__version__)\n\nimport tensorrt\nprint(tensorrt.__version__)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-10T06:57:03.875966Z","iopub.execute_input":"2026-10-10T06:57:03.87671Z","iopub.status.idle":"2026-10-10T06:57:03.882611Z","shell.execute_reply.started":"2026-10-10T06:57:03.876681Z","shell.execute_reply":"2026-10-10T06:57:03.881323Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"SEED = 42\nHEIGHT, WIDTH = (1280, 1920)\nSIZE = (HEIGHT, WIDTH)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-10T06:57:03.884074Z","iopub.execute_input":"2026-10-10T06:57:03.884468Z","iopub.status.idle":"2026-10-10T06:57:04.027433Z","shell.execute_reply.started":"2026-10-10T06:57:03.884437Z","shell.execute_reply":"2026-10-10T06:57:04.025822Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Data","metadata":{}},{"cell_type":"code","source":"DATASET_DIR = '/kaggle/input/competitions/carvana-image-masking-challenge/'\nWORKING_DIR = '/kaggle/working/'\ntrain_dir = WORKING_DIR + 'train_hq/'\nval_dir = WORKING_DIR + 'val/'\ntrain_masks_dir = WORKING_DIR + 'train_masks/'\nval_masks_dir = WORKING_DIR + 'val_masks/'","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-10T06:57:04.028694Z","iopub.execute_input":"2026-10-10T06:57:04.029074Z","iopub.status.idle":"2026-10-10T06:57:07.73509Z","shell.execute_reply.started":"2026-10-10T06:57:04.029036Z","shell.execute_reply":"2026-10-10T06:57:07.734033Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if len(os.listdir(WORKING_DIR)) <= 1:\n    with zipfile.ZipFile(DATASET_DIR + 'train_hq.zip', 'r') as zip_file:\n        zip_file.extractall(WORKING_DIR)\n    \n    with zipfile.ZipFile(DATASET_DIR + 'train_masks.zip', 'r') as zip_file:\n        zip_file.extractall(WORKING_DIR)\n    \n    print(\n        len(os.listdir(WORKING_DIR + 'train_hq')),\n        len(os.listdir(WORKING_DIR + 'train_masks'))\n    )\n    \n    os.mkdir(val_dir)\n    for file in sorted(os.listdir(train_dir))[4600:]:\n      shutil.move(train_dir + file, val_dir)\n    \n    os.mkdir(val_masks_dir)\n    for file in sorted(os.listdir(train_masks_dir))[4600:]:\n      shutil.move(train_masks_dir + file, val_masks_dir)\n    \n    os.mkdir(WORKING_DIR + 'saved_images')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-10T06:57:07.736286Z","iopub.execute_input":"2026-10-10T06:57:07.737455Z","iopub.status.idle":"2026-10-10T06:57:27.845954Z","shell.execute_reply.started":"2026-10-10T06:57:07.737428Z","shell.execute_reply":"2026-10-10T06:57:27.845032Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%writefile carvana_dataset.py\n\nimport torchvision\nfrom torchvision.transforms.v2 import functional as TF\n\nfrom torch.utils.data import Dataset\nfrom PIL import Image\nimport numpy as np\nfrom tqdm import tqdm\nimport os\n\nclass PreloadedCarvanaDataset(Dataset):\n    def __init__(self, image_dir, mask_dir, transform=None):\n        self.image_dir = image_dir\n        self.mask_dir = mask_dir\n        self.transform = transform\n        self.images = os.listdir(image_dir)\n\n        # if dataset fits in memory fully\n        self.targets = [] \n\n        for i in tqdm(range(len(self))):\n            self.targets.append(self.process_path(i))\n\n    def __len__(self):\n        return len(self.images)\n\n    def process_path(self, idx):\n        img_path = os.path.join(self.image_dir, self.images[idx])\n        mask_path = os.path.join(self.mask_dir, self.images[idx].replace('.jpg', '_mask.gif'))\n        image = np.array(Image.open(img_path).convert('RGB'))\n        mask = np.array(Image.open(mask_path).convert('L'), dtype=np.float32)\n        mask[mask == 255.0] = 1.0\n\n        target = {\"image\": image, \"mask\": mask}\n        return target\n    \n    def __getitem__(self, idx):\n        target = self.targets[idx]\n\n        if self.transform is not None:\n            augmentations = self.transform(image=target['image'], mask=target['mask'])\n            image = augmentations['image']\n            mask = augmentations['mask']\n\n        return image, mask\n\nclass CarvanaMockDataset(Dataset):\n    def __init__(self, image_dir, mask_dir, transform=None):\n        self.image_dir = image_dir\n        self.mask_dir = mask_dir\n        self.transform = transform\n        self.images = os.listdir(image_dir)\n\n    def __len__(self):\n        return len(self.images)\n\n    def __getitem__(self, index):\n        img_path = os.path.join(self.image_dir, self.images[index])\n        mask_path = os.path.join(self.mask_dir, self.images[index].replace('.jpg', '_mask.gif'))\n        image = np.array(Image.open(img_path).convert('RGB'))\n        mask = np.array(Image.open(mask_path).convert('L'), dtype=np.float32)\n        mask[mask == 255.0] = 1.0\n\n        if self.transform is not None:\n            augmentations = self.transform(image=image, mask=mask)\n            transformed_image = augmentations['image']\n            transformed_mask = augmentations['mask']\n\n        return image, transformed_image, mask, transformed_mask\n\nclass CarvanaTestDataset(Dataset):\n    def __init__(self, image_dir, transform=None):\n        self.image_dir = image_dir\n        self.images = sorted([f for f in os.listdir(image_dir) if f.endswith(('.jpg', '.jpeg'))])\n        self.transform = transform\n\n    def __len__(self):\n        return len(self.images)\n\n    def __getitem__(self, index):\n        img_name = self.images[index]\n        img_path = os.path.join(self.image_dir, img_name)\n        \n        image = np.array(Image.open(img_path).convert('RGB'))\n        if self.transform is not None:\n            augmentations = self.transform(image=image)\n            image = augmentations['image']\n        \n        return img_name, image\n\nclass CarvanaDataset(Dataset):\n    def __init__(self, image_dir, mask_dir, transform=None):\n        self.image_dir = image_dir\n        self.mask_dir = mask_dir\n        self.transform = transform\n        self.images = os.listdir(image_dir)\n\n    def __len__(self):\n        return len(self.images)\n\n    def __getitem__(self, index):\n        img_path = os.path.join(self.image_dir, self.images[index])\n        mask_path = os.path.join(self.mask_dir, self.images[index].replace('.jpg', '_mask.gif'))\n        image = np.array(Image.open(img_path).convert('RGB'))\n        mask = np.array(Image.open(mask_path).convert('L'), dtype=np.float32)\n        mask[mask == 255.0] = 1.0\n\n        if self.transform is not None:\n            augmentations = self.transform(image=image, mask=mask)\n            image = augmentations['image']\n            mask = augmentations['mask']\n\n        return image, mask","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-10T06:57:27.848446Z","iopub.execute_input":"2026-10-10T06:57:27.848771Z","iopub.status.idle":"2026-10-10T06:57:27.857315Z","shell.execute_reply.started":"2026-10-10T06:57:27.848741Z","shell.execute_reply":"2026-10-10T06:57:27.856468Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from carvana_dataset import CarvanaDataset, CarvanaMockDataset, CarvanaTestDataset, PreloadedCarvanaDataset","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-10T06:57:27.858593Z","iopub.execute_input":"2026-10-10T06:57:27.859239Z","iopub.status.idle":"2026-10-10T06:57:27.888617Z","shell.execute_reply.started":"2026-10-10T06:57:27.859204Z","shell.execute_reply":"2026-10-10T06:57:27.887842Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"mock_transform = A.Compose(\n    [\n        ToTensorV2(),\n    ],\n    strict=True,\n    seed=SEED,\n)\n\nmock_dataset = CarvanaDataset(\n    train_dir, \n    train_masks_dir,\n    mock_transform\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-10T06:57:27.889815Z","iopub.execute_input":"2026-10-10T06:57:27.890119Z","iopub.status.idle":"2026-10-10T06:57:27.914144Z","shell.execute_reply.started":"2026-10-10T06:57:27.890095Z","shell.execute_reply":"2026-10-10T06:57:27.91321Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def show(imgs):\n    if not isinstance(imgs, list):\n        imgs = [imgs]\n    fig, axs = plt.subplots(ncols=len(imgs), squeeze=False, figsize=(12, 6))\n    for i, img in enumerate(imgs):\n        img = img.detach()\n        img = TF.to_pil_image(img)\n        axs[0, i].imshow(np.asarray(img))\n        axs[0, i].set(xticklabels=[], yticklabels=[], xticks=[], yticks=[])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-10T06:57:27.915535Z","iopub.execute_input":"2026-10-10T06:57:27.915899Z","iopub.status.idle":"2026-10-10T06:57:27.923874Z","shell.execute_reply.started":"2026-10-10T06:57:27.915858Z","shell.execute_reply":"2026-10-10T06:57:27.923009Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"for i in range(4):\n    image, masks = mock_dataset[i]\n    result = draw_segmentation_masks(image, masks.to(bool), colors=['cyan'] * len(masks))\n\n    show([image, result])\n    plt.tight_layout()\n    plt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-10T06:57:27.924985Z","iopub.execute_input":"2026-10-10T06:57:27.925851Z","iopub.status.idle":"2026-10-10T06:57:30.925313Z","shell.execute_reply.started":"2026-10-10T06:57:27.925816Z","shell.execute_reply":"2026-10-10T06:57:30.924523Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%writefile recolor.py\n\nimport cv2\nimport numpy as np\nfrom sklearn.cluster import MiniBatchKMeans\n\ndef recolor(\n    image: np.ndarray,\n    mask: np.ndarray,\n    tolerance: float = 0.35,\n    n_clusters: int = 3,\n    hue_shift: float = 0.0,\n    neutral_sat: float = 130.0,\n    lightness_weight: float = 0.2,\n    neutral_threshold: float = 18.0,\n    preserve_highlights: bool = True,\n) -> np.ndarray:\n    car_mask = (mask > 0) if mask.ndim == 2 else (mask[..., 0] > 0)\n    if np.count_nonzero(car_mask) < 50:\n        return image.copy()\n\n    lab = cv2.cvtColor(image, cv2.COLOR_RGB2LAB).astype(np.float32)\n    masked_lab = lab[car_mask]\n\n    sample_size = min(len(masked_lab), 10000)\n    indices = np.random.choice(len(masked_lab), sample_size, replace=False)\n    samples = masked_lab[indices]\n\n    feat_a = (samples[:, 1] - 128.0) / 128.0\n    feat_b = (samples[:, 2] - 128.0) / 128.0\n    feat_l = (samples[:, 0] / 255.0) * lightness_weight\n    features = np.column_stack([feat_a, feat_b, feat_l])\n\n    k = min(n_clusters, len(samples))\n    kmeans = MiniBatchKMeans(n_clusters=k, batch_size=1024, n_init=3, random_state=42)\n    labels = kmeans.fit_predict(features)\n\n    counts = np.bincount(labels, minlength=k)\n    chosen_idx = np.random.choice(k, p=counts / counts.sum())\n    center_lab = samples[labels == chosen_idx].mean(axis=0)\n\n    chroma = np.hypot(center_lab[1] - 128.0, center_lab[2] - 128.0)\n    is_neutral = chroma < neutral_threshold\n\n    da = (lab[:, :, 1] - center_lab[1]) / 128.0\n    db = (lab[:, :, 2] - center_lab[2]) / 128.0\n    dl = (lab[:, :, 0] - center_lab[0]) / 255.0\n\n    l_scale = 0.85 if is_neutral else lightness_weight\n    dist = np.sqrt(da**2 + db**2 + (l_scale * dl)**2)\n\n    u = np.clip(1.0 - (dist / max(tolerance, 1e-4)), 0.0, 1.0)\n    selection = u * u * (3.0 - 2.0 * u)\n    selection *= car_mask.astype(np.float32)\n\n    hsv = cv2.cvtColor(image, cv2.COLOR_RGB2HSV).astype(np.float32)\n    h_channel = hsv[:, :, 0]\n    s_channel = hsv[:, :, 1]\n    v_channel = hsv[:, :, 2]\n\n    if is_neutral:\n        target_h = (hue_shift % 360.0) / 2.0\n        rec_h = np.full_like(h_channel, target_h)\n        rec_s = np.full_like(s_channel, neutral_sat)\n        rec_v = v_channel.copy()\n    else:\n        delta_h = (hue_shift / 2.0) % 180.0\n        rec_h = (h_channel + delta_h) % 180.0\n        rec_s = s_channel.copy()\n        rec_v = v_channel.copy()\n\n    if preserve_highlights:\n        highlight_taper = np.clip((255.0 - v_channel) / 35.0, 0.0, 1.0)\n        shadow_taper = np.clip(v_channel / 25.0, 0.0, 1.0)\n        selection *= (highlight_taper * shadow_taper)\n\n    rec_hsv = np.stack([rec_h, rec_s, rec_v], axis=-1)\n    rec_rgb = cv2.cvtColor(rec_hsv.astype(np.uint8), cv2.COLOR_HSV2RGB).astype(np.float32)\n\n    weight = selection[:, :, np.newaxis]\n    final_rgb = (1.0 - weight) * image.astype(np.float32) + weight * rec_rgb\n    return np.clip(final_rgb, 0, 255).astype(np.uint8)\n\nfrom albumentations.core.transforms_interface import DualTransform\nfrom albumentations.core.transform_params import SampledParams\n\nclass ColorReplace(DualTransform):\n    def __init__(\n        self,\n        tolerance: float = 0.35,\n        n_clusters: int = 3,\n        hue_shift_limit: tuple = (-180, 180),\n        lightness_weight: float = 0.2,\n        neutral_saturation_range: tuple = (100, 160),\n        neutral_threshold: float = 18.0,\n        preserve_highlights: bool = True,\n        p: float = 0.75,\n    ):\n        super().__init__(p=p)\n        self.tolerance = tolerance\n        self.n_clusters = n_clusters\n        self.hue_shift_limit = hue_shift_limit\n        self.lightness_weight = lightness_weight\n        self.neutral_saturation_range = neutral_saturation_range\n        self.neutral_threshold = neutral_threshold\n        self.preserve_highlights = preserve_highlights\n\n    @property\n    def targets_as_params(self):\n        return [\"mask\"]\n\n    def sample_parameters(self, params, data, targets, sampling):\n        mask = data.get(\"mask\", None)\n        hue_shift = sampling.py_random.uniform(self.hue_shift_limit[0], self.hue_shift_limit[1])\n        neutral_sat = sampling.py_random.uniform(\n            self.neutral_saturation_range[0], self.neutral_saturation_range[1]\n        )\n        return SampledParams(params={\n            \"replace_mask\": mask,\n            \"hue_shift\": hue_shift,\n            \"neutral_sat\": neutral_sat,\n        })\n\n    def apply(\n        self,\n        img: np.ndarray,\n        replace_mask: np.ndarray = None,\n        hue_shift: float = 0.0,\n        neutral_sat: float = 130.0,\n        **params,\n    ) -> np.ndarray:\n        if replace_mask is None:\n            return img\n\n        return recolor(\n            image=img,\n            mask=replace_mask,\n            tolerance=self.tolerance,\n            n_clusters=self.n_clusters,\n            hue_shift=hue_shift,\n            neutral_sat=neutral_sat,\n            lightness_weight=self.lightness_weight,\n            neutral_threshold=self.neutral_threshold,\n            preserve_highlights=self.preserve_highlights,\n        )\n\n    def apply_to_mask(self, mask: np.ndarray, **params) -> np.ndarray:\n        return mask","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-10T06:57:30.9263Z","iopub.execute_input":"2026-10-10T06:57:30.926625Z","iopub.status.idle":"2026-10-10T06:57:30.933949Z","shell.execute_reply.started":"2026-10-10T06:57:30.926599Z","shell.execute_reply":"2026-10-10T06:57:30.933134Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from recolor import recolor, ColorReplace","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-10T06:57:30.935135Z","iopub.execute_input":"2026-10-10T06:57:30.935485Z","iopub.status.idle":"2026-10-10T06:57:31.947244Z","shell.execute_reply.started":"2026-10-10T06:57:30.935461Z","shell.execute_reply":"2026-10-10T06:57:31.946592Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"simple_augmentation_pipeline = A.Compose(\n    [\n        A.Pad(padding=(2, 0, 0, 0), border_mode=cv2.BORDER_CONSTANT, fill=0, fill_mask=0), # + 2 pixels on left\n        \n        # Densification on zoom, small movement, small rotation to simulate different positions of a car on the platform\n        A.HorizontalFlip(p=0.5),\n        A.Affine(scale=(0.9, 1.1), rotate=(-5, 5), translate_px=(100, 100), balanced_scale=True, p=0.5),\n\n        # Change color per instance transform\n        ColorReplace(tolerance=0.45, n_clusters=2, hue_shift_limit=(-180, 180), lightness_weight=0.3,  p=0.5),\n        \n        # Exposure\n        A.RandomGamma(gamma_range=(80, 120), p=0.2),\n        A.RandomBrightnessContrast(brightness_range=(-0.15, 0.15), contrast_range=(-0.15, 0.15), p=0.2),\n        A.RandomToneCurve(scale=0.1, p=0.1),\n\n        A.OneOf([\n            A.LensFlare(intensity_range=(0.3, 0.6), p=1.0),\n            A.RandomShadow(p=1.0)\n        ], p=0.2),\n        A.CLAHE(clip_range=(1, 4), tile_grid_size=(8, 8), p=0.05),\n        \n        A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),\n        ToTensorV2(),\n    ],\n    strict=True,\n    seed=SEED,\n)\n\nval_pipeline = A.Compose(\n    [\n        A.Pad(padding=(2, 0, 0, 0), border_mode=cv2.BORDER_CONSTANT, fill=0, fill_mask=0),\n        A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),\n        ToTensorV2(),\n    ],\n    strict=True,\n    seed=SEED,\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-10T06:57:31.948354Z","iopub.execute_input":"2026-10-10T06:57:31.948921Z","iopub.status.idle":"2026-10-10T06:57:31.967951Z","shell.execute_reply.started":"2026-10-10T06:57:31.948896Z","shell.execute_reply":"2026-10-10T06:57:31.967185Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_transform = simple_augmentation_pipeline\nval_transform = val_pipeline\n\ntrain_dataset = CarvanaMockDataset(\n    image_dir=train_dir,\n    mask_dir=train_masks_dir,\n    transform=train_transform\n)\nval_ds = CarvanaMockDataset(\n    image_dir=val_dir,\n    mask_dir=val_masks_dir,\n    transform=val_transform\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-10T06:57:31.969127Z","iopub.execute_input":"2026-10-10T06:57:31.969508Z","iopub.status.idle":"2026-10-10T06:57:31.976751Z","shell.execute_reply.started":"2026-10-10T06:57:31.969474Z","shell.execute_reply":"2026-10-10T06:57:31.97603Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"for i in range(8):\n    image, tf_image, masks, tf_masks = train_dataset[i]\n    image = torch.from_numpy(image).permute(2, 0, 1)\n    masks = torch.from_numpy(masks)\n    tf_image = tf_image / 4.8 + 0.45\n    \n    raw_result = draw_segmentation_masks(image, masks.to(bool), colors=['cyan'] * len(masks))\n    tf_result = draw_segmentation_masks(tf_image, tf_masks.to(bool), colors=['cyan'] * len(masks))\n    show([image, raw_result, tf_image, tf_result])\n    plt.tight_layout()\n    plt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-10T06:57:31.977787Z","iopub.execute_input":"2026-10-10T06:57:31.978164Z","iopub.status.idle":"2026-10-10T06:57:41.223788Z","shell.execute_reply.started":"2026-10-10T06:57:31.978137Z","shell.execute_reply":"2026-10-10T06:57:41.222951Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_dataset = CarvanaDataset(\n    image_dir=train_dir,\n    mask_dir=train_masks_dir,\n    transform=train_transform\n)\nval_ds = CarvanaDataset(\n    image_dir=val_dir,\n    mask_dir=val_masks_dir,\n    transform=val_transform\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-10T06:57:41.224887Z","iopub.execute_input":"2026-10-10T06:57:41.225296Z","iopub.status.idle":"2026-10-10T06:57:41.234046Z","shell.execute_reply.started":"2026-10-10T06:57:41.225237Z","shell.execute_reply":"2026-10-10T06:57:41.2331Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# U-Net","metadata":{}},{"cell_type":"code","source":"%%writefile unet.py\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torchvision\nimport torchvision.models as models\n\nclass ConvRelu(nn.Module):\n    def __init__(self, in_channels, out_channels):\n        super(ConvRelu, self).__init__()\n\n        self.block = nn.Sequential(\n            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1, bias=False),\n            nn.BatchNorm2d(out_channels),\n            nn.ReLU(inplace=True)\n        )\n\n    def forward(self, x):\n        return self.block(x)\n\n\nclass DecoderBlock(nn.Module):\n    def __init__(self, in_channels, middle_channels, out_channels, is_deconv=True):\n        super(DecoderBlock, self).__init__()\n        self.in_channels = in_channels\n\n        if is_deconv:\n            self.block = nn.Sequential(\n                ConvRelu(in_channels, middle_channels),\n                nn.ConvTranspose2d(middle_channels, out_channels, kernel_size=3, stride=2,\n                                   padding=1, output_padding=1, bias=False),\n                nn.BatchNorm2d(out_channels),\n                nn.ReLU(inplace=True)\n            )\n        else:\n            self.block = nn.Sequential(\n                nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True),\n                ConvRelu(in_channels, middle_channels),\n                ConvRelu(middle_channels, out_channels),\n            )\n\n    def forward(self, x):\n        return self.block(x)\n\nclass UnetVgg16Bn(nn.Module):\n    def __init__(self, n_classes=1, num_filters=64):\n        super(UnetVgg16Bn, self).__init__()\n\n        try:\n            vgg = models.vgg16_bn(weights=models.VGG16_BN_Weights.DEFAULT)\n        except AttributeError:\n            vgg = models.vgg16_bn(pretrained=True)\n            \n        encoder = vgg.features\n\n        self.block1 = encoder[0:6]\n        self.block2 = encoder[7:13]\n        self.block3 = encoder[14:23]\n        self.block4 = encoder[24:33]\n        self.block5 = encoder[34:43]\n        \n        self.pool = nn.MaxPool2d(2, 2)\n\n        self.center = DecoderBlock(512, num_filters * 16, num_filters * 8)\n\n        self.dec5 = DecoderBlock(512 + 512, num_filters * 8, num_filters * 8)\n        self.dec4 = DecoderBlock(512 + 512, num_filters * 8, num_filters * 4)\n        self.dec3 = DecoderBlock(256 + 256, num_filters * 4, num_filters * 2)\n        self.dec2 = DecoderBlock(128 + 128, num_filters * 2, num_filters)\n        self.dec1 = ConvRelu(64 + 64, num_filters)\n\n        self.final = nn.Conv2d(num_filters, n_classes, kernel_size=1)\n\n    def forward(self, x):\n        conv1 = self.block1(x)\n        conv2 = self.block2(self.pool(conv1))\n        conv3 = self.block3(self.pool(conv2))\n        conv4 = self.block4(self.pool(conv3))\n        conv5 = self.block5(self.pool(conv4))\n\n        center = self.center(self.pool(conv5))\n\n        dec5 = self.dec5(torch.cat([center, conv5], 1))\n        dec4 = self.dec4(torch.cat([dec5, conv4], 1))\n        dec3 = self.dec3(torch.cat([dec4, conv3], 1))\n        dec2 = self.dec2(torch.cat([dec3, conv2], 1))\n        dec1 = self.dec1(torch.cat([dec2, conv1], 1))\n        \n        return self.final(dec1)\n\nclass UnetVgg11Bn(nn.Module):\n\n  def __init__(self, n_classes=1, num_filters=64):\n    super(UnetVgg11Bn, self).__init__()\n\n    try:\n      vgg = models.vgg11_bn(weights=models.VGG11_BN_Weights.DEFAULT)\n    except AttributeError:\n      vgg = models.vgg11_bn(pretrained=True)\n\n    encoder = vgg.features\n\n    self.block1 = encoder[0:3]\n    self.block2 = encoder[4:7]\n    self.block3 = encoder[8:14]\n    self.block4 = encoder[15:21]\n    self.block5 = encoder[22:28]\n\n    self.pool = nn.MaxPool2d(2, 2)\n\n    # Channel counts match VGG16 exactly (512, 512, 256, 128, 64)\n    self.center = DecoderBlock(512, num_filters * 16, num_filters * 8)\n    self.dec5 = DecoderBlock(512 + 512, num_filters * 8, num_filters * 8)\n    self.dec4 = DecoderBlock(512 + 512, num_filters * 8, num_filters * 4)\n    self.dec3 = DecoderBlock(256 + 256, num_filters * 4, num_filters * 2)\n    self.dec2 = DecoderBlock(128 + 128, num_filters * 2, num_filters)\n    self.dec1 = ConvRelu(64 + 64, num_filters)\n\n    self.final = nn.Conv2d(num_filters, n_classes, kernel_size=1)\n\n  def forward(self, x):\n    conv1 = self.block1(x)\n    conv2 = self.block2(self.pool(conv1))\n    conv3 = self.block3(self.pool(conv2))\n    conv4 = self.block4(self.pool(conv3))\n    conv5 = self.block5(self.pool(conv4))\n\n    center = self.center(self.pool(conv5))  # (B, 512, H/32, W/32)\n\n    dec5 = self.dec5(torch.cat([center, conv5], 1))\n    dec4 = self.dec4(torch.cat([dec5, conv4], 1))\n    dec3 = self.dec3(torch.cat([dec4, conv3], 1))\n    dec2 = self.dec2(torch.cat([dec3, conv2], 1))\n    dec1 = self.dec1(torch.cat([dec2, conv1], 1))\n\n    return self.final(dec1)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-10T06:57:41.235171Z","iopub.execute_input":"2026-10-10T06:57:41.235457Z","iopub.status.idle":"2026-10-10T06:57:41.248679Z","shell.execute_reply.started":"2026-10-10T06:57:41.235433Z","shell.execute_reply":"2026-10-10T06:57:41.248014Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%writefile unet_no_bn.py\nimport torch\nimport torch.nn as nn\nfrom torchvision import models\n\n\nclass ConvRelu(nn.Module):\n    def __init__(self, in_channels: int, out_channels: int):\n        super().__init__()\n        self.conv = nn.Sequential(\n            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),\n            nn.ReLU(inplace=True),\n        )\n\n    def forward(self, x):\n        return self.conv(x)\n\n\nclass DecoderBlock(nn.Module):\n    def __init__(self, in_channels: int, middle_channels: int, out_channels: int):\n        super().__init__()\n        self.block = nn.Sequential(\n            nn.Conv2d(in_channels, middle_channels, kernel_size=3, padding=1),\n            nn.ReLU(inplace=True),\n            nn.Conv2d(middle_channels, out_channels, kernel_size=3, padding=1),\n            nn.ReLU(inplace=True),\n        )\n\n    def forward(self, x):\n        return self.block(x)\n\n\nclass UnetVgg11(nn.Module):\n    def __init__(self, n_classes: int = 1, num_filters: int = 64):\n        super().__init__()\n\n        try:\n            vgg = models.vgg11(weights=models.VGG11_Weights.DEFAULT)\n        except (AttributeError, TypeError):\n            vgg = models.vgg11(pretrained=True)\n\n        encoder = vgg.features\n\n        # VGG11 without BatchNorm (slices exclude MaxPool layers at 2, 5, 10, 15, 20)\n        self.block1 = encoder[0:2]    # Conv(3->64), ReLU\n        self.block2 = encoder[3:5]    # Conv(64->128), ReLU\n        self.block3 = encoder[6:10]   # Conv(128->256), ReLU, Conv(256->256), ReLU\n        self.block4 = encoder[11:15]  # Conv(256->512), ReLU, Conv(512->512), ReLU\n        self.block5 = encoder[16:20]  # Conv(512->512), ReLU, Conv(512->512), ReLU\n\n        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)\n\n        self.center = DecoderBlock(512, num_filters * 16, num_filters * 8)\n        self.dec5 = DecoderBlock(512 + num_filters * 8, num_filters * 8, num_filters * 8)\n        self.dec4 = DecoderBlock(512 + num_filters * 8, num_filters * 8, num_filters * 4)\n        self.dec3 = DecoderBlock(256 + num_filters * 4, num_filters * 4, num_filters * 2)\n        self.dec2 = DecoderBlock(128 + num_filters * 2, num_filters * 2, num_filters)\n        self.dec1 = ConvRelu(64 + num_filters, num_filters)\n        self.final = nn.Conv2d(num_filters, n_classes, kernel_size=1)\n\n        # Kaiming normal initialization for decoder stages\n        self._init_decoder_weights()\n\n    def _init_decoder_weights(self):\n        decoder_modules = [self.center, self.dec5, self.dec4, self.dec3, self.dec2, self.dec1, self.final]\n        for m in decoder_modules:\n            for sub in m.modules():\n                if isinstance(sub, nn.Conv2d):\n                    nn.init.kaiming_normal_(sub.weight, mode='fan_out', nonlinearity='relu')\n                    if sub.bias is not None:\n                        nn.init.constant_(sub.bias, 0.0)\n\n    def forward(self, x):\n        conv1 = self.block1(x)\n        conv2 = self.block2(self.pool(conv1))\n        conv3 = self.block3(self.pool(conv2))\n        conv4 = self.block4(self.pool(conv3))\n        conv5 = self.block5(self.pool(conv4))\n\n        center = self.center(self.pool(conv5))\n\n        up5 = nn.functional.interpolate(center, size=conv5.shape[2:], mode=\"bilinear\", align_corners=True)\n        dec5 = self.dec5(torch.cat([up5, conv5], dim=1))\n\n        up4 = nn.functional.interpolate(dec5, size=conv4.shape[2:], mode=\"bilinear\", align_corners=True)\n        dec4 = self.dec4(torch.cat([up4, conv4], dim=1))\n\n        up3 = nn.functional.interpolate(dec4, size=conv3.shape[2:], mode=\"bilinear\", align_corners=True)\n        dec3 = self.dec3(torch.cat([up3, conv3], dim=1))\n\n        up2 = nn.functional.interpolate(dec3, size=conv2.shape[2:], mode=\"bilinear\", align_corners=True)\n        dec2 = self.dec2(torch.cat([up2, conv2], dim=1))\n\n        up1 = nn.functional.interpolate(dec2, size=conv1.shape[2:], mode=\"bilinear\", align_corners=True)\n        dec1 = self.dec1(torch.cat([up1, conv1], dim=1))\n\n        return self.final(dec1)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-10T06:57:41.249832Z","iopub.execute_input":"2026-10-10T06:57:41.250162Z","iopub.status.idle":"2026-10-10T06:57:41.274325Z","shell.execute_reply.started":"2026-10-10T06:57:41.250124Z","shell.execute_reply":"2026-10-10T06:57:41.273555Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from unet import UnetVgg11Bn, UnetVgg16Bn\nfrom unet_no_bn import UnetVgg11","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-10T06:57:41.275435Z","iopub.execute_input":"2026-10-10T06:57:41.275824Z","iopub.status.idle":"2026-10-10T06:57:41.295661Z","shell.execute_reply.started":"2026-10-10T06:57:41.275779Z","shell.execute_reply":"2026-10-10T06:57:41.294703Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"model = UnetVgg11(n_classes=1)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-10T06:57:41.296656Z","iopub.execute_input":"2026-10-10T06:57:41.29712Z","iopub.status.idle":"2026-10-10T06:57:45.259958Z","shell.execute_reply.started":"2026-10-10T06:57:41.297088Z","shell.execute_reply":"2026-10-10T06:57:45.259338Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"summary(model, input_size=(1, 3, 1920, 1280))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-10T06:57:45.264409Z","iopub.execute_input":"2026-10-10T06:57:45.264739Z","iopub.status.idle":"2026-10-10T06:57:46.305973Z","shell.execute_reply.started":"2026-10-10T06:57:45.264714Z","shell.execute_reply":"2026-10-10T06:57:46.305212Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Training","metadata":{}},{"cell_type":"code","source":"BATCH_SIZE = 1\nNUM_WORKERS = 2\nNUM_EPOCHS = 20\nPIN_MEMORY = True\nLEARNING_RATE = 0.001\nDEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-10T06:57:46.306954Z","iopub.execute_input":"2026-10-10T06:57:46.307241Z","iopub.status.idle":"2026-10-10T06:57:46.311752Z","shell.execute_reply.started":"2026-10-10T06:57:46.307217Z","shell.execute_reply":"2026-10-10T06:57:46.310894Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%writefile carvana_losses.py\nimport numpy as np\nimport scipy.ndimage\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\n\nclass WeightedBCELoss2d(nn.Module):\n    def __init__(self):\n        super().__init__()\n\n    def forward(self, logits: torch.Tensor, labels: torch.Tensor, weights: torch.Tensor = None) -> torch.Tensor:\n        logits_f32 = logits.float()\n        labels_f32 = labels.float()\n        loss = F.binary_cross_entropy_with_logits(logits_f32, labels_f32, reduction=\"none\")\n        if weights is not None:\n            weights_f32 = weights.float()\n            loss = loss * weights_f32\n            return loss.sum() / weights_f32.sum().clamp(min=1e-7)\n        return loss.mean()\n\n\nclass SoftDiceLoss(nn.Module):\n    def __init__(self, smooth: float = 1.0):\n        super().__init__()\n        self.smooth = smooth\n\n    def forward(self, logits: torch.Tensor, labels: torch.Tensor, weights: torch.Tensor = None) -> torch.Tensor:\n        logits_f32 = logits.float()\n        labels_f32 = labels.float()\n        probs = torch.sigmoid(logits_f32)\n        num = labels.size(0)\n\n        probs_flat = probs.view(num, -1)\n        labels_flat = labels_f32.view(num, -1)\n\n        if weights is not None:\n            w = weights.view(num, -1).float()\n            intersection = (w * probs_flat * labels_flat).sum(-1)\n            cardinality = (w * probs_flat).sum(-1) + (w * labels_flat).sum(-1)\n        else:\n            intersection = (probs_flat * labels_flat).sum(-1)\n            cardinality = probs_flat.sum(-1) + labels_flat.sum(-1)\n\n        dice = (2.0 * intersection + self.smooth) / (cardinality + self.smooth)\n        # Clamped to guarantee non-negative loss\n        return torch.clamp(1.0 - dice.mean(), min=0.0)\n\n\nclass EdgeLoss(nn.Module):\n    def __init__(self):\n        super().__init__()\n        sobel_x = torch.tensor([[-1.0, 0.0, 1.0], [-2.0, 0.0, 2.0], [-1.0, 0.0, 1.0]], dtype=torch.float32).view(1, 1, 3, 3)\n        sobel_y = torch.tensor([[-1.0, -2.0, -1.0], [0.0, 0.0, 0.0], [1.0, 2.0, 1.0]], dtype=torch.float32).view(1, 1, 3, 3)\n        self.register_buffer(\"kernel_x\", sobel_x)\n        self.register_buffer(\"kernel_y\", sobel_y)\n\n    def _extract_edges(self, x: torch.Tensor) -> torch.Tensor:\n        x_f32 = x.float()\n        kx = self.kernel_x.to(device=x_f32.device, dtype=x_f32.dtype)\n        ky = self.kernel_y.to(device=x_f32.device, dtype=x_f32.dtype)\n        gx = F.conv2d(x_f32, kx, padding=1)\n        gy = F.conv2d(x_f32, ky, padding=1)\n        mag_sq = torch.clamp(gx**2 + gy**2, min=1e-7)\n        return torch.sqrt(mag_sq)\n\n    def forward(self, logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:\n        if logits.dim() == 3:\n            logits = logits.unsqueeze(1)\n        if labels.dim() == 3:\n            labels = labels.unsqueeze(1)\n\n        probs = torch.sigmoid(logits.float())\n        # Normalize Sobel magnitude to [0, 1] range to avoid Dice > 1.0\n        pred_edges = torch.clamp(self._extract_edges(probs) / 4.0, 0.0, 1.0)\n\n        with torch.no_grad():\n            labels_f32 = labels.float()\n            gt_edges = self._extract_edges(labels_f32)\n            gt_edges = (gt_edges > 0.1).float()\n\n        pred_flat = pred_edges.view(logits.size(0), -1)\n        gt_flat = gt_edges.view(labels.size(0), -1)\n\n        intersection = (pred_flat * gt_flat).sum(-1)\n        dice_edge = (2.0 * intersection + 1.0) / (pred_flat.sum(-1) + gt_flat.sum(-1) + 1.0)\n        return torch.clamp(1.0 - dice_edge.mean(), min=0.0)\n\n\nclass CombinedLoss(nn.Module):\n    def __init__(self, mode: str = \"none\", is_log_dice: bool = False, edge_weight: float = 1.0):\n        super().__init__()\n        self.mode = mode\n        self.is_log_dice = is_log_dice\n        self.edge_weight = edge_weight\n\n        self.bce = WeightedBCELoss2d()\n        self.dice = SoftDiceLoss()\n        if self.mode == \"edge\":\n            self.edge_criterion = EdgeLoss()\n\n    def forward(self, logits: torch.Tensor, labels: torch.Tensor):\n        if logits.dim() == 4 and logits.size(1) == 1:\n            logits = logits.squeeze(1)\n        if labels.dim() == 4 and labels.size(1) == 1:\n            labels = labels.squeeze(1)\n\n        logits = logits.float()\n        labels = labels.float()\n\n        bce_loss = self.bce(logits, labels)\n        dice_loss = self.dice(logits, labels)\n\n        if self.mode == \"edge\":\n            aux_loss = self.edge_criterion(logits, labels)\n        else:\n            aux_loss = torch.tensor(0.0, device=logits.device)\n\n        primary_dice = -(1.0 - dice_loss).clamp(min=1e-7).log() if self.is_log_dice else dice_loss\n        total_loss = bce_loss + primary_dice + (self.edge_weight * aux_loss)\n\n        return total_loss, bce_loss, dice_loss, aux_loss","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-10T06:57:46.312879Z","iopub.execute_input":"2026-10-10T06:57:46.313209Z","iopub.status.idle":"2026-10-10T06:57:47.15944Z","shell.execute_reply.started":"2026-10-10T06:57:46.313185Z","shell.execute_reply":"2026-10-10T06:57:47.158456Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%writefile trainer.py\nimport json\nimport os\nimport cv2\nimport numpy as np\nimport torch\nimport torch.distributed as dist\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.nn.parallel import DistributedDataParallel as DDP\nfrom torch.utils.data import DataLoader\nfrom torch.utils.data.distributed import DistributedSampler\nfrom tqdm import tqdm\n\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\n\nfrom carvana_dataset import CarvanaDataset\nfrom carvana_losses import CombinedLoss\nfrom recolor import ColorReplace\nfrom unet_no_bn import UnetVgg11\n\nWORKING_DIR = '/kaggle/working/'\ntrain_dir = os.path.join(WORKING_DIR, 'train_hq/')\nval_dir = os.path.join(WORKING_DIR, 'val/')\ntrain_masks_dir = os.path.join(WORKING_DIR, 'train_masks/')\nval_masks_dir = os.path.join(WORKING_DIR, 'val_masks/')\n\n\nclass HistoryLogger:\n    \"\"\"Logs individual loss components for fine-grained trajectory inspection.\"\"\"\n\n    def __init__(self, filepath=os.path.join(WORKING_DIR, \"train_history.json\")):\n        self.filepath = filepath\n        self.history = {\n            \"train_snapshots\": [],\n            \"epoch_evaluations\": [],\n        }\n\n    def log_train_snapshot(self, epoch: int, step: int, losses: dict, lr: float):\n        self.history[\"train_snapshots\"].append({\n            \"epoch\": epoch,\n            \"step\": step,\n            \"total_loss\": round(float(losses[\"total\"]), 5),\n            \"bce_loss\": round(float(losses[\"bce\"]), 5),\n            \"dice_loss\": round(float(losses[\"dice\"]), 5),\n            \"edge_loss\": round(float(losses[\"edge\"]), 5),\n            \"lr\": float(f\"{lr:.6e}\"),\n        })\n\n    def log_epoch_eval(self, epoch: int, val_acc: float, val_dice: float, lr: float):\n        self.history[\"epoch_evaluations\"].append({\n            \"epoch\": epoch,\n            \"val_acc\": round(float(val_acc), 3),\n            \"val_dice\": round(float(val_dice), 6),\n            \"val_dice_loss\": round(1.0 - float(val_dice), 6),\n            \"lr\": float(f\"{lr:.6e}\"),\n        })\n        with open(self.filepath, \"w\") as f:\n            json.dump(self.history, f, indent=2)\n\n\nclass Trainer:\n    def __init__(\n        self,\n        model: DDP,\n        raw_model: nn.Module,\n        train_data: DataLoader,\n        val_data: DataLoader,\n        optimizer: torch.optim.Optimizer,\n        scheduler: torch.optim.lr_scheduler.ReduceLROnPlateau,\n        scaler: torch.amp.GradScaler,\n        gpu_id: int,\n        accum_steps: int = 4,\n        steps_per_epoch: int = 250,\n        max_norm: float = 1.0,\n    ) -> None:\n        self.gpu_id = gpu_id\n        self.model = model\n        self.raw_model = raw_model\n        self.train_data = train_data\n        self.val_data = val_data\n        self.optimizer = optimizer\n        self.scheduler = scheduler\n        self.scaler = scaler\n        # Initialize directly on GPU to prevent CPU buffer mismatch\n        self.loss_fn = CombinedLoss(mode=\"edge\", is_log_dice=False, edge_weight=1.0).to(gpu_id)\n        self.accum_steps = accum_steps\n        self.steps_per_epoch = steps_per_epoch\n        self.max_norm = max_norm\n\n        self.best_dice = 0.0\n        self.logger = HistoryLogger() if gpu_id == 0 else None\n        self.train_iter = None\n        self.dataset_pass = 0\n\n    def _get_next_train_batch(self):\n        if self.train_iter is None:\n            self.train_data.sampler.set_epoch(self.dataset_pass)\n            self.train_iter = iter(self.train_data)\n        try:\n            return next(self.train_iter)\n        except StopIteration:\n            self.dataset_pass += 1\n            if self.gpu_id == 0:\n                print(f\"\\n>>> [Dataset Stream] Completed Full Pass #{self.dataset_pass} <<<\")\n            self.train_data.sampler.set_epoch(self.dataset_pass)\n            self.train_iter = iter(self.train_data)\n            return next(self.train_iter)\n\n    def _run_batch(self, source, targets, is_accumulating: bool):\n        source = source.to(self.gpu_id, dtype=torch.float32)\n        targets = targets.to(self.gpu_id, dtype=torch.float32).unsqueeze(1)\n\n        with torch.amp.autocast('cuda'):\n            predictions = self.model(source)\n\n        predictions_fp32 = predictions.float()\n        total_loss, bce_loss, dice_loss, aux_loss = self.loss_fn(predictions_fp32, targets)\n        loss_to_scale = total_loss / self.accum_steps\n\n        self.scaler.scale(loss_to_scale).backward()\n\n        if not is_accumulating:\n            self.scaler.unscale_(self.optimizer)\n            torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.max_norm)\n            self.scaler.step(self.optimizer)\n            self.scaler.update()\n            self.optimizer.zero_grad(set_to_none=True)\n\n        return {\n            \"total\": total_loss.item(),\n            \"bce\": bce_loss.item(),\n            \"dice\": dice_loss.item(),\n            \"edge\": aux_loss.item(),\n        }\n\n    def evaluate(self):\n        self.model.eval()\n        local_correct = torch.tensor(0, dtype=torch.int64, device=self.gpu_id)\n        local_pixels = torch.tensor(0, dtype=torch.int64, device=self.gpu_id)\n        local_dice_sum = torch.tensor(0.0, dtype=torch.float64, device=self.gpu_id)\n        local_count = torch.tensor(0, dtype=torch.int64, device=self.gpu_id)\n\n        with torch.no_grad():\n            for x, y in self.val_data:\n                x = x.to(self.gpu_id, dtype=torch.float32)\n                y = y.to(self.gpu_id, dtype=torch.float32).unsqueeze(1)\n\n                with torch.amp.autocast('cuda'):\n                    preds = torch.sigmoid(self.model(x).float())\n                    preds = (preds > 0.5).float()\n\n                local_correct += (preds == y).sum()\n                local_pixels += torch.numel(preds)\n                local_dice_sum += ((2.0 * (preds * y).sum()) / ((preds + y).sum() + 1e-8)).to(torch.float64)\n                local_count += 1\n\n        dist.all_reduce(local_correct, op=dist.ReduceOp.SUM)\n        dist.all_reduce(local_pixels, op=dist.ReduceOp.SUM)\n        dist.all_reduce(local_dice_sum, op=dist.ReduceOp.SUM)\n        dist.all_reduce(local_count, op=dist.ReduceOp.SUM)\n\n        val_acc = (local_correct.item() / local_pixels.item()) * 100.0\n        val_dice = local_dice_sum.item() / max(local_count.item(), 1)\n\n        if self.gpu_id == 0:\n            print(f\"\\n[Validation] Pixel Acc: {val_acc:.2f}% | Mean Dice: {val_dice:.6f}\")\n\n        self.model.train()\n        return val_acc, val_dice\n\n    def _run_epoch(self, epoch: int):\n        self.model.train()\n        self.optimizer.zero_grad(set_to_none=True)\n\n        loop = tqdm(range(self.steps_per_epoch), disable=(self.gpu_id != 0), desc=f\"Epoch {epoch}\")\n        \n        running_losses = {\"total\": 0.0, \"bce\": 0.0, \"dice\": 0.0, \"edge\": 0.0}\n        snapshot_interval = max(self.steps_per_epoch // 5, 1)\n\n        for step in loop:\n            step_losses = {\"total\": 0.0, \"bce\": 0.0, \"dice\": 0.0, \"edge\": 0.0}\n\n            for acc_idx in range(self.accum_steps):\n                source, targets = self._get_next_train_batch()\n                is_accumulating = (acc_idx != self.accum_steps - 1)\n\n                if is_accumulating:\n                    with self.model.no_sync():\n                        batch_losses = self._run_batch(source, targets, is_accumulating=True)\n                else:\n                    batch_losses = self._run_batch(source, targets, is_accumulating=False)\n\n                for k in step_losses:\n                    step_losses[k] += batch_losses[k] / self.accum_steps\n\n            for k in running_losses:\n                running_losses[k] = (running_losses[k] * step + step_losses[k]) / (step + 1)\n\n            if self.gpu_id == 0:\n                current_lr = self.optimizer.param_groups[0]['lr']\n                loop.set_postfix(\n                    tot=f\"{running_losses['total']:.3f}\",\n                    bce=f\"{running_losses['bce']:.3f}\",\n                    dce=f\"{running_losses['dice']:.3f}\",\n                    edg=f\"{running_losses['edge']:.3f}\",\n                    lr=f\"{current_lr:.1e}\"\n                )\n\n                if (step + 1) % snapshot_interval == 0:\n                    self.logger.log_train_snapshot(epoch, step + 1, running_losses, current_lr)\n\n    def _save_checkpoint(self, path: str, epoch: int, val_dice: float):\n        if self.gpu_id != 0:\n            return\n        checkpoint = {\n            \"epoch\": epoch,\n            \"dataset_pass\": self.dataset_pass,\n            \"model_state\": self.raw_model.state_dict(),\n            \"optimizer_state\": self.optimizer.state_dict(),\n            \"val_dice\": val_dice,\n        }\n        torch.save(checkpoint, path)\n        print(f\"--> Saved checkpoint to {path}\")\n\n    def train(self, max_epochs: int):\n        for epoch in range(max_epochs):\n            self._run_epoch(epoch)\n            val_acc, val_dice = self.evaluate()\n\n            val_dice_loss = 1.0 - val_dice\n            self.scheduler.step(val_dice_loss)\n\n            current_lr = self.optimizer.param_groups[0]['lr']\n\n            if self.gpu_id == 0:\n                self.logger.log_epoch_eval(epoch, val_acc, val_dice, current_lr)\n                self._save_checkpoint(os.path.join(WORKING_DIR, \"last.pt\"), epoch, val_dice)\n\n                if val_dice > self.best_dice:\n                    self.best_dice = val_dice\n                    self._save_checkpoint(os.path.join(WORKING_DIR, \"best.pt\"), epoch, val_dice)\n                    print(f\"--> [New Best] Score: {self.best_dice:.6f}\")\n\n\ndef ddp_setup(rank: int, world_size: int):\n    os.environ[\"MASTER_ADDR\"] = \"localhost\"\n    os.environ[\"MASTER_PORT\"] = \"12355\"\n    torch.cuda.set_device(rank)\n    dist.init_process_group(backend=\"nccl\", rank=rank, world_size=world_size)\n\n\ndef load_train_objs(gpu_id: int, batch_size: int):\n    train_pipeline = A.Compose([\n        A.Pad(padding=(2, 0, 0, 0), border_mode=cv2.BORDER_CONSTANT, fill=0, fill_mask=0),\n        A.HorizontalFlip(p=0.5),\n        A.Affine(scale=(0.9, 1.1), rotate=(-5, 5), translate_px=(50, 50), balanced_scale=True, p=0.5),\n        ColorReplace(tolerance=0.45, n_clusters=2, hue_shift_limit=(-180, 180), lightness_weight=0.3, p=0.5),\n        A.RandomGamma(gamma_range=(80, 120), p=0.2),\n        A.RandomBrightnessContrast(brightness_range=(-0.15, 0.15), contrast_range=(-0.15, 0.15), p=0.2),\n        A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),\n        ToTensorV2(),\n    ], strict=True, seed=42)\n\n    val_pipeline = A.Compose([\n        A.Pad(padding=(2, 0, 0, 0), border_mode=cv2.BORDER_CONSTANT, fill=0, fill_mask=0),\n        A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),\n        ToTensorV2(),\n    ], strict=True, seed=42)\n\n    train_dataset = CarvanaDataset(image_dir=train_dir, mask_dir=train_masks_dir, transform=train_pipeline)\n    val_dataset = CarvanaDataset(image_dir=val_dir, mask_dir=val_masks_dir, transform=val_pipeline)\n\n    train_loader = DataLoader(\n        train_dataset,\n        batch_size=batch_size,\n        sampler=DistributedSampler(train_dataset, shuffle=True),\n        num_workers=4,\n        pin_memory=True\n    )\n    val_loader = DataLoader(\n        val_dataset,\n        batch_size=batch_size,\n        sampler=DistributedSampler(val_dataset, shuffle=False),\n        num_workers=4,\n        pin_memory=True\n    )\n\n    raw_model = UnetVgg11(n_classes=1).to(gpu_id)\n\n    optimizer = optim.AdamW(raw_model.parameters(), lr=1.5e-4, weight_decay=1e-4)\n\n    scheduler = optim.lr_scheduler.ReduceLROnPlateau(\n        optimizer,\n        mode=\"min\",\n        factor=0.2,\n        patience=2,\n        threshold=1e-4,\n        min_lr=1e-6\n    )\n\n    ddp_model = DDP(raw_model, device_ids=[gpu_id])\n    scaler = torch.amp.GradScaler('cuda')\n\n    return train_loader, val_loader, ddp_model, raw_model, optimizer, scheduler, scaler\n\n\ndef train_distributed(\n    rank: int, \n    world_size: int, \n    total_epochs: int, \n    batch_size: int, \n    accum_steps: int = 4, \n    steps_per_epoch: int = 250\n):\n    ddp_setup(rank, world_size)\n    train_loader, val_loader, ddp_model, raw_model, optimizer, scheduler, scaler = load_train_objs(rank, batch_size)\n    trainer = Trainer(\n        model=ddp_model,\n        raw_model=raw_model,\n        train_data=train_loader,\n        val_data=val_loader,\n        optimizer=optimizer,\n        scheduler=scheduler,\n        scaler=scaler,\n        gpu_id=rank,\n        accum_steps=accum_steps,\n        steps_per_epoch=steps_per_epoch,\n    )\n    trainer.train(total_epochs)\n    dist.destroy_process_group()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-10T06:57:47.16067Z","iopub.execute_input":"2026-10-10T06:57:47.161Z","iopub.status.idle":"2026-10-10T06:57:47.181243Z","shell.execute_reply.started":"2026-10-10T06:57:47.160976Z","shell.execute_reply":"2026-10-10T06:57:47.180545Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def clean_memory():\n    import gc\n    \n    try:\n        del model\n    except:\n        pass\n    \n    try:\n        if 'trt_model' in globals():\n            del globals()['trt_model']\n        \n        del trt_model\n    except:\n        pass\n    \n    torch._dynamo.reset()\n    torch.cuda.empty_cache()\n    gc.collect()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-10T06:57:47.182399Z","iopub.execute_input":"2026-10-10T06:57:47.182924Z","iopub.status.idle":"2026-10-10T06:57:47.198697Z","shell.execute_reply.started":"2026-10-10T06:57:47.182871Z","shell.execute_reply":"2026-10-10T06:57:47.198022Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"clean_memory()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-10T06:57:47.199916Z","iopub.execute_input":"2026-10-10T06:57:47.200255Z","iopub.status.idle":"2026-10-10T06:57:47.661826Z","shell.execute_reply.started":"2026-10-10T06:57:47.200219Z","shell.execute_reply":"2026-10-10T06:57:47.660825Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"checkpoint = '/kaggle/input/datasets/oleksiistreltsov/carvana-artifacts/best (13).pt'","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-10T06:57:47.663106Z","iopub.execute_input":"2026-10-10T06:57:47.663453Z","iopub.status.idle":"2026-10-10T06:57:47.667625Z","shell.execute_reply.started":"2026-10-10T06:57:47.6634Z","shell.execute_reply":"2026-10-10T06:57:47.667007Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from trainer import train_distributed\n\nif checkpoint is None:\n    world_size = torch.cuda.device_count()\n    mp.spawn(train_distributed, args=(world_size, NUM_EPOCHS, BATCH_SIZE, 4, 250), nprocs=world_size)\nelse:\n    shutil.copyfile(checkpoint, '/kaggle/working/best.pt')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-10T06:57:47.668585Z","iopub.execute_input":"2026-10-10T06:57:47.668868Z","iopub.status.idle":"2026-10-10T06:58:01.22809Z","shell.execute_reply.started":"2026-10-10T06:57:47.668846Z","shell.execute_reply":"2026-10-10T06:58:01.227462Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import json\nimport matplotlib.pyplot as plt\n\nhistory_path = \"/kaggle/working/train_history.json\"\n\nif os.path.exists(history_path):\n    with open(history_path, \"r\") as f:\n        history = json.load(f)\n    \n    snapshots = history[\"train_snapshots\"]\n    steps = [s[\"step\"] + s[\"epoch\"] * 250 for s in snapshots]\n    \n    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 5))\n    \n    # Individual loss components\n    ax1.plot(steps, [s[\"total_loss\"] for s in snapshots], label=\"Total Loss\", color=\"black\", linewidth=2)\n    ax1.plot(steps, [s[\"bce_loss\"] for s in snapshots], label=\"BCE Loss\", color=\"cornflowerblue\", linestyle=\"--\")\n    ax1.plot(steps, [s[\"dice_loss\"] for s in snapshots], label=\"Dice Loss\", color=\"crimson\", linestyle=\"--\")\n    ax1.plot(steps, [s[\"edge_loss\"] for s in snapshots], label=\"Edge Loss\", color=\"darkorange\", linestyle=\":\")\n    ax1.set_xlabel(\"Update Step\")\n    ax1.set_ylabel(\"Loss\")\n    ax1.set_title(\"Training Loss Breakdown\")\n    ax1.legend()\n    ax1.grid(True, linestyle=\"--\", alpha=0.5)\n    \n    # Validation Dice Score\n    evals = history[\"epoch_evaluations\"]\n    epochs = [e[\"epoch\"] for e in evals]\n    dices = [e[\"val_dice\"] for e in evals]\n    ax2.plot(epochs, dices, label=\"Validation Dice\", color=\"forestgreen\", marker=\"o\")\n    ax2.set_xlabel(\"Epoch\")\n    ax2.set_ylabel(\"Dice Score\")\n    ax2.set_title(\"Validation Dice Metric Progression\")\n    ax2.legend()\n    ax2.grid(True, linestyle=\"--\", alpha=0.5)\n    \n    plt.tight_layout()\n    plt.savefig(\"/kaggle/working/loss_breakdown.png\", dpi=200)\n    plt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-10T06:58:01.229099Z","iopub.execute_input":"2026-10-10T06:58:01.229509Z","iopub.status.idle":"2026-10-10T06:58:01.238387Z","shell.execute_reply.started":"2026-10-10T06:58:01.229471Z","shell.execute_reply":"2026-10-10T06:58:01.237544Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Test and Submit","metadata":{}},{"cell_type":"code","source":"def rle_encode(img):\n    '''\n    img: numpy array, 1 - mask, 0 - background\n    Returns run length as string formated\n    '''\n    pixels = img.flatten()\n    pixels[0] = 0\n    pixels[-1] = 0\n    runs = np.where(pixels[1:] != pixels[:-1])[0] + 2\n    runs[1::2] -= runs[:-1:2]\n    \n    return ' '.join(str(x) for x in runs)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-10T06:58:01.239536Z","iopub.execute_input":"2026-10-10T06:58:01.239849Z","iopub.status.idle":"2026-10-10T06:58:01.26333Z","shell.execute_reply.started":"2026-10-10T06:58:01.239814Z","shell.execute_reply":"2026-10-10T06:58:01.262675Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if 'test_hq' not in os.listdir(WORKING_DIR):\n    if os.path.isfile(WORKING_DIR + 'submission.csv'):\n        os.remove(WORKING_DIR + 'submission.csv')\n    with zipfile.ZipFile(DATASET_DIR + 'test_hq.zip', 'r') as zip_file:\n        zip_file.extractall(WORKING_DIR)\nif 'test_images' not in os.listdir(WORKING_DIR):\n    os.mkdir(WORKING_DIR + 'test_images')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-10T06:58:01.264575Z","iopub.execute_input":"2026-10-10T06:58:01.26508Z","iopub.status.idle":"2026-10-10T07:03:29.859148Z","shell.execute_reply.started":"2026-10-10T06:58:01.265044Z","shell.execute_reply":"2026-10-10T07:03:29.858089Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"best_t = 0.27372","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-10T07:03:29.860811Z","iopub.execute_input":"2026-10-10T07:03:29.861841Z","iopub.status.idle":"2026-10-10T07:03:29.866584Z","shell.execute_reply.started":"2026-10-10T07:03:29.86179Z","shell.execute_reply":"2026-10-10T07:03:29.865769Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if best_t is None:\n    NUM_BINS = 1000\n    bins = np.linspace(0.0, 1.0, NUM_BINS + 1)\n    \n    fg_histogram = np.zeros(NUM_BINS, dtype=np.int64)\n    bg_histogram = np.zeros(NUM_BINS, dtype=np.int64)\n    \n    model = UnetVgg11(n_classes=1)\n    checkpoint = torch.load(WORKING_DIR + \"best.pt\", map_location=DEVICE)\n    model.load_state_dict(checkpoint[\"model_state\"])\n    model.to(DEVICE).eval()\n    \n    val_ds = CarvanaDataset(image_dir=val_dir, mask_dir=val_masks_dir, transform=val_transform)\n    val_loader = DataLoader(val_ds, batch_size=1, shuffle=False, num_workers=4, pin_memory=True)\n    \n    with torch.no_grad():\n        for x, y in tqdm(val_loader):\n            x = x.to(DEVICE).to(torch.float32)\n            y = y.to(DEVICE).unsqueeze(1).to(torch.bool)\n            \n            preds_prob = torch.sigmoid(model(x))\n            \n            fg_probs = preds_prob[y].cpu().numpy()\n            bg_probs = preds_prob[~y].cpu().numpy()\n            \n            fg_hist, _ = np.histogram(fg_probs, bins=bins)\n            bg_hist, _ = np.histogram(bg_probs, bins=bins)\n            \n            fg_histogram += fg_hist\n            bg_histogram += bg_hist\n    \n    total_true_pixels = np.sum(fg_histogram)\n\n    def objective(trial):\n        t = trial.suggest_float(\"threshold\", 0.1, 0.9)\n        \n        bin_idx = np.digitize(t, bins) - 1\n        bin_idx = max(0, min(bin_idx, NUM_BINS - 1))\n    \n        intersection = np.sum(fg_histogram[bin_idx:])\n        false_positives = np.sum(bg_histogram[bin_idx:])\n        \n        pred_sum = intersection + false_positives\n        \n        dice = (2.0 * intersection) / (pred_sum + total_true_pixels + 1e-8)\n        return dice\n\n    optuna.logging.set_verbosity(optuna.logging.WARNING)\n    study = optuna.create_study(direction=\"maximize\")\n    \n    study.optimize(objective, n_trials=150) \n    \n    best_t = study.best_params[\"threshold\"]\n    best_d = study.best_value\n    print(f\"Best threshold: {best_t:.5f} | Best Dice: {best_d:.6f}\")\n\n    trials_df = study.trials_dataframe()\n\n    x_thresholds = trials_df['params_threshold']\n    y_dice_scores = trials_df['value']\n    \n    sns.set_theme(style=\"whitegrid\")\n    plt.figure(figsize=(12, 6), dpi=100)\n    \n    sns.scatterplot(\n        x=x_thresholds, \n        y=y_dice_scores, \n        color='royalblue', \n        alpha=0.6, \n        label='Trials',\n        edgecolor=None\n    )\n    \n    sns.lineplot(\n        x=x_thresholds, \n        y=y_dice_scores, \n        color='darkblue', \n        alpha=0.3, \n        ci=None\n    )\n    \n    best_t = study.best_params[\"threshold\"]\n    best_d = study.best_value\n    \n    plt.scatter(\n        best_t, best_d, \n        color='crimson', \n        s=150, \n        zorder=5, \n        label=f'Best threshold: {best_t:.5f}, Dice: {best_d:.6f})',\n        marker='*'\n    )\n    \n    plt.title('Dice Score vs threshold', fontsize=14, pad=15)\n    plt.xlabel('Threshold', fontsize=12)\n    plt.ylabel('Validation Dice Score', fontsize=12)\n    plt.xlim(0.1, 0.9)\n    \n    y_min = y_dice_scores.min() - (best_d - y_dice_scores.min()) * 0.05\n    plt.ylim(max(0, y_min), best_d + (best_d - y_min) * 0.1)\n    \n    plt.legend(loc='lower center', fontsize=11, frameon=True)\n    plt.tight_layout()\n    \n    plt.savefig('optuna_threshold_exploration.png', bbox_inches='tight')\n    plt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-10T07:03:29.867862Z","iopub.execute_input":"2026-10-10T07:03:29.868342Z","iopub.status.idle":"2026-10-10T07:03:30.957023Z","shell.execute_reply.started":"2026-10-10T07:03:29.868303Z","shell.execute_reply":"2026-10-10T07:03:30.956138Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"clean_memory()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-10T07:03:30.958173Z","iopub.execute_input":"2026-10-10T07:03:30.958523Z","iopub.status.idle":"2026-10-10T07:03:31.37417Z","shell.execute_reply.started":"2026-10-10T07:03:30.958486Z","shell.execute_reply":"2026-10-10T07:03:31.373485Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"TEST_DIR = WORKING_DIR + 'test_hq'\nTHRESHOLD = best_t\n\ntest_transform = A.Compose(\n    [\n        A.Pad(padding=(2, 0, 0, 0), border_mode=cv2.BORDER_CONSTANT, fill=0, fill_mask=0),\n        A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),\n        ToTensorV2(), \n    ]\n)\n   \ntest_set = CarvanaTestDataset(\n    image_dir=TEST_DIR,\n    transform=test_transform\n)    \n\n\ntest_loader = DataLoader(\n    test_set, \n    batch_size=BATCH_SIZE * 8, \n    shuffle=False,\n    num_workers=4,          \n    pin_memory=True,        \n    persistent_workers=True\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-10T07:03:31.375338Z","iopub.execute_input":"2026-10-10T07:03:31.375783Z","iopub.status.idle":"2026-10-10T07:03:31.497842Z","shell.execute_reply.started":"2026-10-10T07:03:31.375735Z","shell.execute_reply":"2026-10-10T07:03:31.496938Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"engine = '/kaggle/input/datasets/oleksiistreltsov/carvana-artifacts/carvana_unet_vgg_trt (1).ep'\nnew_model = False","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-10T07:03:31.498975Z","iopub.execute_input":"2026-10-10T07:03:31.499384Z","iopub.status.idle":"2026-10-10T07:03:31.503232Z","shell.execute_reply.started":"2026-10-10T07:03:31.499357Z","shell.execute_reply":"2026-10-10T07:03:31.502283Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"dummy_input = torch.randn((4, 3, 1280, 1920), dtype=torch.float16).cuda()\n\nif engine is None:\n    base_model = UnetVgg11(n_classes=1).half().eval().cuda()\n    base_model.load_state_dict(torch.load(\"/kaggle/working/best.pt\")[\"model_state\"])\n      \n    orig_exp_program = torch.export.export(base_model, (dummy_input,))\n    \n    trt_model = torch_tensorrt.compile(\n        orig_exp_program,\n        ir=\"dynamo\",\n        inputs=[dummy_input],\n        enabled_precisions={torch.float16}, \n        immutable_weights=False,          \n        num_avg_timing_iters=4,           \n        workspace_size=4 << 30 \n    )\n    \n    torch_tensorrt.save(trt_model, \"carvana_unet_vgg_trt.ep\", arg_inputs=[dummy_input])\n    engine = \"carvana_unet_vgg_trt.ep\"\n    print(\"Model saved successfully\")\nelif new_model == True:\n    trt_model = torch_tensorrt.load(engine)\n    \n    new_weights_model = UnetVgg11(n_classes=1).half().eval().cuda()\n    new_weights_model.load_state_dict(torch.load(\"/kaggle/working/best.pt\")[\"model_state\"])\n    \n    trt_model.load_state_dict(new_weights_model.state_dict())\n    \n    _ = trt_model(dummy_input)\n    \n    print(\"Engine is automatically refitted and ready for inference\")\nelse:\n    trt_model = torch_tensorrt.load(engine).module()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-10T07:03:31.504232Z","iopub.execute_input":"2026-10-10T07:03:31.504604Z","iopub.status.idle":"2026-10-10T07:03:35.38122Z","shell.execute_reply.started":"2026-10-10T07:03:31.50458Z","shell.execute_reply":"2026-10-10T07:03:35.38032Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"THRESHOLD = best_t\n\nlarge_range = range(len(test_set))\nsample = random.sample(large_range, k=8)\n\nfor i in tqdm(sample):    \n    img_names, x = test_set[i]\n    \n    x = x.unsqueeze(0).to(DEVICE)\n    \n    with torch.no_grad():\n        z = x.expand(4, -1, -1, -1).to(torch.float16)\n        preds = torch.sigmoid(trt_model(z))\n        preds = preds[0] > THRESHOLD\n    \n    x = TF.crop(\n        x, top=2, left=0, height=1280, width=1918\n    )\n    preds = TF.crop(\n        preds, top=2, left=0, height=1280, width=1918\n    )\n\n    for b_idx in range(1):\n        result = draw_segmentation_masks(x[b_idx] / 4.8 + 0.45, masks=preds[b_idx].to(bool), colors=['cyan'] * 1, alpha=1.0)\n        img = x[b_idx] / 4.8 + 0.45\n        \n        show([img, result])\n        plt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-10T07:03:35.382216Z","iopub.execute_input":"2026-10-10T07:03:35.382607Z","iopub.status.idle":"2026-10-10T07:03:43.604622Z","shell.execute_reply.started":"2026-10-10T07:03:35.382564Z","shell.execute_reply":"2026-10-10T07:03:43.603661Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"clean_memory()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-10T07:03:43.605628Z","iopub.execute_input":"2026-10-10T07:03:43.606048Z","iopub.status.idle":"2026-10-10T07:03:44.015569Z","shell.execute_reply.started":"2026-10-10T07:03:43.606021Z","shell.execute_reply":"2026-10-10T07:03:44.014858Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%writefile ddp_inference.py\nimport os\nimport sys\nimport glob\nimport time\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.distributed as dist\nimport torch.multiprocessing as mp\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.utils.data.distributed import DistributedSampler\nimport torchvision.transforms.functional as TF\nimport torchvision\nfrom tqdm import tqdm\nfrom unet_no_bn import UnetVgg11\n\nimport albumentations as A\nimport cv2\nfrom albumentations.pytorch import ToTensorV2\n\nimport torch_tensorrt\n\nfrom carvana_dataset import CarvanaDataset, CarvanaTestDataset\n\nWORKING_DIR = '/kaggle/working/'\nTEST_DIR = os.path.join(WORKING_DIR, 'test_hq')\nBATCH_SIZE_PER_GPU = 4\n\ndef rle_encode(img):\n    '''\n    img: torch.Tensor or numpy array, 1 - mask, 0 - background\n    Returns run length as string formatted\n    '''\n    if isinstance(img, torch.Tensor):\n        pixels = img.flatten().cpu().numpy()\n    else:\n        pixels = img.flatten()\n        \n    pixels[0] = 0\n    pixels[-1] = 0\n    runs = np.where(pixels[1:] != pixels[:-1])[0] + 2\n    runs[1::2] -= runs[:-1:2]\n    \n    return ' '.join(str(x) for x in runs)\n\ndef ddp_worker(rank, threshold, world_size, checkpoint_path, engine_path):\n    os.environ['MASTER_ADDR'] = 'localhost'\n    os.environ['MASTER_PORT'] = '24681'\n    dist.init_process_group(\"nccl\", rank=rank, world_size=world_size)\n    torch.cuda.set_device(rank)\n\n    test_transform = A.Compose(\n        [\n            A.Pad(padding=(2, 0, 0, 0), border_mode=cv2.BORDER_CONSTANT, fill=0, fill_mask=0),\n            A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),\n            ToTensorV2(), \n        ]\n    )\n    \n    test_set = CarvanaTestDataset(image_dir=TEST_DIR, transform=test_transform)\n    sampler = DistributedSampler(test_set, num_replicas=world_size, rank=rank, shuffle=False, drop_last=False)\n    \n    test_loader = DataLoader(\n        test_set, \n        batch_size=BATCH_SIZE_PER_GPU, \n        shuffle=False,\n        sampler=sampler,\n        num_workers=4,          \n        pin_memory=True,        \n        persistent_workers=True \n    )\n\n    try:\n        trt_model = torch_tensorrt.load(engine_path).module()\n    except NameError:\n        print(f\"[GPU {rank}] Error: Unet is not defined in this file. Please attach model definition.\")\n        sys.exit(1)\n        \n    checkpoint = torch.load(checkpoint_path, map_location=f'cuda:{rank}')[\"model_state\"]\n\n    local_img_names = []\n    local_encodings = []\n    flush_interval = 500\n    shard_file = f\"submission_gpu_{rank}.csv\"\n    \n    pd.DataFrame(columns=[\"img\", \"rle_mask\"]).to_csv(shard_file, index=False)\n    progress_bar = tqdm(test_loader, desc=f\"GPU {rank} Processing\", disable=(rank != 0))\n\n    with torch.no_grad():\n        for idx, (img_names, x) in enumerate((progress_bar if rank == 0 else test_loader)):                        \n            x = x.to(rank).to(torch.float16)\n            \n            preds = torch.sigmoid(trt_model(x))\n            preds = (preds > threshold).float()\n            \n            preds = TF.crop(preds, top=0, left=2, height=1280, width=1918)\n            \n            for b_idx in range(len(img_names)):\n                mask_tensor = preds[b_idx].squeeze()\n                encoding = rle_encode(mask_tensor)\n                local_img_names.append(img_names[b_idx])\n                local_encodings.append(encoding)\n\n            if len(local_img_names) >= flush_interval:\n                df = pd.DataFrame({\"img\": local_img_names, \"rle_mask\": local_encodings})\n                df.to_csv(shard_file, mode='a', header=False, index=False)\n                local_img_names, local_encodings = [], []\n\n        if local_img_names:\n            df = pd.DataFrame({\"img\": local_img_names, \"rle_mask\": local_encodings})\n            df.to_csv(shard_file, mode='a', header=False, index=False)\n\n    dist.destroy_process_group()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-10T07:21:16.927748Z","iopub.execute_input":"2026-10-10T07:21:16.928427Z","iopub.status.idle":"2026-10-10T07:21:16.935898Z","shell.execute_reply.started":"2026-10-10T07:21:16.928395Z","shell.execute_reply":"2026-10-10T07:21:16.934867Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from ddp_inference import ddp_worker","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-10T07:21:18.672206Z","iopub.execute_input":"2026-10-10T07:21:18.672677Z","iopub.status.idle":"2026-10-10T07:21:18.676628Z","shell.execute_reply.started":"2026-10-10T07:21:18.672647Z","shell.execute_reply":"2026-10-10T07:21:18.675969Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"checkpoint_path = \"best.pt\"\nworld_size = torch.cuda.device_count()\n\nif world_size < 2:\n    print(f\"DDP script expects multiple devices. Detected: {world_size}\")\n    \nmp.spawn(ddp_worker, args=(best_t, world_size, checkpoint_path, engine), nprocs=world_size, join=True)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-10T07:21:19.79768Z","iopub.execute_input":"2026-10-10T07:21:19.798153Z","iopub.status.idle":"2026-10-10T07:23:45.376729Z","shell.execute_reply.started":"2026-10-10T07:21:19.798124Z","shell.execute_reply":"2026-10-10T07:23:45.369934Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def merge_csv_chunks_in_place(chunks, output_file = \"submission.csv\", buffer_size_mb = 16):\n    chunks = list(map(lambda p: Path(p), chunks))\n    \n    output_path = output_file\n    buf_size = buffer_size_mb * 1024 * 1024\n\n    first_chunk = chunks[0]\n    first_chunk.rename(output_path)\n\n    with open(output_path, \"ab\", buffering=buf_size) as dst:\n        for chunk in chunks[1:]:\n            with open(chunk, \"rb\", buffering=buf_size) as src:\n                src.readline()\n                shutil.copyfileobj(src, dst, length=buf_size)\n            \n            chunk.unlink()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-10T07:23:49.033644Z","iopub.execute_input":"2026-10-10T07:23:49.034557Z","iopub.status.idle":"2026-10-10T07:23:49.042446Z","shell.execute_reply.started":"2026-10-10T07:23:49.034511Z","shell.execute_reply":"2026-10-10T07:23:49.041769Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"merge_csv_chunks_in_place(\n    chunks = ['/kaggle/working/submission_gpu_0.csv', '/kaggle/working/submission_gpu_1.csv'], \n    output_file=\"/kaggle/working/submission.csv\"\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-10T07:23:51.35986Z","iopub.execute_input":"2026-10-10T07:23:51.36039Z","iopub.status.idle":"2026-10-10T07:23:51.37023Z","shell.execute_reply.started":"2026-10-10T07:23:51.36036Z","shell.execute_reply":"2026-10-10T07:23:51.369367Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"shutil.rmtree(WORKING_DIR + 'train_hq')\nshutil.rmtree(WORKING_DIR + 'test_hq')\nshutil.rmtree(WORKING_DIR + 'saved_images')\nshutil.rmtree(WORKING_DIR + 'test_images')\nshutil.rmtree(WORKING_DIR + 'train_masks')\nshutil.rmtree(WORKING_DIR + 'val')\nshutil.rmtree(WORKING_DIR + 'val_masks')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-10T07:23:52.615326Z","iopub.execute_input":"2026-10-10T07:23:52.616064Z","iopub.status.idle":"2026-10-10T07:23:55.190602Z","shell.execute_reply.started":"2026-10-10T07:23:52.616032Z","shell.execute_reply":"2026-10-10T07:23:55.189984Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"os.chdir('/kaggle/working')\nFileLink('submission.csv')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-10T07:23:56.440011Z","iopub.execute_input":"2026-10-10T07:23:56.440475Z","iopub.status.idle":"2026-10-10T07:23:56.446253Z","shell.execute_reply.started":"2026-10-10T07:23:56.440447Z","shell.execute_reply":"2026-10-10T07:23:56.44548Z"}},"outputs":[],"execution_count":null}]}