{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":103103,"databundleVersionId":13042974,"sourceType":"competition"}],"dockerImageVersionId":31091,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-09-05T13:08:07.381897Z","iopub.execute_input":"2025-09-05T13:08:07.382710Z","iopub.status.idle":"2025-09-05T13:08:10.927403Z","shell.execute_reply.started":"2025-09-05T13:08:07.382684Z","shell.execute_reply":"2025-09-05T13:08:10.926643Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport sys\nimport time\nimport glob\nimport cv2\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nimport seaborn as sns\nfrom pathlib import Path\nfrom tqdm.auto import tqdm\nimport yaml\nimport random\nfrom PIL import Image\nimport warnings\nwarnings.filterwarnings('ignore')\n\n# Disable wandb\nos.environ['WANDB_DISABLED'] = 'true'\n\n# Set random seeds for reproducibility\nrandom.seed(42)\nnp.random.seed(42)\n\n# Install required packages\nprint(\"Installing required packages...\")\nos.system('pip install -q segmentation-models-pytorch')\nos.system('pip install -q albumentations')\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nimport segmentation_models_pytorch as smp\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\nfrom sklearn.model_selection import train_test_split\n\n# Set deterministic behavior for PyTorch\ntorch.manual_seed(42)\nif torch.cuda.is_available():\n    torch.cuda.manual_seed(42)\n    torch.backends.cudnn.deterministic = True\n    torch.backends.cudnn.benchmark = False\n\nprint(f'\\nPyTorch Version: {torch.__version__}')\nprint(f'CUDA Available: {torch.cuda.is_available()}')\nif torch.cuda.is_available():\n    print(f'CUDA Device: {torch.cuda.get_device_name(0)}')\n\n# Configuration Class\nclass Config:\n    # Dataset paths\n    BASE_PATH = '/kaggle/input/alpha-dent/AlphaDent'\n    OUTPUT_DIR = '/kaggle/working/'\n    \n    # Image preprocessing\n    XRAY_SIZE = 512\n    CLAHE_CLIP_LIMIT = 3.0\n    CLAHE_TILE_SIZE = (8, 8)\n    BLUR_KERNEL = (3, 3)\n    \n    # Training parameters\n    BATCH_SIZE = 8 if torch.cuda.is_available() else 4\n    EPOCHS = 20\n    LEARNING_RATE = 1e-4\n    WEIGHT_DECAY = 1e-5\n    NUM_CLASSES = 9\n    \n    # Model parameters\n    ENCODER = 'efficientnet-b4'\n    ENCODER_WEIGHTS = 'imagenet'\n    \n    # Training settings\n    PATIENCE = 10\n    DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n\n# Define paths\nTRAIN_IMAGES_PATH = f'{Config.BASE_PATH}/images/train'\nVALID_IMAGES_PATH = f'{Config.BASE_PATH}/images/valid'\nTEST_IMAGES_PATH = f'{Config.BASE_PATH}/images/test'\nTRAIN_LABELS_PATH = f'{Config.BASE_PATH}/labels/train'\nVALID_LABELS_PATH = f'{Config.BASE_PATH}/labels/valid'\n\n# Output paths\nos.makedirs(Config.OUTPUT_DIR, exist_ok=True)\n\n# Define class information\nCLASS_INFO = {\n    0: {'name': 'Abrasion', 'description': 'Teeth with mechanical wear of hard tissues'},\n    1: {'name': 'Filling', 'description': 'Dental fillings of various types'},\n    2: {'name': 'Crown', 'description': 'Dental crown (restoration)'},\n    3: {'name': 'Caries Class 1', 'description': 'Caries in fissures and pits'},\n    4: {'name': 'Caries Class 2', 'description': 'Caries on proximal surfaces of molars/premolars'},\n    5: {'name': 'Caries Class 3', 'description': 'Caries on proximal surfaces of incisors/canines without incisal edge'},\n    6: {'name': 'Caries Class 4', 'description': 'Caries on proximal surfaces of incisors/canines with incisal edge'},\n    7: {'name': 'Caries Class 5', 'description': 'Cervical caries (buccal/lingual surfaces)'},\n    8: {'name': 'Caries Class 6', 'description': 'Caries on incisal edges or cusps'}\n}\n\nfrom sklearn.metrics import average_precision_score\n\ndef compute_map50(y_true_masks, y_pred_masks):\n    \"\"\"\n    Compute mAP@50 from segmentation masks.\n    - y_true_masks: list of ground-truth binary masks (numpy arrays)\n    - y_pred_masks: list of predicted binary masks (numpy arrays, same shape as true masks)\n    \"\"\"\n    aps = []\n    for y_true, y_pred in zip(y_true_masks, y_pred_masks):\n        # Flatten masks\n        y_true_flat = y_true.flatten()\n        y_pred_flat = y_pred.flatten()\n\n        # Compute IoU\n        intersection = (y_true_flat & y_pred_flat).sum()\n        union = (y_true_flat | y_pred_flat).sum()\n        iou = intersection / union if union > 0 else 0.0\n\n        # Binary \"detection success\" at IoU≥0.5\n        success = 1 if iou >= 0.5 else 0\n\n        # Average precision needs probs, so we just use IoU as confidence\n        aps.append(success)\n\n    # mAP@50 is mean of all APs\n    return sum(aps) / len(aps) if aps else 0.0\n\n\n# Preprocessing functions\ndef preprocess_xray(image, augment=False):\n    \"\"\"\n    Enhanced X-ray preprocessing for dental images\n    Args:\n        image: RGB numpy array\n        augment: Whether to apply augmentation\n    Returns:\n        Preprocessed grayscale image as float32\n    \"\"\"\n    # --- Standardization: resize + grayscale ---\n    image = cv2.resize(image, (Config.XRAY_SIZE, Config.XRAY_SIZE))\n    gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)\n    \n    # --- Invert intensities (mimic radiodensity: teeth = white, background = dark) ---\n    gray = cv2.bitwise_not(gray)\n    \n    # --- Enhance local contrast (CLAHE) ---\n    clahe = cv2.createCLAHE(clipLimit=Config.CLAHE_CLIP_LIMIT, tileGridSize=Config.CLAHE_TILE_SIZE)\n    gray = clahe.apply(gray)\n    \n    # --- Add soft blur to mimic X-ray scattering ---\n    gray = cv2.GaussianBlur(gray, Config.BLUR_KERNEL, 0)\n    \n    # --- Normalize intensities 0-1 ---\n    gray = cv2.normalize(gray, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8)\n\n    \n    # --- Pose / Angle: optional augmentation ---\n    if augment:\n        # small rotation\n        angle = np.random.uniform(-10, 10)\n        center = (Config.XRAY_SIZE // 2, Config.XRAY_SIZE // 2)\n        M = cv2.getRotationMatrix2D(center, angle, 1)\n        gray = cv2.warpAffine(gray, M, (Config.XRAY_SIZE, Config.XRAY_SIZE), \n                             borderMode=cv2.BORDER_CONSTANT, borderValue=0)\n        # random flip\n        if np.random.rand() > 0.5:\n            gray = cv2.flip(gray, 1)\n    \n    return gray\n\ndef preprocess_xray_pil(pil_img):\n    \"\"\"Convert PIL image to numpy, apply X-ray preprocessing, return back PIL\"\"\"\n    img = np.array(pil_img.convert(\"RGB\"))  # ensure RGB\n    img = cv2.resize(img, (Config.XRAY_SIZE, Config.XRAY_SIZE))\n    gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)\n    gray = cv2.bitwise_not(gray)  # invert\n    clahe = cv2.createCLAHE(clipLimit=Config.CLAHE_CLIP_LIMIT, tileGridSize=Config.CLAHE_TILE_SIZE)\n    gray = clahe.apply(gray)\n    gray = cv2.GaussianBlur(gray, Config.BLUR_KERNEL, 0)\n    gray = cv2.normalize(gray, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8)\n    return Image.fromarray(gray)  # back to PIL for plotting\n\ndef parse_yolo_labels(label_path, img_height, img_width):\n    \"\"\"Convert YOLO format labels to segmentation masks\"\"\"\n    mask = np.zeros((img_height, img_width), dtype=np.uint8)\n    \n    if os.path.exists(label_path) and os.path.getsize(label_path) > 0:\n        try:\n            with open(label_path, 'r') as f:\n                lines = f.readlines()\n                \n            for line in lines:\n                parts = line.strip().split()\n                if len(parts) >= 5:  # At least class + 2 points (4 coordinates)\n                    class_id = int(parts[0]) + 1  # Add 1 to make classes 1-9 (0 is background)\n                    \n                    # Parse polygon coordinates\n                    coords = list(map(float, parts[1:]))\n                    if len(coords) >= 6 and len(coords) % 2 == 0:  # At least 3 points\n                        # Convert normalized coordinates to pixel coordinates\n                        polygon_points = []\n                        for i in range(0, len(coords), 2):\n                            x = int(coords[i] * img_width)\n                            y = int(coords[i + 1] * img_height)\n                            polygon_points.append([x, y])\n                        \n                        # Fill polygon\n                        polygon_points = np.array(polygon_points, dtype=np.int32)\n                        cv2.fillPoly(mask, [polygon_points], class_id)\n        except Exception as e:\n            print(f\"Error parsing label {label_path}: {e}\")\n    \n    return mask\n\n# Custom Dataset\nclass DentalSegmentationDataset(Dataset):\n    def __init__(self, image_paths, label_paths=None, transform=None, is_test=False):\n        self.image_paths = image_paths\n        self.label_paths = label_paths\n        self.transform = transform\n        self.is_test = is_test\n    \n    def __len__(self):\n        return len(self.image_paths)\n    \n    def __getitem__(self, idx):\n        # Load image\n        image = cv2.imread(self.image_paths[idx])\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n        \n        # Apply X-ray preprocessing\n        processed_image = preprocess_xray(image, augment=not self.is_test)\n        \n        # Convert to 3-channel for pretrained encoders\n        processed_image = np.stack([processed_image] * 3, axis=-1)\n        \n        if self.is_test:\n            if self.transform:\n                transformed = self.transform(image=processed_image)\n                processed_image = transformed['image']\n            return processed_image, os.path.basename(self.image_paths[idx])\n        \n        # Load mask\n        if self.label_paths is not None:\n            label_path = self.label_paths[idx]\n            mask = parse_yolo_labels(label_path, Config.XRAY_SIZE, Config.XRAY_SIZE)\n        else:\n            mask = np.zeros((Config.XRAY_SIZE, Config.XRAY_SIZE), dtype=np.uint8)\n        \n        if self.transform:\n            transformed = self.transform(image=processed_image, mask=mask)\n            processed_image = transformed['image']\n            mask = transformed['mask']\n        \n        return processed_image, mask.long()\n\n# Data transforms\ndef get_train_transforms():\n    return A.Compose([\n        A.HorizontalFlip(p=0.5),\n        A.ShiftScaleRotate(shift_limit=0.1, scale_limit=0.1, rotate_limit=15, p=0.5),\n        A.RandomBrightnessContrast(brightness_limit=0.2, contrast_limit=0.2, p=0.5),\n        A.ElasticTransform(p=0.3),\n        A.GridDistortion(p=0.3),\n        A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n        ToTensorV2()\n    ])\n\ndef get_valid_transforms():\n    return A.Compose([\n        A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n        ToTensorV2()\n    ])\n\n# Loss functions\nclass CombinedLoss(nn.Module):\n    def __init__(self, alpha=0.5, beta=0.5):\n        super(CombinedLoss, self).__init__()\n        self.alpha = alpha\n        self.beta = beta\n        self.dice_loss = smp.losses.DiceLoss(mode='multiclass')\n        self.focal_loss = smp.losses.FocalLoss(mode='multiclass')\n    \n    def forward(self, outputs, targets):\n        dice = self.dice_loss(outputs, targets)\n        focal = self.focal_loss(outputs, targets)\n        return self.alpha * dice + self.beta * focal\n\n# Metrics\ndef dice_coefficient(pred, target, num_classes):\n    dice_scores = []\n    for class_idx in range(1, num_classes + 1):  # Skip background\n        pred_class = (pred == class_idx).float()\n        target_class = (target == class_idx).float()\n        \n        intersection = (pred_class * target_class).sum()\n        union = pred_class.sum() + target_class.sum()\n        \n        if union == 0:\n            dice = 1.0 if intersection == 0 else 0.0\n        else:\n            dice = (2.0 * intersection) / union\n        \n        # Always convert to float\n        dice_scores.append(float(dice))\n    \n    return np.mean(dice_scores)\n\n\n# Training function\ndef train_epoch(model, train_loader, criterion, optimizer, device):\n    model.train()\n    running_loss = 0.0\n    running_dice = 0.0\n    \n    for batch_idx, (images, masks) in enumerate(tqdm(train_loader, desc=\"Training\")):\n        images = images.to(device)\n        masks = masks.to(device)\n        \n        optimizer.zero_grad()\n        outputs = model(images)\n        loss = criterion(outputs, masks)\n        \n        loss.backward()\n        optimizer.step()\n        \n        running_loss += loss.item()\n        \n        # Calculate dice score\n        pred_masks = torch.argmax(outputs, dim=1)\n        dice = dice_coefficient(pred_masks, masks, Config.NUM_CLASSES)\n        running_dice += dice\n    \n    epoch_loss = running_loss / len(train_loader)\n    epoch_dice = running_dice / len(train_loader)\n    \n    return epoch_loss, epoch_dice\n\ndef validate_epoch(model, valid_loader, criterion, device):\n    model.eval()\n    running_loss = 0.0\n    running_dice = 0.0\n    \n    with torch.no_grad():\n        for images, masks in tqdm(valid_loader, desc=\"Validation\"):\n            images = images.to(device)\n            masks = masks.to(device)\n            \n            outputs = model(images)\n            loss = criterion(outputs, masks)\n            \n            running_loss += loss.item()\n            \n            # Calculate dice score\n            pred_masks = torch.argmax(outputs, dim=1)\n            dice = dice_coefficient(pred_masks, masks, Config.NUM_CLASSES)\n            running_dice += dice\n    \n    epoch_loss = running_loss / len(valid_loader)\n    epoch_dice = running_dice / len(valid_loader)\n    \n    return epoch_loss, epoch_dice\n\n# Data preparation\nprint(\"\\n=== Preparing Dataset ===\")\ntrain_images = sorted(glob.glob(f'{TRAIN_IMAGES_PATH}/*.jpg'))\nvalid_images = sorted(glob.glob(f'{VALID_IMAGES_PATH}/*.jpg'))\ntest_images = sorted(glob.glob(f'{TEST_IMAGES_PATH}/*.jpg'))\n\n# Get corresponding label files\ntrain_labels = []\nfor img_path in train_images:\n    img_name = os.path.basename(img_path).replace('.jpg', '.txt')\n    label_path = os.path.join(TRAIN_LABELS_PATH, img_name)\n    train_labels.append(label_path)\n\nvalid_labels = []\nfor img_path in valid_images:\n    img_name = os.path.basename(img_path).replace('.jpg', '.txt')\n    label_path = os.path.join(VALID_LABELS_PATH, img_name)\n    valid_labels.append(label_path)\n\nprint(f\"Training images: {len(train_images)}\")\nprint(f\"Validation images: {len(valid_images)}\")\nprint(f\"Test images: {len(test_images)}\")\n\n# Create datasets\ntrain_dataset = DentalSegmentationDataset(\n    train_images, train_labels, \n    transform=get_train_transforms(),\n    is_test=False\n)\n\nvalid_dataset = DentalSegmentationDataset(\n    valid_images, valid_labels,\n    transform=get_valid_transforms(),\n    is_test=False\n)\n\ntest_dataset = DentalSegmentationDataset(\n    test_images,\n    transform=get_valid_transforms(),\n    is_test=True\n)\n\n# Create data loaders\ntrain_loader = DataLoader(\n    train_dataset, \n    batch_size=Config.BATCH_SIZE, \n    shuffle=True, \n    num_workers=0\n)\n\nvalid_loader = DataLoader(\n    valid_dataset, \n    batch_size=Config.BATCH_SIZE, \n    shuffle=False, \n    num_workers=0\n)\n\ntest_loader = DataLoader(\n    test_dataset, \n    batch_size=Config.BATCH_SIZE, \n    shuffle=False, \n    num_workers=0\n)\n\n# Initialize model\nprint(\"\\n=== Initializing U-Net++ Model ===\")\nmodel = smp.UnetPlusPlus(\n    encoder_name=Config.ENCODER,\n    encoder_weights=Config.ENCODER_WEIGHTS,\n    in_channels=3,\n    classes=Config.NUM_CLASSES + 1,  # +1 for background\n    activation=None,\n)\n\nmodel = model.to(Config.DEVICE)\n\n# Loss and optimizer\ncriterion = CombinedLoss()\noptimizer = torch.optim.AdamW(\n    model.parameters(), \n    lr=Config.LEARNING_RATE, \n    weight_decay=Config.WEIGHT_DECAY\n)\n\nscheduler = torch.optim.lr_scheduler.CosineAnnealingLR(\n    optimizer, T_max=Config.EPOCHS\n)\n\n# Training loop\nprint(\"\\n=== Starting Training ===\")\nbest_dice = 0.0\npatience_counter = 0\ntrain_losses = []\ntrain_dices = []\nvalid_losses = []\nvalid_dices = []\n\nfor epoch in range(Config.EPOCHS):\n    print(f\"\\nEpoch {epoch+1}/{Config.EPOCHS}\")\n    \n    # Train\n    train_loss, train_dice = train_epoch(model, train_loader, criterion, optimizer, Config.DEVICE)\n    \n    # Validate\n    valid_loss, valid_dice = validate_epoch(model, valid_loader, criterion, Config.DEVICE)\n    \n    # Update learning rate\n    scheduler.step()\n    \n    # Store metrics\n    train_losses.append(train_loss)\n    train_dices.append(train_dice)\n    valid_losses.append(valid_loss)\n    valid_dices.append(valid_dice)\n\n    print(f\"Valid Loss: {valid_loss:.4f}, Valid Dice: {valid_dice:.4f}\")\n    print(f\"Train Loss: {train_loss:.4f}, Train Dice: {train_dice:.4f}\")\n    print(f\"Valid Loss: {valid_loss:.4f}, Valid Dice: {valid_dice:.4f}\")\n    \n    # Save best model\n    if valid_dice > best_dice:\n        best_dice = valid_dice\n        patience_counter = 0\n        torch.save({\n            'epoch': epoch,\n            'model_state_dict': model.state_dict(),\n            'optimizer_state_dict': optimizer.state_dict(),\n            'best_dice': best_dice,\n        }, f'{Config.OUTPUT_DIR}/best_model.pth')\n        print(f\"New best model saved with dice: {best_dice:.4f}\")\n    else:\n        patience_counter += 1\n        \n    if patience_counter >= Config.PATIENCE:\n        print(f\"Early stopping after {epoch+1} epochs\")\n        break\n\nprint(f\"\\nTraining completed! Best validation dice: {best_dice:.4f}\")\n\n# Load best model for inference\nprint(\"\\n=== Loading Best Model for Inference ===\")\n# Load best model for inference (fix for PyTorch >=2.6)\ncheckpoint = torch.load(\n    f'{Config.OUTPUT_DIR}/best_model.pth',\n    map_location=Config.DEVICE,\n    weights_only=False  # <-- allow full pickle loading\n)\nmodel.load_state_dict(checkpoint['model_state_dict'])\nmodel.eval()\n\nprint(f\"Loaded model with:\")\nprint(f\"  Best Dice: {checkpoint['best_dice']:.4f}\")\nprint(f\"  Best mAP@50: {checkpoint.get('best_map', 0.0):.4f}\")\n\nif 'class_aps' in checkpoint and checkpoint['class_aps'] is not None:\n    print(\"  Per-class AP:\")\n    for i, ap in enumerate(checkpoint['class_aps']):\n        print(f\"    {CLASS_INFO[i]['name']}: {ap:.4f}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-05T13:08:10.928954Z","iopub.execute_input":"2025-09-05T13:08:10.929474Z","iopub.status.idle":"2025-09-05T15:10:09.389871Z","shell.execute_reply.started":"2025-09-05T13:08:10.929448Z","shell.execute_reply":"2025-09-05T15:10:09.389170Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Final validation\nprint(\"\\n=== Final Model Evaluation ===\")\nfinal_loss, final_dice = validate_epoch(model, valid_loader, criterion, Config.DEVICE)\n\nprint(f\"Final Validation Results:\")\nprint(f\"  Loss: {final_loss:.4f}\")\nprint(f\"  Dice Score: {final_dice:.4f}\")\n\n# Inference function\ndef mask_to_polygons(mask, min_area=100):\n    \"\"\"Convert segmentation mask to polygon format\"\"\"\n    polygons = []\n    \n    for class_id in range(1, Config.NUM_CLASSES + 1):  # Skip background (0)\n        class_mask = (mask == class_id).astype(np.uint8)\n        \n        # Find contours\n        contours, _ = cv2.findContours(class_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)\n        \n        for contour in contours:\n            area = cv2.contourArea(contour)\n            if area > min_area:\n                # Simplify contour\n                epsilon = 0.01 * cv2.arcLength(contour, True)\n                simplified_contour = cv2.approxPolyDP(contour, epsilon, True)\n                \n                if len(simplified_contour) >= 3:  # Valid polygon\n                    # Convert to normalized coordinates\n                    normalized_coords = []\n                    for point in simplified_contour:\n                        x, y = point[0]\n                        x_norm = x / Config.XRAY_SIZE\n                        y_norm = y / Config.XRAY_SIZE\n                        # Ensure coordinates are within [0, 1]\n                        x_norm = max(0.0, min(1.0, x_norm))\n                        y_norm = max(0.0, min(1.0, y_norm))\n                        normalized_coords.extend([x_norm, y_norm])\n                    \n                    if len(normalized_coords) >= 6:  # At least 3 points\n                        poly_str = ' '.join([f'{coord:.6f}' for coord in normalized_coords])\n                        polygons.append({\n                            'class_id': class_id - 1,  # Convert back to 0-based indexing\n                            'confidence': 0.8,  # Default confidence\n                            'poly': poly_str\n                        })\n    \n    return polygons\n\n# Run inference on test set\nprint(\"\\n=== Running Inference on Test Set ===\")\nall_predictions = []\n\nwith torch.no_grad():\n    for batch_idx, (images, filenames) in enumerate(tqdm(test_loader, desc=\"Inference\")):\n        images = images.to(Config.DEVICE)\n        outputs = model(images)\n        predictions = torch.argmax(outputs, dim=1).cpu().numpy()\n        \n        for i, (pred_mask, filename) in enumerate(zip(predictions, filenames)):\n            patient_id = filename.replace('.jpg', '')\n            \n            # Convert mask to polygons\n            polygons = mask_to_polygons(pred_mask)\n            \n            if not polygons:  # No detections, add dummy prediction\n                polygons = [{\n                    'patient_id': patient_id,\n                    'class_id': 0,\n                    'confidence': 0.01,\n                    'poly': '0.1 0.1 0.1 0.2 0.2 0.2 0.2 0.1'\n                }]\n            \n            for poly_data in polygons:\n                poly_data['patient_id'] = patient_id\n                all_predictions.append(poly_data)\n\nprint(f\"Generated {len(all_predictions)} predictions\")\n\n# Create submission\nprint(\"\\n=== Creating Submission File ===\")\nsubmission_df = pd.DataFrame(all_predictions)\n\n# Ensure all test images have predictions\nall_test_ids = [os.path.basename(img).replace('.jpg', '') for img in test_images]\nif len(submission_df) > 0:\n    predicted_ids = submission_df['patient_id'].unique()\n    missing_ids = set(all_test_ids) - set(predicted_ids)\nelse:\n    missing_ids = set(all_test_ids)\n\n# Add dummy predictions for missing images\nif missing_ids:\n    print(f\"Adding dummy predictions for {len(missing_ids)} images\")\n    dummy_rows = []\n    for img_id in missing_ids:\n        dummy_rows.append({\n            'patient_id': img_id,\n            'class_id': 0,\n            'confidence': 0.01,\n            'poly': '0.1 0.1 0.1 0.2 0.2 0.2 0.2 0.1'\n        })\n    \n    submission_df = pd.concat([submission_df, pd.DataFrame(dummy_rows)], ignore_index=True)\n\n# Sort and save\nsubmission_df = submission_df.sort_values(['patient_id', 'confidence'], ascending=[True, False])\nsubmission_df = submission_df[['patient_id', 'class_id', 'confidence', 'poly']]\nsubmission_df.to_csv('submission_unet++.csv', index=False)\n\nprint(f\"\\nSubmission created with {len(submission_df)} predictions\")\nprint(f\"Unique images: {submission_df['patient_id'].nunique()}\")\nprint(f\"All test images included: {submission_df['patient_id'].nunique() == len(test_images)}\")\n\nprint(\"\\n=== U-Net++ Pipeline Completed Successfully! ===\")\nprint(\"Files created:\")\nprint(\"  - best_model.pth (trained model)\")\nprint(\"  - submission_unet++.csv (submission file)\")\nprint(\"\\nReady to submit to Kaggle!\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-05T15:12:55.573401Z","iopub.execute_input":"2025-09-05T15:12:55.574045Z","iopub.status.idle":"2025-09-05T15:13:37.420643Z","shell.execute_reply.started":"2025-09-05T15:12:55.574019Z","shell.execute_reply":"2025-09-05T15:13:37.420056Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"submission_df.to_csv(\"/kaggle/working/submission.csv\", index=False)\nfrom IPython.display import FileLink\nFileLink(\"submission_unet++.csv\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-05T15:14:01.545645Z","iopub.execute_input":"2025-09-05T15:14:01.546430Z","iopub.status.idle":"2025-09-05T15:14:01.554856Z","shell.execute_reply.started":"2025-09-05T15:14:01.546396Z","shell.execute_reply":"2025-09-05T15:14:01.554264Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"img, mask = train_dataset[0]\nplt.subplot(1,2,1)\nplt.imshow(img.permute(1,2,0))   # image\nplt.subplot(1,2,2)\nplt.imshow(mask)                 # mask\nplt.show()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-05T15:14:14.011870Z","iopub.execute_input":"2025-09-05T15:14:14.012115Z","iopub.status.idle":"2025-09-05T15:14:14.466466Z","shell.execute_reply.started":"2025-09-05T15:14:14.012098Z","shell.execute_reply":"2025-09-05T15:14:14.465763Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def visualize_submission_on_test(submission_csv, test_images_path, image_id, conf_thresh=0.5):\n    \"\"\"\n    Visualize predicted masks from submission CSV on a test image.\n    Args:\n        submission_csv: path to submission CSV\n        test_images_path: path to test images folder\n        image_id: image filename without extension, e.g., 'test_000'\n        conf_thresh: minimum confidence to show prediction\n    \"\"\"\n    import matplotlib.pyplot as plt\n    import pandas as pd\n    import cv2\n    import numpy as np\n    import random\n\n    # Load submission\n    df = pd.read_csv(submission_csv)\n    df = df[df['patient_id'] == image_id]\n    df = df[df['confidence'] >= conf_thresh]\n\n    # Find the image (try multiple extensions)\n    img_path = None\n    for ext in ['.jpg', '.jpeg', '.png', '.JPG', '.JPEG']:\n        candidate = os.path.join(test_images_path, image_id + ext)\n        if os.path.exists(candidate):\n            img_path = candidate\n            break\n    if img_path is None:\n        raise FileNotFoundError(f\"No image found for {image_id} in {test_images_path}\")\n\n    # Load image\n    img = cv2.imread(img_path)\n    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n\n    # Draw polygons\n    h, w, _ = img.shape  # image height and width\n    for _, row in df.iterrows():\n        class_id = int(row['class_id'])\n        poly_str = row['poly']  # \"x1 y1 x2 y2 ...\"\n        pts = np.array(list(map(float, poly_str.split()))).reshape(-1, 2)\n\n    # Convert normalized coordinates (0-1) to pixel coordinates\n    pts[:, 0] = pts[:, 0] * w\n    pts[:, 1] = pts[:, 1] * h\n    pts = pts.astype(np.int32)\n\n    # Random color per class\n    color = tuple([int(c) for c in np.random.randint(0, 255, size=3)])\n    cv2.polylines(img, [pts], isClosed=True, color=color, thickness=2)\n    cv2.fillPoly(img, [pts], color=color + (50,))\n\n\n    # Show image\n    plt.figure(figsize=(8, 8))\n    plt.imshow(img)\n    plt.axis('off')\n    plt.title(f\"Predictions for {image_id}\")\n    plt.show()\n\n\nvisualize_submission_on_test(\"submission_unet++.csv\", TEST_IMAGES_PATH, \"test_114\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-05T15:14:18.347082Z","iopub.execute_input":"2025-09-05T15:14:18.347387Z","iopub.status.idle":"2025-09-05T15:14:21.134846Z","shell.execute_reply.started":"2025-09-05T15:14:18.347346Z","shell.execute_reply":"2025-09-05T15:14:21.134164Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import cv2\nimport numpy as np\nimport pandas as pd\nfrom skimage import measure\n\ndef masks_to_submission(model, test_loader, submission_path):\n    \"\"\"\n    Convert U-Net++ predictions to submission CSV\n    \"\"\"\n    submission_data = []\n\n    model.eval()\n    with torch.no_grad():\n        for images, image_names in test_loader:\n            images = images.to(Config.DEVICE)\n            outputs = model(images)  # shape: (B, C, H, W)\n            probs = torch.softmax(outputs, dim=1)  # for confidence\n\n            preds = torch.argmax(probs, dim=1).cpu().numpy()  # (B, H, W)\n            probs = probs.cpu().numpy()  # (B, C, H, W)\n\n            for b in range(len(image_names)):\n                image_id = image_names[b].split('.')[0]\n                pred_mask = preds[b]  # (H, W)\n                prob_mask = probs[b]   # (C, H, W)\n\n                for class_id in range(1, Config.NUM_CLASSES + 1):  # skip background\n                    class_mask = (pred_mask == class_id).astype(np.uint8)\n                    if class_mask.sum() == 0:\n                        continue  # no instances\n\n                    # Find connected components\n                    labeled_mask = measure.label(class_mask, connectivity=1)\n                    for region_label in np.unique(labeled_mask):\n                        if region_label == 0:\n                            continue\n                        instance_mask = (labeled_mask == region_label).astype(np.uint8)\n\n                        # Get contours\n                        contours, _ = cv2.findContours(\n                            instance_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE\n                        )\n\n                        if len(contours) == 0:\n                            continue\n\n                        # Choose the largest contour\n                        contour = max(contours, key=lambda x: cv2.contourArea(x))\n\n                        # Compute confidence as max probability inside the instance\n                        conf = float(prob_mask[class_id][instance_mask==1].max())\n\n                        # Normalize contour points to [0,1]\n                        h, w = class_mask.shape\n                        poly = []\n                        for point in contour.squeeze():\n                            x_norm = point[0] / w\n                            y_norm = point[1] / h\n                            poly.append(x_norm)\n                            poly.append(y_norm)\n\n                        poly_str = \" \".join([f\"{p:.6f}\" for p in poly])\n                        submission_data.append([image_id, class_id-1, conf, poly_str])\n\n    # Save CSV\n    submission_df = pd.DataFrame(submission_data, columns=['patient_id','class_id','confidence','poly'])\n    submission_df.to_csv(submission_path, index=False)\n    print(f\"Submission saved to {submission_path}\")\n\n# After loading your best U-Net++ model\nmasks_to_submission(model, test_loader, \"submission_unetpp.csv\")\n","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}