{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.11","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceType":"competition","sourceId":97838,"databundleVersionId":11654191}],"dockerImageVersionId":31041,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# ======================================================================================\n# 1. IMPORTS\n# ======================================================================================\n\n# --- Standard libraries ---\nimport os\nimport random\n\n# --- Data libraries ---\nimport numpy as np\nimport pandas as pd\nimport cv2\nfrom PIL import Image\nimport matplotlib.pyplot as plt\nfrom tqdm import tqdm\n\n# --- PyTorch Core ---\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.utils.data import Dataset, DataLoader\n\n# --- Transforms and albumentations ---\nimport torchvision.transforms as T\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\n\n\n# ======================================================================================\n# 2. Env. Codes\n# ======================================================================================\n\nDATASET_PATH = \"/kaggle/input/a0-2025-medical-image-segmentation/Dataset\" # <- BU SATIR DEĞİŞTİRİLDİ\n\nTRAIN_IMAGE_DIR = os.path.join(DATASET_PATH, \"Train/Image\")\nTRAIN_MASK_DIR = os.path.join(DATASET_PATH, \"Train/Mask\")\nTEST_DIR = os.path.join(DATASET_PATH, \"Test/Image\")\n\n# --- Constants for the image size\nIMAGE_HEIGHT = 256\nIMAGE_WIDTH = 256\n\n# --- Validation of the path\n# print(\"Dataset print...\")\n# !ls -R \"{DATASET_PATH}\"\n\n# Gpu availability check.\nprint(\"\\nGPU check:\")\ngpu_info = !nvidia-smi\ngpu_info = '\\n'.join(gpu_info)\nif 'failed' in gpu_info or 'command not found' in gpu_info:\n  print('-> There is no gpu...')\nelse:\n  print(gpu_info)","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-06-26T06:22:07.953644Z","iopub.execute_input":"2025-06-26T06:22:07.954250Z","iopub.status.idle":"2025-06-26T06:22:08.073479Z","shell.execute_reply.started":"2025-06-26T06:22:07.954226Z","shell.execute_reply":"2025-06-26T06:22:08.072887Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ======================================================================================\n# 3. DATA PREPARATION\n# This cell finds all image-mask pairs, shuffles them, and splits them\n# into training and validation sets.\n# ======================================================================================\n\n# --- 3.1. Find and Pair Images with Masks ---\n\n# Get a list of all mask filenames from the mask directory.\n# We assume that for every mask file, a corresponding image file exists.\nmask_files = sorted([f for f in os.listdir(TRAIN_MASK_DIR) if f.lower().endswith('.png')])\n\n# Create a list of tuples, where each tuple contains the full path\n# to a mask and its corresponding image.\n# We assume image files are .jpg and have the same name as the .png masks.\nimage_mask_pairs = [\n    (os.path.join(TRAIN_MASK_DIR, name), os.path.join(TRAIN_IMAGE_DIR, name.replace(\".png\", \".jpg\")))\n    for name in mask_files\n]\n\nprint(f\"Total image-mask pairs found: {len(image_mask_pairs)}\")\n\n\n# --- 3.2. Shuffle and Split the Dataset (80/20) ---\n\n# Set a random seed for reproducibility. This ensures the shuffle is the same every time.\nrandom.seed(42)\nrandom.shuffle(image_mask_pairs)\n\n# Calculate the split index for an 80% training and 20% validation set.\nsplit_idx = int(len(image_mask_pairs) * 0.8)\ntrain_pairs = image_mask_pairs[:split_idx]\nval_pairs = image_mask_pairs[split_idx:]\n\n\n# --- 3.3. Verify the Split and Display a Sample ---\n\n# Print the number of pairs in each set to verify the split.\nprint(f\"Number of training pairs: {len(train_pairs)}\")\nprint(f\"Number of validation pairs: {len(val_pairs)}\")\nprint(f\"\\nExample training pair (Mask, Image): {train_pairs[0]}\")\n\n# Visualize the first training sample to ensure data is loaded correctly.\nprint(\"\\nDisplaying a sample pair from the training set...\")\nfig, axs = plt.subplots(1, 2, figsize=(12, 6))\n\n# Display the mask\nmask_sample = Image.open(train_pairs[0][0])\naxs[0].imshow(mask_sample, cmap=\"gray\")\naxs[0].set_title(\"Image Mask\")\naxs[0].axis(\"off\")\n\n# Display the original image\nimage_sample = Image.open(train_pairs[0][1])\naxs[1].imshow(image_sample)\naxs[1].set_title(\"Original Image\")\naxs[1].axis(\"off\")\n\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-26T06:22:11.201843Z","iopub.execute_input":"2025-06-26T06:22:11.202130Z","iopub.status.idle":"2025-06-26T06:22:11.599758Z","shell.execute_reply.started":"2025-06-26T06:22:11.202110Z","shell.execute_reply":"2025-06-26T06:22:11.599038Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ======================================================================================\n# 4. DATA PIPELINE SETUP\n# This cell defines the transformations, the custom Dataset class,\n# and creates the DataLoader instances for training and validation.\n# ======================================================================================\n\n# --- 4.1. Define Image and Mask Transformations using Albumentations ---\n\nBATCH_SIZE = 8\nNUM_WORKERS = 2\n\n# Define the augmentations for the training set.\n# These random transformations help the model generalize better and prevent overfitting.\ntrain_transforms = A.Compose([\n    A.Resize(height=IMAGE_HEIGHT, width=IMAGE_WIDTH),\n    A.HorizontalFlip(p=0.5),\n    A.VerticalFlip(p=0.5),\n    A.Rotate(limit=35, p=0.5),\n    A.ColorJitter(p=0.2), # This is applied only to the image.\n    A.Normalize(\n        mean=[0.485, 0.456, 0.406],\n        std=[0.229, 0.224, 0.225],\n        max_pixel_value=255.0,\n    ),\n    ToTensorV2(), # Converts image and mask to PyTorch Tensors.\n])\n\n# Define the transformations for the validation set.\n# No random augmentations are applied to ensure consistent evaluation.\nval_transforms = A.Compose([\n    A.Resize(height=IMAGE_HEIGHT, width=IMAGE_WIDTH),\n    A.Normalize(\n        mean=[0.485, 0.456, 0.406],\n        std=[0.229, 0.224, 0.225],\n        max_pixel_value=255.0,\n    ),\n    ToTensorV2(),\n])\n\n\n# --- 4.2. Create the Custom Dataset Class ---\n\nclass SegmentationDataset(Dataset):\n    \"\"\"\n    Custom PyTorch Dataset for loading image-mask pairs.\n    It uses Albumentations for synchronized transformations on both\n    the image and the mask.\n    \"\"\"\n    def __init__(self, image_mask_pairs, transform=None):\n        self.image_mask_pairs = image_mask_pairs\n        self.transform = transform\n\n    def __len__(self):\n        # Returns the total number of samples in the dataset.\n        return len(self.image_mask_pairs)\n\n    def __getitem__(self, idx):\n        # Get the file paths for the image and mask at the given index.\n        mask_path, image_path = self.image_mask_pairs[idx]\n\n        # Load image and mask using OpenCV.\n        image = cv2.imread(image_path)\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) # Convert from BGR to RGB.\n        mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)\n\n        # Normalize mask to be binary (0.0 or 1.0) and of type float32.\n        # This is a critical step to ensure the loss function works correctly.\n        mask = (mask / 255.0 > 0.5).astype(np.float32)\n\n        # Apply transformations if they are provided.\n        if self.transform:\n            # Albumentations applies the same geometric transforms to both image and mask.\n            transformed = self.transform(image=image, mask=mask)\n            image = transformed['image']\n            mask = transformed['mask']\n            # Add a channel dimension to the mask for PyTorch [C, H, W] format.\n            mask = mask.unsqueeze(0)\n\n        return image, mask\n\n\n# --- 4.3. Instantiate Datasets and DataLoaders ---\n# We use the previously defined variables: BATCH_SIZE, NUM_WORKERS,\n# train_pairs, and val_pairs.\n\n# Create the training dataset instance\ntrain_dataset = SegmentationDataset(\n    image_mask_pairs=train_pairs,\n    transform=train_transforms\n)\n# Create the training data loader\ntrain_loader = DataLoader(\n    train_dataset,\n    batch_size=BATCH_SIZE,\n    shuffle=True,\n    num_workers=NUM_WORKERS\n)\n\n# Create the validation dataset instance\nval_dataset = SegmentationDataset(\n    image_mask_pairs=val_pairs,\n    transform=val_transforms\n)\n# Create the validation data loader\nval_loader = DataLoader(\n    val_dataset,\n    batch_size=BATCH_SIZE,\n    shuffle=False, # No need to shuffle validation data\n    num_workers=NUM_WORKERS\n)\n\n\n# --- 4.4. Verify a Sample Batch ---\n# It's good practice to check the output of the DataLoader to ensure\n# tensor shapes are correct before starting the training.\ntry:\n    print(\"DataLoaders created successfully.\")\n    images, masks = next(iter(train_loader))\n    print(f\"Image batch shape from train_loader: {images.shape}\")\n    print(f\"Mask batch shape from train_loader: {masks.shape}\")\n    # Expected shapes:\n    # Images: [BATCH_SIZE, 3, IMAGE_HEIGHT, IMAGE_WIDTH]\n    # Masks:  [BATCH_SIZE, 1, IMAGE_HEIGHT, IMAGE_WIDTH]\nexcept Exception as e:\n    print(f\"An error occurred while verifying the DataLoader: {e}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-26T06:22:26.813552Z","iopub.execute_input":"2025-06-26T06:22:26.814251Z","iopub.status.idle":"2025-06-26T06:22:27.990253Z","shell.execute_reply.started":"2025-06-26T06:22:26.814229Z","shell.execute_reply":"2025-06-26T06:22:27.989351Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ======================================================================================\n# 5. MODEL ARCHITECTURE\n# This cell defines the building blocks and the main U-Net model architecture.\n# ======================================================================================\n\n# --- 5.1. U-Net Building Block: Double Convolution ---\n\nclass DoubleConv(nn.Module):\n    \"\"\"A block consisting of two sequential convolutional layers,\n    each followed by BatchNorm and ReLU activation.\n    (Convolution => [BatchNorm] => ReLU) * 2\n    \"\"\"\n    def __init__(self, in_channels, out_channels):\n        super().__init__()\n        self.double_conv = nn.Sequential(\n            # First convolutional layer\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            # Second convolutional layer\n            nn.Conv2d(out_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\n\n# --- 5.2. The U-Net Model ---\n\nclass RollasUnet(nn.Module):\n    \"\"\"\n    The U-Net architecture for image segmentation.\n    It consists of an encoder (contracting path), a bottleneck,\n    and a decoder (expansive path) with skip connections.\n    \"\"\"\n    def __init__(self, in_channels=3, out_channels=1):\n        super(RollasUnet, self).__init__()\n        self.in_channels = in_channels\n        self.out_channels = out_channels\n\n        # --- Encoder (Contracting Path) ---\n        # Each 'down' block consists of a DoubleConv followed by a MaxPool.\n        self.inc = DoubleConv(in_channels, 64)\n        self.down1 = DoubleConv(64, 128)\n        self.down2 = DoubleConv(128, 256)\n        self.down3 = DoubleConv(256, 512)\n        self.pool = nn.MaxPool2d(2)\n\n        # --- Bottleneck ---\n        self.bottleneck = DoubleConv(512, 1024)\n\n        # --- Decoder (Expansive Path) ---\n        # Each 'up' step consists of an up-convolution, a skip connection concatenation,\n        # and a DoubleConv block.\n        self.up1 = nn.ConvTranspose2d(1024, 512, kernel_size=2, stride=2)\n        self.conv1 = DoubleConv(1024, 512) # Input is 512 (from upconv) + 512 (from skip)\n\n        self.up2 = nn.ConvTranspose2d(512, 256, kernel_size=2, stride=2)\n        self.conv2 = DoubleConv(512, 256) # Input is 256 + 256\n\n        self.up3 = nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2)\n        self.conv3 = DoubleConv(256, 128) # Input is 128 + 128\n\n        self.up4 = nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2)\n        self.conv4 = DoubleConv(128, 64) # Input is 64 + 64\n\n        # --- Final Output Layer ---\n        # This final 1x1 convolution maps the 64-channel feature map to the\n        # desired number of output channels (1 for binary segmentation).\n        self.outc = nn.Conv2d(64, out_channels, kernel_size=1)\n\n    def forward(self, x):\n        # --- Encoder Path ---\n        # Pass input through each encoder block and save the output\n        # for the skip connections before the pooling operation.\n        skip1 = self.inc(x)\n        skip2 = self.down1(self.pool(skip1))\n        skip3 = self.down2(self.pool(skip2))\n        skip4 = self.down3(self.pool(skip3))\n\n        # --- Bottleneck ---\n        x = self.bottleneck(self.pool(skip4))\n\n        # --- Decoder Path ---\n        # At each step, upsample, concatenate with the corresponding\n        # skip connection, and pass through the DoubleConv block.\n        x = self.up1(x)\n        x = torch.cat([x, skip4], dim=1)\n        x = self.conv1(x)\n\n        x = self.up2(x)\n        x = torch.cat([x, skip3], dim=1)\n        x = self.conv2(x)\n\n        x = self.up3(x)\n        x = torch.cat([x, skip2], dim=1)\n        x = self.conv3(x)\n\n        x = self.up4(x)\n        x = torch.cat([x, skip1], dim=1)\n        x = self.conv4(x)\n\n        # --- Final Output ---\n        # The output 'logits' are the raw scores from the model.\n        # They will be passed to a loss function like BCEWithLogitsLoss.\n        logits = self.outc(x)\n        return logits\n\n\n# --- 5.3. Verify Model Architecture ---\n# A quick test to ensure the model is wired correctly and the\n# output shape matches the input shape (except for channels).\n\n# Use the device defined in the first cell\ndevice = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n\n# Instantiate the model and move it to the configured device\nmodel = RollasUnet(in_channels=3, out_channels=1).to(device)\n\n# Create a dummy input tensor with the expected shape\ndummy_input = torch.randn(BATCH_SIZE, 3, IMAGE_HEIGHT, IMAGE_WIDTH).to(device)\n\n# Get the model's prediction\npreds = model(dummy_input)\n\n# Print the shapes to verify\nprint(\"--- Model Sanity Check ---\")\nprint(f\"Input tensor shape:  {dummy_input.shape}\")\nprint(f\"Output tensor shape: {preds.shape}\")\nprint(f\"Model successfully created and tested.\")\n# The output shape should be [BATCH_SIZE, 1, IMAGE_HEIGHT, IMAGE_WIDTH]\nLEARNING_RATE = 1e-4 # Learning rate\nNUM_EPOCHS = 25    # Number of epochs in neural network.","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-26T06:23:34.005065Z","iopub.execute_input":"2025-06-26T06:23:34.005942Z","iopub.status.idle":"2025-06-26T06:23:34.351561Z","shell.execute_reply.started":"2025-06-26T06:23:34.005916Z","shell.execute_reply":"2025-06-26T06:23:34.350814Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ======================================================================================\n# 6. TRAINING SETUP: Loss, Optimizer, and Helper Functions\n# This cell defines the components required for training the model.\n# ======================================================================================\n\n# --- 6.1. Custom Loss Function and Evaluation Metric ---\n\nclass DiceLoss(nn.Module):\n    \"\"\"\n    Implements Dice Loss, calculated as 1 - Dice Coefficient.\n    Useful for directly optimizing the Dice score metric.\n    \"\"\"\n    def __init__(self):\n        super(DiceLoss, self).__init__()\n\n    def forward(self, inputs, targets, smooth=1):\n        # Apply sigmoid to model outputs to get probabilities\n        inputs = torch.sigmoid(inputs)\n\n        # Flatten label and prediction tensors\n        inputs = inputs.view(-1)\n        targets = targets.view(-1)\n\n        intersection = (inputs * targets).sum()\n        dice_coeff = (2. * intersection + smooth) / (inputs.sum() + targets.sum() + smooth)\n\n        return 1 - dice_coeff\n\ndef dice_coefficient(preds, targets, smooth=1e-6):\n    \"\"\"\n    Calculates the Dice Coefficient metric for evaluation.\n    Args:\n        preds (torch.Tensor): The model's raw output logits.\n        targets (torch.Tensor): The ground truth masks.\n    Returns:\n        float: The Dice Coefficient score.\n    \"\"\"\n    preds = torch.sigmoid(preds)\n    preds = (preds > 0.5).float()\n    intersection = (preds * targets).sum()\n    dice = (2. * intersection + smooth) / (preds.sum() + targets.sum() + smooth)\n    return dice.item()\n\ndef combined_loss(pred, target):\n    \"\"\"\n    A combined loss function that balances between BCE and Dice loss.\n    This helps with both pixel-level accuracy and segmentation overlap.\n    \"\"\"\n    bce = nn.BCEWithLogitsLoss()\n    dice = DiceLoss()\n    # You can experiment with the weights, e.g., 0.8 for BCE and 0.2 for Dice\n    return 0.5 * bce(pred, target) + 0.5 * dice(pred, target)\n\n\n# ======================================================================================\n# 7. INITIALIZE MODEL, LOSS, AND OPTIMIZER\n# Prepare the main components for the training process.\n# ======================================================================================\n\n# Use variables defined in the first refactored cell\nDEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nLEARNING_RATE = 1e-4\nNUM_EPOCHS = 25 # You can increase this if the model is still learning\n\n# Instantiate the model (assuming UNET class is defined in a previous cell)\nmodel = RollasUnet(in_channels=3, out_channels=1).to(DEVICE)\n\n# Define the loss function and optimizer\nloss_fn = nn.BCEWithLogitsLoss()\noptimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)\n\n# Lists to store metrics for plotting later\ntrain_losses = []\nval_losses = []\nval_dice_scores = []\nbest_val_dice = 0.0 # Variable to track the best score for model saving\n\n\n# ======================================================================================\n# 8. EXECUTE THE TRAINING LOOP\n# ======================================================================================\n\nprint(\"--- Starting Training ---\")\n\nfor epoch in range(NUM_EPOCHS):\n    # --- Training Phase ---\n    model.train()\n    total_train_loss = 0.0\n    train_loop = tqdm(train_loader, desc=f\"Epoch {epoch+1}/{NUM_EPOCHS} [Training]\", leave=False)\n\n    for data, targets in train_loop:\n        data = data.to(device=DEVICE)\n        targets = targets.float().to(device=DEVICE)\n\n        # Forward pass\n        predictions = model(data)\n        loss = loss_fn(predictions, targets)\n        total_train_loss += loss.item()\n\n        # Backward pass and optimization\n        optimizer.zero_grad()\n        loss.backward()\n        optimizer.step()\n\n        # Update progress bar\n        train_loop.set_postfix(loss=loss.item())\n\n    avg_train_loss = total_train_loss / len(train_loader)\n    train_losses.append(avg_train_loss)\n\n    # --- Validation Phase ---\n    model.eval()\n    total_val_dice = 0.0\n    total_val_loss = 0.0\n\n    with torch.no_grad():\n        for data, targets in val_loader:\n            data = data.to(device=DEVICE)\n            targets = targets.float().to(device=DEVICE)\n            predictions = model(data)\n\n            # Calculate validation loss (using the same combined loss)\n            val_loss = loss_fn(predictions, targets)\n            total_val_loss += val_loss.item()\n\n            # Calculate validation Dice score\n            dice_score = dice_coefficient(predictions, targets)\n            total_val_dice += dice_score\n\n    avg_val_loss = total_val_loss / len(val_loader)\n    avg_val_dice = total_val_dice / len(val_loader)\n    val_losses.append(avg_val_loss)\n    val_dice_scores.append(avg_val_dice)\n\n    # --- Save the best model based on validation Dice score ---\n    if avg_val_dice > best_val_dice:\n        best_val_dice = avg_val_dice\n        torch.save(model.state_dict(), \"best_model.pth\")\n        print(f\"--> New best model saved with Dice score: {avg_val_dice:.4f}\")\n\n    # --- Print Epoch Summary ---\n    print(f\"Epoch [{epoch+1}/{NUM_EPOCHS}] | \"\n          f\"Train Loss: {avg_train_loss:.4f} | \"\n          f\"Val Loss: {avg_val_loss:.4f} | \"\n          f\"Val Dice: {avg_val_dice:.4f}\")\n\nprint(\"\\n--- Training Finished ---\")\nprint(f\"Best validation Dice score achieved: {best_val_dice:.4f}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-26T06:23:59.462977Z","iopub.execute_input":"2025-06-26T06:23:59.463589Z","iopub.status.idle":"2025-06-26T06:49:21.867752Z","shell.execute_reply.started":"2025-06-26T06:23:59.463567Z","shell.execute_reply":"2025-06-26T06:49:21.866803Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ======================================================================================\n# 9. VISUALIZE PREDICTIONS\n# This cell defines a function to visualize the model's predictions on a few\n# samples from a given data loader.\n# ======================================================================================\n\ndef visualize_predictions(model, loader, device, num_images_to_show=3):\n    \"\"\"\n    Visualizes model predictions against the ground truth masks.\n\n    Args:\n        model (torch.nn.Module): The trained model.\n        loader (DataLoader): The DataLoader to get samples from (e.g., val_loader).\n        device (str): The device to run the model on (\"cuda\" or \"cpu\").\n        num_images_to_show (int): The number of sample images to display.\n    \"\"\"\n    print(\"\\n--- Visualizing a few predictions... ---\")\n    model.eval() # Set the model to evaluation mode\n\n    # Get a single batch of data from the loader\n    try:\n        images, masks = next(iter(loader))\n    except StopIteration:\n        print(\"DataLoader is empty. Cannot visualize.\")\n        return\n\n    images = images.to(device)\n\n    with torch.no_grad():\n        # Get model predictions\n        preds = torch.sigmoid(model(images))\n        # Apply threshold to get binary masks\n        preds = (preds > 0.5).float()\n\n    # Move tensors to CPU for plotting and convert to numpy arrays\n    images = images.cpu().numpy()\n    masks = masks.cpu().numpy()\n    preds = preds.cpu().numpy()\n\n    # Determine how many images to show from the batch\n    num_to_show = min(num_images_to_show, len(images))\n\n    # Plot the images, ground truth masks, and predicted masks\n    fig, axs = plt.subplots(num_to_show, 3, figsize=(15, num_to_show * 5))\n\n    for i in range(num_to_show):\n        # --- Denormalize the input image for better visualization ---\n        # NOTE: These are the std and mean used in the transforms.\n        # If you changed them, you must update these values too.\n        mean = np.array([0.485, 0.456, 0.406])\n        std = np.array([0.229, 0.224, 0.225])\n\n        # Transpose image from [C, H, W] to [H, W, C] for matplotlib\n        img_display = np.transpose(images[i], (1, 2, 0))\n\n        # Apply the reverse transformation: (pixel * std) + mean\n        img_display = std * img_display + mean\n\n        # Clip values to [0, 1] range to handle potential floating point inaccuracies\n        img_display = np.clip(img_display, 0, 1)\n\n        # --- Plotting ---\n        # If only one image is shown, axs is not a 2D array\n        ax_row = axs[i] if num_to_show > 1 else axs\n\n        # Original Image\n        ax_row[0].imshow(img_display)\n        ax_row[0].set_title(\"Input Image\")\n        ax_row[0].axis(\"off\")\n\n        # Ground Truth Mask\n        ax_row[1].imshow(masks[i].squeeze(), cmap=\"gray\")\n        ax_row[1].set_title(\"Ground Truth Mask\")\n        ax_row[1].axis(\"off\")\n\n        # Predicted Mask\n        ax_row[2].imshow(preds[i].squeeze(), cmap=\"gray\")\n        ax_row[2].set_title(\"Predicted Mask\")\n        ax_row[2].axis(\"off\")\n\n    plt.tight_layout()\n    plt.show()\n\n\n# --- Execute the visualization ---\n# We call the function using the trained model and the validation loader.\n# The model and val_loader should already be defined and available in the notebook.\nvisualize_predictions(model, val_loader, device=DEVICE)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-26T06:50:09.870893Z","iopub.execute_input":"2025-06-26T06:50:09.871233Z","iopub.status.idle":"2025-06-26T06:50:12.333675Z","shell.execute_reply.started":"2025-06-26T06:50:09.871205Z","shell.execute_reply":"2025-06-26T06:50:12.332772Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\n\nclass TestDataset(Dataset):\n    \"\"\"\n    Custom PyTorch Dataset for loading test images.\n    It returns the transformed image, its filename, and its original dimensions.\n    \"\"\"\n    def __init__(self, test_dir, transform=None):\n        self.test_dir = test_dir\n        self.image_names = sorted(os.listdir(test_dir))\n        self.transform = transform\n\n    def __len__(self):\n        return len(self.image_names)\n\n    def __getitem__(self, idx):\n        image_name = self.image_names[idx]\n        image_path = os.path.join(self.test_dir, image_name)\n\n        image = cv2.imread(image_path)\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n\n        original_height, original_width, _ = image.shape\n\n        if self.transform:\n            transformed = self.transform(image=image)\n            image = transformed['image']\n\n        return image, image_name, (original_height, original_width)\n\n\n# --- 10.2. Main Inference and Visualization Function ---\n\ndef generate_submission_and_visualize(model, device, test_dir, num_to_show=5):\n    \"\"\"\n    Uses the trained model to generate predictions on the test set,\n    creates a submission file, and visualizes some results.\n    \"\"\"\n    # Define transformations for the test set (no augmentation)\n    test_transforms = A.Compose([\n        A.Resize(height=IMAGE_HEIGHT, width=IMAGE_WIDTH),\n        A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225], max_pixel_value=255.0),\n        ToTensorV2(),\n    ])\n\n    # Create the test dataset and dataloader\n    test_dataset = TestDataset(test_dir, transform=test_transforms)\n    test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=NUM_WORKERS)\n\n    # Set model to evaluation mode\n    model.eval()\n\n    # Lists to store results\n    \n    images_to_display = []\n    masks_to_display = []\n\n    print(\"\\n--- Generating predictions on the test set... ---\")\n    with torch.no_grad():\n        for images, image_names, original_dims in tqdm(test_loader, desc=\"Predicting\"):\n            images = images.to(device)\n\n            # Get predictions (logits) and convert to binary masks\n            preds = torch.sigmoid(model(images))\n            preds = (preds > 0.5).cpu().numpy()\n\n            for i in range(len(preds)):\n                # --- Resize mask to original image dimensions ---\n                pred_mask_256 = preds[i].squeeze().astype(np.uint8)\n                original_h, original_w = original_dims[0][i].item(), original_dims[1][i].item()\n                resized_mask = cv2.resize(pred_mask_256, (original_w, original_h), interpolation=cv2.INTER_NEAREST)\n\n                \n\n                # Store some samples for visualization\n                if len(images_to_display) < num_to_show:\n                    original_image_path = os.path.join(test_dir, image_names[i])\n                    images_to_display.append(original_image_path)\n                    masks_to_display.append(resized_mask)\n\n    \n\n    # --- Visualize the stored samples ---\n    print(\"\\n--- Visualizing a few test predictions (in original size) ---\")\n    for i in range(len(images_to_display)):\n        original_image = cv2.imread(images_to_display[i])\n        original_image = cv2.cvtColor(original_image, cv2.COLOR_BGR2RGB)\n        predicted_mask = masks_to_display[i]\n\n        fig, axs = plt.subplots(1, 2, figsize=(15, 7))\n        axs[0].imshow(original_image)\n        axs[0].set_title(f\"Original Test Image: {os.path.basename(images_to_display[i])}\")\n        axs[0].axis('off')\n\n        axs[1].imshow(predicted_mask, cmap='gray')\n        axs[1].set_title(\"Predicted Mask (Resized)\")\n        axs[1].axis('off')\n        plt.show()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-26T06:51:45.905255Z","iopub.execute_input":"2025-06-26T06:51:45.905957Z","iopub.status.idle":"2025-06-26T06:51:45.916883Z","shell.execute_reply.started":"2025-06-26T06:51:45.905933Z","shell.execute_reply":"2025-06-26T06:51:45.916132Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Instantiate the model architecture\nmodel = RollasUnet(in_channels=3, out_channels=1).to(DEVICE)\n\n# Load the weights of the best performing model saved during training\nprint(\"Loading best model weights from 'best_model.pth'...\")\ntry:\n    model.load_state_dict(torch.load(\"best_model.pth\", map_location=DEVICE))\nexcept FileNotFoundError:\n    print(\"Error: 'best_model.pth' not found. Please ensure the model was trained and saved correctly.\")\n    # Handle error, maybe by stopping or using the current model in memory\nelse:\n    # Run the submission and visualization pipeline\n    generate_submission_and_visualize(model, DEVICE, TEST_DIR)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-26T06:52:09.375156Z","iopub.execute_input":"2025-06-26T06:52:09.375453Z","iopub.status.idle":"2025-06-26T06:52:16.835614Z","shell.execute_reply.started":"2025-06-26T06:52:09.375431Z","shell.execute_reply":"2025-06-26T06:52:16.834813Z"}},"outputs":[],"execution_count":null}]}