{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":6927,"databundleVersionId":45059,"sourceType":"competition"}],"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},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Install required packages\n!pip install -q albumentations segmentation-models-pytorch\n\n# Import packages\nimport os\nimport numpy as np\nimport pandas as pd\nimport zipfile\nfrom glob import glob\nfrom os.path import basename, splitext\nimport shutil\nimport seaborn as sns\nimport matplotlib.pyplot as plt\nimport cv2\nfrom PIL import Image\nfrom tqdm import tqdm\nimport warnings\nwarnings.filterwarnings('ignore')\n\n# Sklearn imports\nfrom sklearn.model_selection import train_test_split\nfrom sklearn.metrics import confusion_matrix, classification_report, jaccard_score\n\n# PyTorch imports\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nfrom torchvision import transforms, models\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\n\n# Set device\nDEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nprint(f\"Using device: {DEVICE}\")\n\n# ====================== DATA PREPARATION ======================\n\ndef unzip_data():\n    \"\"\"Unzip all data files\"\"\"\n    train_zip = '/kaggle/input/carvana-image-masking-challenge/train.zip'\n    test_zip = '/kaggle/input/carvana-image-masking-challenge/test.zip'\n    mask_zip = '/kaggle/input/carvana-image-masking-challenge/train_masks.zip'\n    \n    os.makedirs(\"train\", exist_ok=True)\n    os.makedirs(\"train_masks\", exist_ok=True)\n    os.makedirs(\"test\", exist_ok=True)\n    \n    with zipfile.ZipFile(train_zip, 'r') as zip_ref:\n        zip_ref.extractall(\"train\")\n    with zipfile.ZipFile(mask_zip, 'r') as zip_ref:\n        zip_ref.extractall(\"train_masks\")\n    with zipfile.ZipFile(test_zip, 'r') as zip_ref:\n        zip_ref.extractall(\"test\")\n    \n    print(\"✅ Data unzipped successfully\")\n\n# Unzip data\nunzip_data()\n\n# Get file paths\nTRAIN_IMG_DIR = 'train/train'\nTRAIN_MASK_DIR = 'train_masks/train_masks'\nTEST_IMG_DIR = 'test/test'\n\ntrain_images = sorted(glob(os.path.join(TRAIN_IMG_DIR, '*.jpg')))\ntrain_masks = sorted(glob(os.path.join(TRAIN_MASK_DIR, '*.gif')))\ntest_images = sorted(glob(os.path.join(TEST_IMG_DIR, '*.jpg')))\n\nprint(f\"Train images: {len(train_images)}\")\nprint(f\"Train masks: {len(train_masks)}\")\nprint(f\"Test images: {len(test_images)}\")\n\n# ====================== DATASET CLASS ======================\n\nclass CarvanaDataset(Dataset):\n    \"\"\"Enhanced dataset with albumentations support\"\"\"\n    def __init__(self, image_paths, mask_paths=None, img_size=(256, 256), \n                 augmentation=None, preprocessing=None):\n        self.image_paths = image_paths\n        self.mask_paths = mask_paths\n        self.img_size = img_size\n        self.augmentation = augmentation\n        self.preprocessing = preprocessing\n    \n    def __len__(self):\n        return len(self.image_paths)\n    \n    def __getitem__(self, idx):\n        # Read image\n        img = cv2.imread(self.image_paths[idx])\n        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n        \n        # Read mask if available\n        if self.mask_paths:\n            mask_gif = Image.open(self.mask_paths[idx]).convert(\"L\")\n            mask = np.array(mask_gif)\n        else:\n            mask = np.zeros((img.shape[0], img.shape[1]), dtype=np.uint8)\n        \n        # Resize\n        img = cv2.resize(img, self.img_size)\n        mask = cv2.resize(mask, self.img_size, interpolation=cv2.INTER_NEAREST)\n        \n        # Apply augmentations\n        if self.augmentation:\n            sample = self.augmentation(image=img, mask=mask)\n            img, mask = sample['image'], sample['mask']\n        \n        # Apply preprocessing\n        if self.preprocessing:\n            sample = self.preprocessing(image=img, mask=mask)\n            img, mask = sample['image'], sample['mask']\n        else:\n            img = img.astype(np.float32) / 255.0\n            img = torch.from_numpy(img).permute(2, 0, 1)\n            \n        mask = (mask > 127).astype(np.float32)\n        mask = torch.from_numpy(mask).unsqueeze(0)\n        \n        return img, mask\n\ndef get_augmentation():\n    \"\"\"Define augmentation pipeline\"\"\"\n    train_transform = A.Compose([\n        A.RandomRotate90(p=0.5),\n        A.Flip(p=0.5),\n        A.ShiftScaleRotate(shift_limit=0.1, scale_limit=0.1, \n                          rotate_limit=10, p=0.5),\n        A.OneOf([\n            A.GaussNoise(p=1),\n            A.GaussianBlur(p=1),\n            A.MotionBlur(p=1),\n        ], p=0.3),\n        A.OneOf([\n            A.RandomBrightnessContrast(p=1),\n            A.HueSaturationValue(p=1),\n        ], p=0.3),\n    ])\n    return train_transform\n\ndef get_preprocessing():\n    \"\"\"Preprocessing pipeline\"\"\"\n    _transform = A.Compose([\n        A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),\n        ToTensorV2(),\n    ])\n    return _transform\n\n# ====================== MODEL ARCHITECTURES ======================\n\nclass DoubleConv(nn.Module):\n    \"\"\"Double convolution block\"\"\"\n    def __init__(self, in_channels, out_channels, mid_channels=None):\n        super().__init__()\n        if not mid_channels:\n            mid_channels = out_channels\n        self.double_conv = nn.Sequential(\n            nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1, bias=False),\n            nn.BatchNorm2d(mid_channels),\n            nn.ReLU(inplace=True),\n            nn.Conv2d(mid_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.double_conv(x)\n\nclass AttentionBlock(nn.Module):\n    \"\"\"Attention gate for U-Net\"\"\"\n    def __init__(self, F_g, F_l, F_int):\n        super().__init__()\n        self.W_g = nn.Sequential(\n            nn.Conv2d(F_g, F_int, kernel_size=1, stride=1, padding=0, bias=True),\n            nn.BatchNorm2d(F_int)\n        )\n        self.W_x = nn.Sequential(\n            nn.Conv2d(F_l, F_int, kernel_size=1, stride=1, padding=0, bias=True),\n            nn.BatchNorm2d(F_int)\n        )\n        self.psi = nn.Sequential(\n            nn.Conv2d(F_int, 1, kernel_size=1, stride=1, padding=0, bias=True),\n            nn.BatchNorm2d(1),\n            nn.Sigmoid()\n        )\n        self.relu = nn.ReLU(inplace=True)\n    \n    def forward(self, g, x):\n        g1 = self.W_g(g)\n        x1 = self.W_x(x)\n        psi = self.relu(g1 + x1)\n        psi = self.psi(psi)\n        return x * psi\n\nclass ImprovedUNet(nn.Module):\n    \"\"\"U-Net with attention gates and deep supervision\"\"\"\n    def __init__(self, n_channels=3, n_classes=1, features=[64, 128, 256, 512]):\n        super().__init__()\n        self.n_channels = n_channels\n        self.n_classes = n_classes\n        \n        # Encoder\n        self.encoder1 = DoubleConv(n_channels, features[0])\n        self.pool1 = nn.MaxPool2d(2)\n        self.encoder2 = DoubleConv(features[0], features[1])\n        self.pool2 = nn.MaxPool2d(2)\n        self.encoder3 = DoubleConv(features[1], features[2])\n        self.pool3 = nn.MaxPool2d(2)\n        self.encoder4 = DoubleConv(features[2], features[3])\n        self.pool4 = nn.MaxPool2d(2)\n        \n        # Bottleneck\n        self.bottleneck = DoubleConv(features[3], features[3]*2)\n        \n        # Decoder with attention\n        self.upconv4 = nn.ConvTranspose2d(features[3]*2, features[3], kernel_size=2, stride=2)\n        self.att4 = AttentionBlock(F_g=features[3], F_l=features[3], F_int=features[3]//2)\n        self.decoder4 = DoubleConv(features[3]*2, features[3])\n        \n        self.upconv3 = nn.ConvTranspose2d(features[3], features[2], kernel_size=2, stride=2)\n        self.att3 = AttentionBlock(F_g=features[2], F_l=features[2], F_int=features[2]//2)\n        self.decoder3 = DoubleConv(features[2]*2, features[2])\n        \n        self.upconv2 = nn.ConvTranspose2d(features[2], features[1], kernel_size=2, stride=2)\n        self.att2 = AttentionBlock(F_g=features[1], F_l=features[1], F_int=features[1]//2)\n        self.decoder2 = DoubleConv(features[1]*2, features[1])\n        \n        self.upconv1 = nn.ConvTranspose2d(features[1], features[0], kernel_size=2, stride=2)\n        self.att1 = AttentionBlock(F_g=features[0], F_l=features[0], F_int=features[0]//2)\n        self.decoder1 = DoubleConv(features[0]*2, features[0])\n        \n        self.conv = nn.Conv2d(features[0], n_classes, kernel_size=1)\n    \n    def forward(self, x):\n        # Encoder\n        enc1 = self.encoder1(x)\n        enc2 = self.encoder2(self.pool1(enc1))\n        enc3 = self.encoder3(self.pool2(enc2))\n        enc4 = self.encoder4(self.pool3(enc3))\n        \n        # Bottleneck\n        bottleneck = self.bottleneck(self.pool4(enc4))\n        \n        # Decoder with attention\n        dec4 = self.upconv4(bottleneck)\n        enc4 = self.att4(g=dec4, x=enc4)\n        dec4 = torch.cat((enc4, dec4), dim=1)\n        dec4 = self.decoder4(dec4)\n        \n        dec3 = self.upconv3(dec4)\n        enc3 = self.att3(g=dec3, x=enc3)\n        dec3 = torch.cat((enc3, dec3), dim=1)\n        dec3 = self.decoder3(dec3)\n        \n        dec2 = self.upconv2(dec3)\n        enc2 = self.att2(g=dec2, x=enc2)\n        dec2 = torch.cat((enc2, dec2), dim=1)\n        dec2 = self.decoder2(dec2)\n        \n        dec1 = self.upconv1(dec2)\n        enc1 = self.att1(g=dec1, x=enc1)\n        dec1 = torch.cat((enc1, dec1), dim=1)\n        dec1 = self.decoder1(dec1)\n        \n        return self.conv(dec1)\n\n# ====================== LOSS FUNCTIONS ======================\n\nclass DiceBCELoss(nn.Module):\n    \"\"\"Combined Dice and BCE loss\"\"\"\n    def __init__(self, weight=None, size_average=True):\n        super().__init__()\n        self.bce = nn.BCEWithLogitsLoss()\n        \n    def forward(self, inputs, targets, smooth=1):\n        # BCE loss\n        bce = self.bce(inputs, targets)\n        \n        # Dice loss\n        inputs = torch.sigmoid(inputs)\n        inputs = inputs.view(-1)\n        targets = targets.view(-1)\n        intersection = (inputs * targets).sum()\n        dice_loss = 1 - (2.*intersection + smooth)/(inputs.sum() + targets.sum() + smooth)\n        \n        return bce + dice_loss\n\nclass FocalTverskyLoss(nn.Module):\n    \"\"\"Focal Tversky Loss for handling class imbalance\"\"\"\n    def __init__(self, alpha=0.7, beta=0.3, gamma=0.75, smooth=1e-6):\n        super().__init__()\n        self.alpha = alpha\n        self.beta = beta\n        self.gamma = gamma\n        self.smooth = smooth\n    \n    def forward(self, inputs, targets):\n        inputs = torch.sigmoid(inputs)\n        inputs = inputs.view(-1)\n        targets = targets.view(-1)\n        \n        # True Positives, False Positives & False Negatives\n        TP = (inputs * targets).sum()\n        FP = ((1-targets) * inputs).sum()\n        FN = (targets * (1-inputs)).sum()\n        \n        Tversky = (TP + self.smooth) / (TP + self.alpha*FP + self.beta*FN + self.smooth)\n        FocalTversky = (1 - Tversky)**self.gamma\n        \n        return FocalTversky\n\n# ====================== METRICS ======================\n\ndef calculate_metrics(pred, target, threshold=0.5):\n    \"\"\"Calculate IoU, Dice, and Pixel Accuracy\"\"\"\n    pred = torch.sigmoid(pred)\n    pred = (pred > threshold).float()\n    \n    # Flatten\n    pred = pred.view(-1).cpu().numpy()\n    target = target.view(-1).cpu().numpy()\n    \n    # IoU\n    intersection = np.sum(pred * target)\n    union = np.sum(pred) + np.sum(target) - intersection\n    iou = (intersection + 1e-6) / (union + 1e-6)\n    \n    # Dice\n    dice = (2 * intersection + 1e-6) / (np.sum(pred) + np.sum(target) + 1e-6)\n    \n    # Pixel Accuracy\n    correct = np.sum(pred == target)\n    total = len(pred)\n    accuracy = correct / total\n    \n    return iou, dice, accuracy\n\n# ====================== TRAINING FUNCTIONS ======================\n\nclass EarlyStopping:\n    \"\"\"Early stopping to prevent overfitting\"\"\"\n    def __init__(self, patience=7, verbose=False, delta=0):\n        self.patience = patience\n        self.verbose = verbose\n        self.counter = 0\n        self.best_score = None\n        self.early_stop = False\n        self.val_loss_min = np.Inf\n        self.delta = delta\n    \n    def __call__(self, val_loss, model, path='checkpoint.pth'):\n        score = -val_loss\n        \n        if self.best_score is None:\n            self.best_score = score\n            self.save_checkpoint(val_loss, model, path)\n        elif score < self.best_score + self.delta:\n            self.counter += 1\n            if self.verbose:\n                print(f'EarlyStopping counter: {self.counter} out of {self.patience}')\n            if self.counter >= self.patience:\n                self.early_stop = True\n        else:\n            self.best_score = score\n            self.save_checkpoint(val_loss, model, path)\n            self.counter = 0\n    \n    def save_checkpoint(self, val_loss, model, path):\n        if self.verbose:\n            print(f'Validation loss decreased ({self.val_loss_min:.6f} --> {val_loss:.6f}). Saving model ...')\n        torch.save(model.state_dict(), path)\n        self.val_loss_min = val_loss\n\ndef train_model(model, train_loader, val_loader, num_epochs=30, lr=1e-3, patience=10):\n    \"\"\"Enhanced training function with early stopping and learning rate scheduling\"\"\"\n    \n    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-4)\n    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(\n        optimizer, mode='min', factor=0.5, patience=5, verbose=True\n    )\n    criterion = DiceBCELoss()\n    early_stopping = EarlyStopping(patience=patience, verbose=True)\n    \n    train_losses = []\n    val_losses = []\n    val_ious = []\n    val_dices = []\n    \n    for epoch in range(num_epochs):\n        # Training\n        model.train()\n        train_loss = 0.0\n        train_bar = tqdm(train_loader, desc=f'Epoch {epoch+1}/{num_epochs} [Train]')\n        \n        for imgs, masks in train_bar:\n            imgs = imgs.to(DEVICE)\n            masks = masks.to(DEVICE)\n            \n            optimizer.zero_grad()\n            outputs = model(imgs)\n            loss = criterion(outputs, masks)\n            loss.backward()\n            \n            # Gradient clipping\n            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n            \n            optimizer.step()\n            \n            train_loss += loss.item() * imgs.size(0)\n            train_bar.set_postfix(loss=loss.item())\n        \n        avg_train_loss = train_loss / len(train_loader.dataset)\n        train_losses.append(avg_train_loss)\n        \n        # Validation\n        model.eval()\n        val_loss = 0.0\n        ious, dices, accs = [], [], []\n        \n        with torch.no_grad():\n            val_bar = tqdm(val_loader, desc=f'Epoch {epoch+1}/{num_epochs} [Val]')\n            for imgs, masks in val_bar:\n                imgs = imgs.to(DEVICE)\n                masks = masks.to(DEVICE)\n                \n                outputs = model(imgs)\n                loss = criterion(outputs, masks)\n                val_loss += loss.item() * imgs.size(0)\n                \n                # Calculate metrics\n                iou, dice, acc = calculate_metrics(outputs, masks)\n                ious.append(iou)\n                dices.append(dice)\n                accs.append(acc)\n                \n                val_bar.set_postfix(loss=loss.item())\n        \n        avg_val_loss = val_loss / len(val_loader.dataset)\n        avg_iou = np.mean(ious)\n        avg_dice = np.mean(dices)\n        avg_acc = np.mean(accs)\n        \n        val_losses.append(avg_val_loss)\n        val_ious.append(avg_iou)\n        val_dices.append(avg_dice)\n        \n        print(f'\\nEpoch {epoch+1}/{num_epochs}:')\n        print(f'Train Loss: {avg_train_loss:.4f}')\n        print(f'Val Loss: {avg_val_loss:.4f} | IoU: {avg_iou:.4f} | Dice: {avg_dice:.4f} | Acc: {avg_acc:.4f}')\n        print('-' * 50)\n        \n        # Scheduler step\n        scheduler.step(avg_val_loss)\n        \n        # Early stopping\n        early_stopping(avg_val_loss, model)\n        if early_stopping.early_stop:\n            print(\"Early stopping triggered\")\n            break\n    \n    # Load best model\n    model.load_state_dict(torch.load('checkpoint.pth'))\n    \n    return model, train_losses, val_losses, val_ious, val_dices\n\n# ====================== VISUALIZATION ======================\n\ndef plot_training_history(train_losses, val_losses, val_ious, val_dices):\n    \"\"\"Plot training metrics\"\"\"\n    fig, axes = plt.subplots(1, 3, figsize=(15, 4))\n    \n    # Loss\n    axes[0].plot(train_losses, label='Train Loss')\n    axes[0].plot(val_losses, label='Val Loss')\n    axes[0].set_xlabel('Epoch')\n    axes[0].set_ylabel('Loss')\n    axes[0].set_title('Training and Validation Loss')\n    axes[0].legend()\n    axes[0].grid(True)\n    \n    # IoU\n    axes[1].plot(val_ious, label='Val IoU', color='green')\n    axes[1].set_xlabel('Epoch')\n    axes[1].set_ylabel('IoU')\n    axes[1].set_title('Validation IoU')\n    axes[1].legend()\n    axes[1].grid(True)\n    \n    # Dice\n    axes[2].plot(val_dices, label='Val Dice', color='orange')\n    axes[2].set_xlabel('Epoch')\n    axes[2].set_ylabel('Dice')\n    axes[2].set_title('Validation Dice Score')\n    axes[2].legend()\n    axes[2].grid(True)\n    \n    plt.tight_layout()\n    plt.show()\n\ndef visualize_predictions(model, dataset, num_samples=4):\n    \"\"\"Visualize model predictions with overlay\"\"\"\n    model.eval()\n    fig, axes = plt.subplots(num_samples, 4, figsize=(16, num_samples*4))\n    \n    indices = np.random.choice(len(dataset), num_samples, replace=False)\n    \n    with torch.no_grad():\n        for i, idx in enumerate(indices):\n            img, mask = dataset[idx]\n            img_input = img.unsqueeze(0).to(DEVICE)\n            \n            # Get prediction\n            pred = model(img_input)\n            pred = torch.sigmoid(pred)\n            pred = (pred > 0.5).float()\n            \n            # Convert to numpy\n            img_np = img.cpu().numpy().transpose(1, 2, 0)\n            # Denormalize if normalized\n            img_np = (img_np * np.array([0.229, 0.224, 0.225]) + \n                     np.array([0.485, 0.456, 0.406]))\n            img_np = np.clip(img_np, 0, 1)\n            \n            mask_np = mask.squeeze().cpu().numpy()\n            pred_np = pred.squeeze().cpu().numpy()\n            \n            # Create overlay\n            overlay = img_np.copy()\n            overlay[:, :, 1] = np.where(mask_np > 0.5, 1, overlay[:, :, 1])  # Green for GT\n            overlay[:, :, 0] = np.where(pred_np > 0.5, 1, overlay[:, :, 0])  # Red for Pred\n            \n            # Plot\n            axes[i, 0].imshow(img_np)\n            axes[i, 0].set_title('Original Image')\n            axes[i, 0].axis('off')\n            \n            axes[i, 1].imshow(mask_np, cmap='gray')\n            axes[i, 1].set_title('Ground Truth')\n            axes[i, 1].axis('off')\n            \n            axes[i, 2].imshow(pred_np, cmap='gray')\n            axes[i, 2].set_title('Prediction')\n            axes[i, 2].axis('off')\n            \n            axes[i, 3].imshow(overlay)\n            axes[i, 3].set_title('Overlay (Green=GT, Red=Pred)')\n            axes[i, 3].axis('off')\n    \n    plt.tight_layout()\n    plt.show()\n\n# ====================== MAIN EXECUTION ======================\n\n# Configuration\nIMG_SIZE = (256, 256)  # Increased from 128x128\nBATCH_SIZE = 8 if torch.cuda.is_available() else 2\nNUM_EPOCHS = 50\nLR = 1e-3\nVAL_SPLIT = 0.15\n\n# Split data\ntrain_imgs, val_imgs, train_masks, val_masks = train_test_split(\n    train_images, train_masks, test_size=VAL_SPLIT, random_state=42\n)\n\nprint(f\"Training samples: {len(train_imgs)}\")\nprint(f\"Validation samples: {len(val_imgs)}\")\n\n# Create datasets\ntrain_dataset = CarvanaDataset(\n    train_imgs, train_masks, \n    img_size=IMG_SIZE,\n    augmentation=get_augmentation(),\n    preprocessing=get_preprocessing()\n)\n\nval_dataset = CarvanaDataset(\n    val_imgs, val_masks,\n    img_size=IMG_SIZE,\n    preprocessing=get_preprocessing()\n)\n\n# Create data loaders\ntrain_loader = DataLoader(\n    train_dataset, \n    batch_size=BATCH_SIZE, \n    shuffle=True, \n    num_workers=2,\n    pin_memory=True\n)\n\nval_loader = DataLoader(\n    val_dataset, \n    batch_size=BATCH_SIZE, \n    shuffle=False,\n    num_workers=2,\n    pin_memory=True\n)\n\n# Initialize model\nmodel = ImprovedUNet(n_channels=3, n_classes=1).to(DEVICE)\nprint(f\"Model parameters: {sum(p.numel() for p in model.parameters()):,}\")\n\n# Train model\nmodel, train_losses, val_losses, val_ious, val_dices = train_model(\n    model, train_loader, val_loader, \n    num_epochs=NUM_EPOCHS, \n    lr=LR, \n    patience=10\n)\n\n# Plot training history\nplot_training_history(train_losses, val_losses, val_ious, val_dices)\n\n# Visualize predictions\nvisualize_predictions(model, val_dataset, num_samples=4)\n\n# ====================== FINAL EVALUATION ======================\n\ndef evaluate_model(model, val_loader):\n    \"\"\"Comprehensive model evaluation\"\"\"\n    model.eval()\n    all_preds = []\n    all_targets = []\n    \n    with torch.no_grad():\n        for imgs, masks in tqdm(val_loader, desc='Evaluating'):\n            imgs = imgs.to(DEVICE)\n            masks = masks.to(DEVICE)\n            \n            preds = model(imgs)\n            preds = torch.sigmoid(preds)\n            preds = (preds > 0.5).float()\n            \n            all_preds.append(preds.cpu().numpy())\n            all_targets.append(masks.cpu().numpy())\n    \n    all_preds = np.concatenate(all_preds).reshape(-1)\n    all_targets = np.concatenate(all_targets).reshape(-1)\n    \n    # Calculate metrics\n    cm = confusion_matrix(all_targets, all_preds)\n    report = classification_report(all_targets, all_preds, digits=4)\n    \n    # Plot confusion matrix\n    plt.figure(figsize=(6, 5))\n    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues',\n                xticklabels=['Background', 'Car'],\n                yticklabels=['Background', 'Car'])\n    plt.xlabel('Predicted')\n    plt.ylabel('Actual')\n    plt.title('Confusion Matrix')\n    plt.show()\n    \n    print(\"\\nClassification Report:\")\n    print(report)\n    \n    # Calculate final IoU\n    final_iou = jaccard_score(all_targets, all_preds, average='binary')\n    print(f\"\\nFinal IoU Score: {final_iou:.4f}\")\n    \n    return final_iou\n\n# Evaluate model\nfinal_iou = evaluate_model(model, val_loader)\n\n# ====================== TEST PREDICTIONS ======================\n\ndef predict_test_set(model, test_images, batch_size=8):\n    \"\"\"Generate predictions for test set\"\"\"\n    test_dataset = CarvanaDataset(\n        test_images, \n        mask_paths=None,\n        img_size=IMG_SIZE,\n        preprocessing=get_preprocessing()\n    )\n    \n    test_loader = DataLoader(\n        test_dataset,\n        batch_size=batch_size,\n        shuffle=False,\n        num_workers=2\n    )\n    \n    model.eval()\n    predictions = []\n    \n    with torch.no_grad():\n        for imgs, _ in tqdm(test_loader, desc='Predicting on test set'):\n            imgs = imgs.to(DEVICE)\n            preds = model(imgs)\n            preds = torch.sigmoid(preds)\n            preds = (preds > 0.5).float()\n            predictions.append(preds.cpu().numpy())\n    \n    predictions = np.concatenate(predictions, axis=0)\n    return predictions\n\n# Generate test predictions (uncomment if needed)\n# test_predictions = predict_test_set(model, test_images[:100])  # Predict first 100 for demo\n# print(f\"Test predictions shape: {test_predictions.shape}\")\n\n# Save final model\ntorch.save(model.state_dict(), 'carvana_unet_final.pth')\nprint(\"\\n✅ Model saved as 'carvana_unet_final.pth'\")\n\nprint(\"\\n\" + \"=\"*50)\nprint(\"Training Complete!\")\nprint(f\"Best Validation IoU: {max(val_ious):.4f}\")\nprint(f\"Best Validation Dice: {max(val_dices):.4f}\")\nprint(\"=\"*50)","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}