{"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":"nvidiaTeslaT4","dataSources":[{"sourceId":1063445,"sourceType":"datasetVersion","datasetId":589983},{"sourceId":2822650,"sourceType":"datasetVersion","datasetId":1715304},{"sourceId":6318833,"sourceType":"datasetVersion","datasetId":3636171},{"sourceId":13048986,"sourceType":"datasetVersion","datasetId":8263116},{"sourceId":13042854,"sourceType":"datasetVersion","datasetId":8258982}],"dockerImageVersionId":31090,"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":"# This cell must be the VERY FIRST cell in your notebook\n!pip install segmentation-models-pytorch albumentations timm --quiet","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-12T20:54:46.790079Z","iopub.execute_input":"2025-09-12T20:54:46.790794Z","iopub.status.idle":"2025-09-12T20:54:50.239722Z","shell.execute_reply.started":"2025-09-12T20:54:46.790765Z","shell.execute_reply":"2025-09-12T20:54:50.238862Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport glob\nimport numpy as np\nimport matplotlib.pyplot as plt\nfrom tqdm import tqdm\nimport cv2 # Albumentations can use this backend\n\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.utils.data import Dataset, DataLoader, random_split, ConcatDataset, Subset\nfrom PIL import Image\nimport segmentation_models_pytorch as smp\n\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\n\n# ===============================\n# 1. Data Preparation at High Resolution\n# ===============================\n\ntrain_transform = A.Compose([\n    A.Resize(512, 512),\n    A.HorizontalFlip(p=0.5),\n    A.VerticalFlip(p=0.5),\n    A.ElasticTransform(p=0.5, alpha=120, sigma=120 * 0.05),\n    A.GridDistortion(p=0.5),\n    A.RandomBrightnessContrast(p=0.8),\n    A.CLAHE(p=0.8),\n    A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),\n    ToTensorV2(),\n])\n\nval_transform = A.Compose([\n    A.Resize(512, 512),\n    A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),\n    ToTensorV2(),\n])\n\nclass RetinaNumpyDataset(Dataset):\n    def __init__(self, image_dir, mask_dir):\n        self.image_paths = sorted(glob.glob(os.path.join(image_dir, \"*\")))\n        self.mask_paths = sorted(glob.glob(os.path.join(mask_dir, \"*\")))\n\n    def __len__(self):\n        return len(self.image_paths)\n\n    def __getitem__(self, idx):\n        image = np.array(Image.open(self.image_paths[idx]).convert(\"RGB\"))\n        mask = np.array(Image.open(self.mask_paths[idx]).convert(\"L\"))\n        return image, mask\n\nclass TransformedRetinaDataset(Dataset):\n    def __init__(self, original_dataset, transform):\n        self.original_dataset = original_dataset\n        self.transform = transform\n\n    def __len__(self):\n        return len(self.original_dataset)\n\n    def __getitem__(self, idx):\n        image_np, mask_np = self.original_dataset[idx]\n        transformed = self.transform(image=image_np, mask=mask_np)\n        image = transformed['image']\n        mask = transformed['mask']\n        mask = (mask > 0).float().unsqueeze(0)\n        return image, mask\n\ndatasets_np = [\n    RetinaNumpyDataset(\"/kaggle/input/drive-digital-retinal-images-for-vessel-extraction/DRIVE/training/images\", \"/kaggle/input/drive-digital-retinal-images-for-vessel-extraction/DRIVE/training/1st_manual\"),\n    RetinaNumpyDataset(\"/kaggle/input/retina-blood-vessel/Data/train/image\", \"/kaggle/input/retina-blood-vessel/Data/train/mask\"),\n]\n\n# --- THIS IS THE CORRECTED PART ---\nfull_np_dataset = ConcatDataset(datasets_np)\n\ntorch.manual_seed(42)\ntrain_size = int(0.8 * len(full_np_dataset))\nval_size = len(full_np_dataset) - train_size\ntrain_indices, val_indices = random_split(range(len(full_np_dataset)), [train_size, val_size])\n\ntrain_dataset = TransformedRetinaDataset(Subset(full_np_dataset, train_indices.indices), transform=train_transform)\nval_dataset = TransformedRetinaDataset(Subset(full_np_dataset, val_indices.indices), transform=val_transform)\n\nBATCH_SIZE = 1 \ntrain_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=0, pin_memory=True)\nval_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=0, pin_memory=True)\n\n# ===============================\n# 2. Model, Loss, and Optimizer\n# ===============================\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\nmodel = smp.UnetPlusPlus(\n    encoder_name=\"efficientnet-b5\",\n    encoder_weights=\"imagenet\",\n    in_channels=3,\n    classes=1,\n).to(device)\n\nclass SuperLoss(nn.Module):\n    def __init__(self, focal_weight=0.5, tversky_weight=0.5):\n        super().__init__()\n        self.focal = smp.losses.FocalLoss(mode='binary')\n        self.tversky = smp.losses.TverskyLoss(mode='binary', alpha=0.3, beta=0.7)\n\n    def forward(self, pred, true):\n        return 0.5 * self.focal(pred, true) + 0.5 * self.tversky(pred, true)\n\nloss_fn = SuperLoss()\noptimizer = optim.AdamW(model.parameters(), lr=1e-4, weight_decay=1e-5)\n\nEPOCHS = 50 \n\nscheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=EPOCHS)\n\n# ===============================\n# 3. Metrics\n# ===============================\n# Helper function to get TP, FP, FN, TN\ndef get_stats(y_pred, y_true):\n    y_pred_sigmoid = torch.sigmoid(y_pred)\n    # Using a threshold of 0.5 for metrics calculation\n    y_pred_binary = (y_pred_sigmoid > 0.5).long()\n    tp, fp, fn, tn = smp.metrics.get_stats(y_pred_binary, y_true.long(), mode='binary')\n    return tp, fp, fn, tn\n\ndef dice_coef(y_pred, y_true):\n    return smp.metrics.f1_score(*get_stats(y_pred, y_true), reduction='micro')\n\ndef iou_coef(y_pred, y_true):\n    return smp.metrics.iou_score(*get_stats(y_pred, y_true), reduction='micro')\n\ndef pixel_accuracy(y_pred, y_true):\n    return smp.metrics.accuracy(*get_stats(y_pred, y_true), reduction='micro')\n\ndef precision(y_pred, y_true):\n    return smp.metrics.precision(*get_stats(y_pred, y_true), reduction='micro')\n\ndef recall(y_pred, y_true):\n    return smp.metrics.recall(*get_stats(y_pred, y_true), reduction='micro')\n\n# ===============================\n# 4. Training Loop\n# ===============================\nbest_val_iou = 0.0 \npatience = 5 \nepochs_no_improve = 0\n\nfor epoch in range(EPOCHS):\n    model.train()\n    train_loss, train_dice, train_iou, train_acc, train_prec, train_rec = 0, 0, 0, 0, 0, 0\n    \n    for images, masks in tqdm(train_loader, desc=f\"Epoch {epoch+1}/{EPOCHS} [Train]\"):\n        images, masks = images.to(device), masks.to(device)\n        outputs = model(images)\n        loss = loss_fn(outputs, masks)\n        \n        optimizer.zero_grad()\n        loss.backward()\n        optimizer.step()\n        \n        train_loss += loss.item()\n        train_dice += dice_coef(outputs, masks).item()\n        train_iou += iou_coef(outputs, masks).item()\n        train_acc += pixel_accuracy(outputs, masks).item()\n        train_prec += precision(outputs, masks).item()\n        train_rec += recall(outputs, masks).item()\n\n    model.eval()\n    val_loss, val_dice, val_iou, val_acc, val_prec, val_rec = 0, 0, 0, 0, 0, 0\n\n    with torch.no_grad():\n        for images, masks in tqdm(val_loader, desc=f\"Epoch {epoch+1}/{EPOCHS} [Val]\"):\n            images, masks = images.to(device), masks.to(device)\n            outputs = model(images)\n            loss = loss_fn(outputs, masks)\n\n            val_loss += loss.item()\n            val_dice += dice_coef(outputs, masks).item()\n            val_iou += iou_coef(outputs, masks).item()\n            val_acc += pixel_accuracy(outputs, masks).item()\n            val_prec += precision(outputs, masks).item()\n            val_rec += recall(outputs, masks).item()\n\n    scheduler.step()\n    \n    avg_train_loss = train_loss/len(train_loader)\n    avg_train_dice = train_dice/len(train_loader)\n    avg_train_iou = train_iou/len(train_loader)\n    avg_train_acc = train_acc/len(train_loader)\n    avg_train_prec = train_prec/len(train_loader)\n    avg_train_rec = train_rec/len(train_loader)\n    \n    avg_val_loss = val_loss/len(val_loader)\n    avg_val_dice = val_dice/len(val_loader)\n    avg_val_iou = val_iou/len(val_loader)\n    avg_val_acc = val_acc/len(val_loader)\n    avg_val_prec = val_prec/len(val_loader)\n    avg_val_rec = val_rec/len(val_loader)\n\n    print(f\"--- Epoch {epoch+1}/{EPOCHS} ---\")\n    print(f\"Train -> Loss: {avg_train_loss:.4f} | Dice: {avg_train_dice:.4f} | IoU: {avg_train_iou:.4f} | Acc: {avg_train_acc:.4f} | Prec: {avg_train_prec:.4f} | Rec: {avg_train_rec:.4f}\")\n    print(f\"Valid -> Loss: {avg_val_loss:.4f} | Dice: {avg_val_dice:.4f} | IoU: {avg_val_iou:.4f} | Acc: {avg_val_acc:.4f} | Prec: {avg_val_prec:.4f} | Rec: {avg_val_rec:.4f}\")\n\n    if avg_val_iou > best_val_iou:\n        best_val_iou = avg_val_iou\n        torch.save(model.state_dict(), \"/kaggle/working/unetplusplus_retina_final.pth\")\n        print(f\"✅ Model Saved! New best Validation IoU: {best_val_iou:.4f}\")\n        epochs_no_improve = 0\n    else:\n        epochs_no_improve += 1\n        print(f\"Validation IoU did not improve for {epochs_no_improve} epoch(s).\")\n\n    if epochs_no_improve >= patience:\n        print(f\"Early stopping triggered after {patience} epochs with no improvement.\")\n        break\n\n# ===============================\n# 5. Prediction Visualization\n# ===============================\n# We will only run this if the training finishes successfully and saves a model\nif os.path.exists(\"/kaggle/working/unetplusplus_retina_final.pth\"):\n    print(\"\\nLoading best model for prediction...\")\n    model.load_state_dict(torch.load(\"/kaggle/working/unetplusplus_retina_final.pth\"))\n    print(\"Best model loaded successfully.\")\n\n    def denormalize(tensor, mean, std):\n        mean = torch.tensor(mean).view(3, 1, 1)\n        std = torch.tensor(std).view(3, 1, 1)\n        return tensor.cpu() * std + mean\n\n    images, masks = next(iter(val_loader))\n    images, masks = images.to(device), masks.to(device)\n\n    with torch.no_grad():\n        outputs = model(images)\n        preds = (torch.sigmoid(outputs) > 0.5).float()\n\n    i = 0\n    plt.figure(figsize=(12, 4))\n    plt.subplot(1, 3, 1)\n    plt.title(\"Original\")\n    img_denorm = denormalize(images[i], mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225))\n    plt.imshow(img_denorm.permute(1, 2, 0).clamp(0, 1)) # Clamp values to [0,1] for proper display\n    plt.subplot(1, 3, 2)\n    plt.title(\"True Mask\")\n    plt.imshow(masks[i][0].cpu(), cmap=\"gray\")\n    plt.subplot(1, 3, 3)\n    plt.title(\"Predicted Mask\")\n    plt.imshow(preds[i][0].cpu(), cmap=\"gray\")\n    plt.show()\nelse:\n    print(\"\\nNo model was saved, skipping prediction visualization.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-12T19:58:47.323190Z","iopub.execute_input":"2025-09-12T19:58:47.323997Z","iopub.status.idle":"2025-09-12T20:14:56.411820Z","shell.execute_reply.started":"2025-09-12T19:58:47.323959Z","shell.execute_reply":"2025-09-12T20:14:56.411060Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install segmentation-models-pytorch timm","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-13T19:39:28.847620Z","iopub.execute_input":"2025-09-13T19:39:28.848341Z","iopub.status.idle":"2025-09-13T19:40:51.353643Z","shell.execute_reply.started":"2025-09-13T19:39:28.848313Z","shell.execute_reply":"2025-09-13T19:40:51.352738Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport glob\nimport numpy as np\nimport matplotlib.pyplot as plt\nfrom tqdm import tqdm\nfrom PIL import Image\nimport torch\nimport segmentation_models_pytorch as smp\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\nfrom torch.utils.data import Dataset, DataLoader\nimport random\n\n# ==========================================================\n# 1. Configuration - الإعدادات الرئيسية\n# ==========================================================\nMODEL_PATH = \"/kaggle/input/unetplusplus-retina-final/unetplusplus_retina_final.pth\"\n\nBASE_INPUT_DIR = \"/kaggle/input/aptos2019/\"\nBASE_OUTPUT_DIR = \"/kaggle/working/\"\n\nTRAIN_INPUT_FOLDER = os.path.join(BASE_INPUT_DIR, \"train_images/train_images\")\nVAL_INPUT_FOLDER = os.path.join(BASE_INPUT_DIR, \"val_images/val_images\")\nTEST_INPUT_FOLDER = os.path.join(BASE_INPUT_DIR, \"test_images/test_images\")\n\nTRAIN_OUTPUT_FOLDER = os.path.join(BASE_OUTPUT_DIR, \"train_images_predicted\")\nVAL_OUTPUT_FOLDER = os.path.join(BASE_OUTPUT_DIR, \"val_images_predicted\")\nTEST_OUTPUT_FOLDER = os.path.join(BASE_OUTPUT_DIR, \"test_images_predicted\")\n\nDEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nBATCH_SIZE = 8\n\n# ==========================================================\n# 2. Definitions - التعريفات\n# ==========================================================\npredict_transform = A.Compose([\n    A.Resize(512, 512),\n    A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),\n    ToTensorV2(),\n])\n\nclass PredictionDataset(Dataset):\n    def __init__(self, image_paths, transform):\n        self.image_paths = image_paths\n        self.transform = transform\n\n    def __len__(self):\n        return len(self.image_paths)\n\n    def __getitem__(self, idx):\n        image_path = self.image_paths[idx]\n        try:\n            image = np.array(Image.open(image_path).convert(\"RGB\"))\n            transformed = self.transform(image=image)\n            image_tensor = transformed['image']\n            return image_tensor, image_path\n        except Exception as e:\n            print(f\"Could not read image {image_path}, skipping. Error: {e}\")\n            return None, None\n\n# ==========================================================\n# 3. Model Loading - تحميل الموديل\n# ==========================================================\nif not os.path.exists(MODEL_PATH):\n    print(f\"ERROR: Model file not found at {MODEL_PATH}\")\n    model = None\nelse:\n    print(\"Loading the champion model...\")\n    model = smp.UnetPlusPlus(\n        encoder_name=\"efficientnet-b5\",\n        encoder_weights=None,\n        in_channels=3,\n        classes=1,\n    ).to(DEVICE)\n    model.load_state_dict(torch.load(MODEL_PATH))\n    model.eval()\n    print(\"Model loaded successfully!\")\n\n# ==========================================================\n# 4. Prediction Function - دالة التوقع\n# ==========================================================\ndef predict_on_folder(input_folder, output_folder, model, transform):\n    print(f\"\\n--- Starting prediction for folder: {input_folder} ---\")\n    os.makedirs(output_folder, exist_ok=True)\n    \n    print(\"Step 1 of 2: Searching for image files...\")\n    image_paths = []\n    for ext in [\"*.png\", \"*.jpeg\", \"*.jpg\", \"*.tif\"]:\n        image_paths.extend(glob.glob(os.path.join(input_folder, ext)))\n    \n    print(f\"Found {len(image_paths)} images.\")\n    \n    if not image_paths:\n        print(\"No images found in this folder, skipping.\")\n        return\n\n    dataset = PredictionDataset(image_paths, transform=transform)\n    loader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=0)\n    \n    print(\"Step 2 of 2: Starting the prediction loop...\")\n    with torch.no_grad():\n        for batch in tqdm(loader, desc=f\"Predicting\"):\n            if batch is None or batch[0] is None:\n                continue\n            image_tensors, image_paths_batch = batch\n            \n            image_tensors = image_tensors.to(DEVICE)\n            outputs = model(image_tensors)\n            preds = (torch.sigmoid(outputs) > 0.5).float()\n\n            for i, pred_tensor in enumerate(preds):\n                pred_mask_np = pred_tensor.squeeze().cpu().numpy()\n                pred_mask_img = Image.fromarray((pred_mask_np * 255).astype(np.uint8))\n                \n                original_filename = os.path.basename(image_paths_batch[i])\n                new_filename = os.path.splitext(original_filename)[0] + \".png\"\n                save_path = os.path.join(output_folder, new_filename)\n                pred_mask_img.save(save_path)\n    \n    print(f\"--- Predictions for {input_folder} are complete and saved to {output_folder} ---\")\n\n# ==========================================================\n# 5. Main Execution - التشغيل الرئيسي\n# ==========================================================\nif model is not None:\n    predict_on_folder(TRAIN_INPUT_FOLDER, TRAIN_OUTPUT_FOLDER, model, predict_transform)\n    predict_on_folder(VAL_INPUT_FOLDER, VAL_OUTPUT_FOLDER, model, predict_transform)\n    predict_on_folder(TEST_INPUT_FOLDER, TEST_OUTPUT_FOLDER, model, predict_transform)\n    print(\"\\nAll predictions are complete!\")\nelse:\n    print(\"\\nSkipping prediction because the model file was not loaded.\")\n\n# ==========================================================\n# 6. Visualization - عرض النتائج\n# ==========================================================\n# بشكل افتراضي، سنعرض نتائج مجلد الاختبار\nORIGINAL_DIR_TO_SHOW = TEST_INPUT_FOLDER\nPREDICTED_DIR_TO_SHOW = TEST_OUTPUT_FOLDER\nNUM_EXAMPLES = 5\n\nprint(f\"\\nDisplaying {NUM_EXAMPLES} random examples from the test set predictions...\")\npredicted_mask_paths = glob.glob(os.path.join(PREDICTED_DIR_TO_SHOW, \"*.png\"))\n\nif not predicted_mask_paths:\n    print(\"No predicted masks found to display!\")\nelse:\n    random_samples = random.sample(predicted_mask_paths, min(NUM_EXAMPLES, len(predicted_mask_paths)))\n    \n    plt.figure(figsize=(10, 5 * NUM_EXAMPLES))\n    \n    for i, mask_path in enumerate(random_samples):\n        predicted_mask = Image.open(mask_path)\n        base_filename = os.path.splitext(os.path.basename(mask_path))[0]\n        \n        original_path = None\n        for ext in ['.png', '.jpeg', '.jpg', '.tif']:\n            path_to_check = os.path.join(ORIGINAL_DIR_TO_SHOW, base_filename + ext)\n            if os.path.exists(path_to_check):\n                original_path = path_to_check\n                break\n        \n        if original_path:\n            original_image = Image.open(original_path)\n            \n            plt.subplot(NUM_EXAMPLES, 2, 2*i + 1)\n            plt.title(f\"Original: {base_filename}\")\n            plt.imshow(original_image)\n            plt.axis('off')\n            \n            plt.subplot(NUM_EXAMPLES, 2, 2*i + 2)\n            plt.title(\"Predicted Mask\")\n            plt.imshow(predicted_mask, cmap='gray')\n            plt.axis('off')\n        else:\n            print(f\"Could not find original image for mask: {mask_path}\")\n\n    plt.tight_layout()\n    plt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-13T19:40:56.855484Z","iopub.execute_input":"2025-09-13T19:40:56.856179Z","iopub.status.idle":"2025-09-13T19:52:55.219245Z","shell.execute_reply.started":"2025-09-13T19:40:56.856149Z","shell.execute_reply":"2025-09-13T19:52:55.218464Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport gc\nimport cv2\nimport numpy as np\nimport pandas as pd\nfrom tqdm.auto import tqdm\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, WeightedRandomSampler\n\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\n\nfrom sklearn.metrics import accuracy_score, cohen_kappa_score\nfrom sklearn.utils.class_weight import compute_class_weight\n\nimport timm\n\n# =============================================================================\n# CONFIGURATION FOR 4-CHANNEL (IMAGE + VESSEL MASK) INPUT\n# =============================================================================\nclass CFG:\n    # --- MODEL & IMAGE SIZE (Same as your baseline EffNet-B3) ---\n    MODEL_NAME = 'efficientnet_b3'\n    IMG_SIZE = 384\n    BATCH_SIZE = 8\n\n    # --- DATA PATHS ---\n    BASE_PATH = \"/kaggle/input/aptos2019\"\n    TRAIN_CSV = os.path.join(BASE_PATH, \"train_1.csv\")\n    VAL_CSV   = os.path.join(BASE_PATH, \"valid.csv\")\n    TRAIN_DIR = os.path.join(BASE_PATH, \"train_images\", \"train_images\")\n    VAL_DIR   = os.path.join(BASE_PATH, \"val_images\", \"val_images\")\n    \n    # --- PATHS TO YOUR NEW SEGMENTED MASKS ---\n    SEG_BASE_PATH = \"/kaggle/working\"\n    SEG_TRAIN_DIR = os.path.join(SEG_BASE_PATH, \"train_images_predicted\")\n    SEG_VAL_DIR   = os.path.join(SEG_BASE_PATH, \"val_images_predicted\")\n\n    # --- TRAINING PIPELINE (Identical to your successful run for fair comparison) ---\n    S1_EPOCHS = 15; S1_LR = 1e-4; S1_USE_MIXUP = True\n    S2_EPOCHS = 15; S2_LR = 3e-5; S2_USE_MIXUP = False\n    \n    # --- GENERAL & SAVING ---\n    DEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n    NUM_WORKERS = 2\n    PATIENCE = 5\n    SEED = 42\n    LABEL_SMOOTHING = 0.05\n    # New save paths for this experiment\n    SAVE_PATH_S1 = \"best_model_effnet_b3_seg_stage1.pth\"\n    SAVE_PATH_FINAL = \"best_model_effnet_b3_seg_final.pth\"\n\ndef seed_everything(seed):\n    os.environ['PYTHONHASHSEED'] = str(seed); np.random.seed(seed); torch.manual_seed(seed); torch.cuda.manual_seed(seed)\n    torch.backends.cudnn.deterministic = True; torch.backends.cudnn.benchmark = True\nseed_everything(CFG.SEED)\n\n# =============================================================================\n# PREPROCESSING & AUGMENTATIONS (CORRECTED LOGIC)\n# =============================================================================\ndef preprocess_ben_graham(image, output_size):\n    # This function only preprocesses the 3-channel image\n    try:\n        gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)\n        if gray.mean() < 15: \n            image = cv2.resize(image, (output_size, output_size), interpolation=cv2.INTER_AREA)\n        else:\n            _, thresh = cv2.threshold(gray, 15, 255, cv2.THRESH_BINARY)\n            contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)\n            if contours:\n                largest_contour = max(contours, key=cv2.contourArea)\n                x, y, w, h = cv2.boundingRect(largest_contour)\n                image = image[y:y+h, x:x+w]\n            image = cv2.resize(image, (output_size, output_size), interpolation=cv2.INTER_AREA)\n    except Exception: \n        image = cv2.resize(image, (output_size, output_size), interpolation=cv2.INTER_AREA)\n    \n    b, g, r = cv2.split(image)\n    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))\n    g = clahe.apply(g)\n    \n    return cv2.merge((b, g, r))\n\ndef get_transforms(is_train=True):\n    # This pipeline now only contains augmentations. Preprocessing happens before.\n    if is_train:\n        return A.Compose([\n            A.HorizontalFlip(p=0.5),\n            A.VerticalFlip(p=0.5),\n            A.ShiftScaleRotate(shift_limit=0.05, scale_limit=0.1, rotate_limit=15, p=0.7),\n            A.RandomBrightnessContrast(brightness_limit=0.2, contrast_limit=0.2, p=0.7),\n            A.HueSaturationValue(hue_shift_limit=10, sat_shift_limit=20, val_shift_limit=10, p=0.5),\n        ])\n    else:\n        # No augmentations for validation/test\n        return None\n\n# =============================================================================\n# UPGRADED DATASET (CORRECTED LOGIC)\n# =============================================================================\nclass Dataset4Channel(Dataset):\n    def __init__(self, df, img_dir, seg_dir, transform=None):\n        self.df = df.reset_index(drop=True)\n        self.img_dir = img_dir\n        self.seg_dir = seg_dir\n        self.transform = transform\n        # The final normalization/tensor conversion is always applied\n        self.post_transform = A.Compose([\n            A.Normalize(mean=[0.485, 0.456, 0.406, 0.5], std=[0.229, 0.224, 0.225, 0.5]),\n            ToTensorV2()\n        ])\n\n    def __len__(self): \n        return len(self.df)\n\n    def __getitem__(self, idx):\n        row = self.df.iloc[idx]\n        img_path = os.path.join(self.img_dir, row['id_code'] + '.png')\n        seg_path = os.path.join(self.seg_dir, row['id_code'] + '.png')\n        \n        img = cv2.imread(img_path)\n        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n        \n        mask = cv2.imread(seg_path, cv2.IMREAD_GRAYSCALE)\n        \n        # Step 1: Apply preprocessing to the 3-channel RGB image first\n        img = preprocess_ben_graham(img, CFG.IMG_SIZE)\n        \n        # Step 2: Resize the mask to the exact same size to ensure alignment\n        mask = cv2.resize(mask, (CFG.IMG_SIZE, CFG.IMG_SIZE), interpolation=cv2.INTER_NEAREST)\n        \n        # Step 3: Apply geometric and color augmentations to the ALIGNED pair\n        if self.transform:\n            augmented = self.transform(image=img, mask=mask)\n            img = augmented['image']\n            mask = augmented['mask']\n            \n        # Step 4: Add the mask as the 4th channel\n        img_4_channel = np.dstack((img, mask))\n        \n        # Step 5: Apply final normalization and convert to tensor\n        img_4_channel = self.post_transform(image=img_4_channel)['image']\n            \n        label = torch.tensor(row['diagnosis'], dtype=torch.long)\n        return img_4_channel, label\n\n# =============================================================================\n# UPGRADED MODEL TO ACCEPT 4 CHANNELS (Unchanged, was already correct)\n# =============================================================================\nclass EfficientNet4ChannelOrdinal(nn.Module):\n    def __init__(self, model_name, num_classes=5, pretrained=True):\n        super().__init__()\n        self.backbone = timm.create_model(model_name, pretrained=pretrained, num_classes=0, global_pool='avg')\n        \n        original_conv = self.backbone.conv_stem\n        original_weights = original_conv.weight.clone()\n\n        new_conv = nn.Conv2d(4, original_conv.out_channels, \n                             kernel_size=original_conv.kernel_size, stride=original_conv.stride, \n                             padding=original_conv.padding, bias=(original_conv.bias is not None))\n        with torch.no_grad():\n            new_conv.weight[:, :3] = original_weights\n            new_conv.weight[:, 3] = original_weights.mean(dim=1)\n        self.backbone.conv_stem = new_conv\n        \n        feature_dim = self.backbone.num_features\n        self.classifier = nn.Sequential(nn.Dropout(0.5), nn.Linear(feature_dim, 256), nn.ReLU(), nn.Dropout(0.3), nn.Linear(256, num_classes - 1))\n\n    def forward(self, x): \n        return self.classifier(self.backbone(x))\n\n# --- Loss functions, training loops, and other utilities are unchanged ---\nclass WeightedOrdinalFocalLoss(nn.Module):\n    def __init__(self, num_classes=5, gamma=2.0, class_weights=None, label_smoothing=0.0):\n        super().__init__()\n        self.num_classes, self.gamma, self.class_weights, self.label_smoothing = num_classes, gamma, class_weights, label_smoothing\n        self.bce = nn.BCEWithLogitsLoss(reduction='none')\n    def forward(self, outputs, targets):\n        ordinal_targets = torch.zeros_like(outputs)\n        for i, t in enumerate(targets):\n            if t > 0: ordinal_targets[i, :t] = 1.0\n        if self.label_smoothing > 0.0: ordinal_targets = ordinal_targets * (1.0 - self.label_smoothing) + 0.5 * self.label_smoothing\n        bce = self.bce(outputs, ordinal_targets)\n        if self.class_weights is not None:\n            weights = self.class_weights[targets].view(-1, 1).expand(-1, outputs.shape[1])\n            bce = bce * weights\n        pt = torch.exp(-bce)\n        focal = (1 - pt) ** self.gamma * bce\n        return focal.mean()\n\nclass SmoothKappaLoss(nn.Module):\n    def __init__(self, num_classes=5, eps=1e-7):\n        super().__init__()\n        self.num_classes, self.eps = num_classes, eps\n        W = torch.zeros(num_classes, num_classes)\n        for i in range(num_classes):\n            for j in range(num_classes): W[i,j] = ((i - j)**2) / ((num_classes - 1)**2)\n        self.register_buffer(\"W\", W)\n    def forward(self, outputs, targets):\n        device = outputs.device; B = outputs.size(0); probs = torch.sigmoid(outputs)\n        class_probs = torch.zeros(B, self.num_classes, device=device)\n        class_probs[:, 0] = 1 - probs[:, 0]\n        for k in range(1, self.num_classes-1): class_probs[:, k] = probs[:, k-1] - probs[:, k]\n        class_probs[:, -1] = probs[:, -1]\n        class_probs = torch.clamp(class_probs, min=self.eps, max=1.0)\n        one_hot = F.one_hot(targets, num_classes=self.num_classes).float().to(device)\n        conf_mat = torch.matmul(one_hot.T, class_probs)\n        hist_true = one_hot.sum(dim=0); hist_pred = class_probs.sum(dim=0)\n        expected = torch.outer(hist_true, hist_pred)\n        W = self.W.to(device); obs = torch.sum(W * conf_mat); exp = torch.sum(W * expected)\n        kappa = 1.0 - (B * obs) / (exp + self.eps)\n        return 1.0 - kappa\n\ndef mixup_data(x, y, alpha=0.4):\n    if alpha > 0: lam = np.random.beta(alpha, alpha)\n    else: lam = 1\n    batch_size = x.size()[0]; index = torch.randperm(batch_size).to(x.device)\n    mixed_x = lam * x + (1 - lam) * x[index, :]\n    y_a, y_b = y, y[index]\n    return mixed_x, y_a, y_b, lam\n\ndef ordinal_to_class(outputs): \n    return torch.sum(torch.sigmoid(outputs) > 0.5, dim=1).long()\n\ndef calculate_metrics(outputs, targets):\n    preds = ordinal_to_class(outputs).cpu().numpy()\n    targets_np = targets.cpu().numpy()\n    return accuracy_score(targets_np, preds), cohen_kappa_score(targets_np, preds, weights='quadratic')\n\ndef clear_memory(): \n    gc.collect()\n    torch.cuda.empty_cache()\n\ndef train_epoch(model, loader, optimizer, criterion, scaler, device, use_mixup):\n    model.train(); running_loss = 0.0; all_out, all_t = [], []\n    pbar = tqdm(loader, desc=\"Training\", leave=False)\n    for images, targets in pbar:\n        images, targets = images.to(device, non_blocking=True), targets.to(device, non_blocking=True)\n        optimizer.zero_grad(set_to_none=True)\n        if use_mixup: images, targets_a, targets_b, lam = mixup_data(images, targets)\n        with torch.cuda.amp.autocast():\n            outputs = model(images)\n            if use_mixup: loss = lam * criterion(outputs, targets_a) + (1 - lam) * criterion(outputs, targets_b)\n            else: loss = criterion(outputs, targets)\n        scaler.scale(loss).backward(); scaler.step(optimizer); scaler.update()\n        running_loss += loss.item(); all_out.append(outputs.detach()); all_t.append(targets.detach())\n        pbar.set_postfix(loss=loss.item())\n    all_out, all_t = torch.cat(all_out), torch.cat(all_t)\n    return running_loss / len(loader), *calculate_metrics(all_out, all_t)\n\ndef validate_epoch(model, loader, criterion, device):\n    model.eval(); running_loss = 0.0; all_out, all_t = [], []\n    with torch.no_grad():\n        pbar = tqdm(loader, desc=\"Validating\", leave=False)\n        for images, targets in pbar:\n            images, targets = images.to(device, non_blocking=True), targets.to(device, non_blocking=True)\n            with torch.cuda.amp.autocast():\n                outputs = model(images)\n                loss = criterion(outputs, targets)\n            running_loss += loss.item()\n            all_out.append(outputs)\n            all_t.append(targets)\n    all_out, all_t = torch.cat(all_out), torch.cat(all_t)\n    return running_loss / len(loader), *calculate_metrics(all_out, all_t)\n\ndef main():\n    print(f\"Device: {CFG.DEVICE}, Model: {CFG.MODEL_NAME} (4-Channel), Image Size: {CFG.IMG_SIZE}\")\n    train_df = pd.read_csv(CFG.TRAIN_CSV)\n    val_df = pd.read_csv(CFG.VAL_CSV)\n    \n    train_tf = get_transforms(is_train=True)\n    val_tf = get_transforms(is_train=False)\n\n    train_ds = Dataset4Channel(train_df, CFG.TRAIN_DIR, CFG.SEG_TRAIN_DIR, transform=train_tf)\n    val_ds   = Dataset4Channel(val_df, CFG.VAL_DIR, CFG.SEG_VAL_DIR, transform=val_tf)\n\n    class_weights_sampler = compute_class_weight('balanced', classes=np.unique(train_df['diagnosis']), y=train_df['diagnosis'])\n    sample_weights = np.array([class_weights_sampler[int(l)] for l in train_df['diagnosis']])\n    sampler = WeightedRandomSampler(sample_weights, num_samples=len(sample_weights), replacement=True)\n    train_loader = DataLoader(train_ds, batch_size=CFG.BATCH_SIZE, sampler=sampler, num_workers=CFG.NUM_WORKERS, pin_memory=True)\n    val_loader = DataLoader(val_ds, batch_size=CFG.BATCH_SIZE*2, shuffle=False, num_workers=CFG.NUM_WORKERS, pin_memory=True)\n    \n    model = EfficientNet4ChannelOrdinal(CFG.MODEL_NAME).to(CFG.DEVICE)\n    class_weights_loss = torch.tensor(class_weights_sampler, dtype=torch.float).to(CFG.DEVICE)\n    focal_loss = WeightedOrdinalFocalLoss(num_classes=5, gamma=2.0, class_weights=class_weights_loss, label_smoothing=CFG.LABEL_SMOOTHING)\n    kappa_loss = SmoothKappaLoss(num_classes=5)\n    \n    def hybrid_loss(outputs, targets): \n        return 0.7 * kappa_loss(outputs, targets) + 0.3 * focal_loss(outputs, targets)\n    \n    scaler = torch.cuda.amp.GradScaler()\n\n    # --- STAGE 1 ---\n    print(\"\\n\" + \"=\"*50 + \"\\n     STARTING STAGE 1 (4-Channel)\\n\" + \"=\"*50)\n    opt = optim.AdamW(model.parameters(), lr=CFG.S1_LR, weight_decay=1e-4)\n    sched = optim.lr_scheduler.CosineAnnealingLR(opt, T_max=CFG.S1_EPOCHS)\n    best_val_qwk, patience_counter = -1, 0\n\n    for epoch in range(CFG.S1_EPOCHS):\n        clear_memory()\n        print(f\"\\nEpoch {epoch+1}/{CFG.S1_EPOCHS}\")\n        train_loss, train_acc, train_qwk = train_epoch(model, train_loader, opt, focal_loss, scaler, CFG.DEVICE, CFG.S1_USE_MIXUP)\n        val_loss, val_acc, val_qwk = validate_epoch(model, val_loader, focal_loss, CFG.DEVICE)\n        sched.step()\n        print(f\"Train -> Loss:{train_loss:.4f} Acc:{train_acc:.4f} QWK:{train_qwk:.4f}\")\n        print(f\"Valid -> Loss:{val_loss:.4f} Acc:{val_acc:.4f} QWK:{val_qwk:.4f}\")\n        if val_qwk > best_val_qwk:\n            print(f\"Val QWK improved from {best_val_qwk:.4f} to {val_qwk:.4f}. Saving model...\")\n            best_val_qwk, patience_counter = val_qwk, 0\n            torch.save(model.state_dict(), CFG.SAVE_PATH_S1)\n        else:\n            patience_counter += 1\n            if patience_counter >= CFG.PATIENCE: \n                print(\"Early stopping in Stage 1.\")\n                break\n    \n    # --- STAGE 2 ---\n    print(\"\\n\" + \"=\"*50 + \"\\n     STARTING STAGE 2 (4-Channel)\\n\" + \"=\"*50)\n    if os.path.exists(CFG.SAVE_PATH_S1):\n        model.load_state_dict(torch.load(CFG.SAVE_PATH_S1))\n    else:\n        print(\"No Stage 1 model was saved. Continuing with the current model.\")\n\n    opt = optim.AdamW(model.parameters(), lr=CFG.S2_LR, weight_decay=1e-5)\n    sched = optim.lr_scheduler.CosineAnnealingLR(opt, T_max=CFG.S2_EPOCHS)\n    best_val_qwk_stage2, patience_counter = best_val_qwk, 0\n\n    for epoch in range(CFG.S2_EPOCHS):\n        clear_memory()\n        print(f\"\\nEpoch {epoch+1}/{CFG.S2_EPOCHS}\")\n        train_loss, train_acc, train_qwk = train_epoch(model, train_loader, opt, hybrid_loss, scaler, CFG.DEVICE, CFG.S2_USE_MIXUP)\n        val_loss, val_acc, val_qwk = validate_epoch(model, val_loader, hybrid_loss, CFG.DEVICE)\n        sched.step()\n        print(f\"Train -> Loss:{train_loss:.4f} Acc:{train_acc:.4f} QWK:{train_qwk:.4f}\")\n        print(f\"Valid -> Loss:{val_loss:.4f} Acc:{val_acc:.4f} QWK:{val_qwk:.4f}\")\n        if val_qwk > best_val_qwk_stage2:\n            print(f\"Val QWK improved from {best_val_qwk_stage2:.4f} to {val_qwk:.4f}. Saving final model...\")\n            best_val_qwk_stage2, patience_counter = val_qwk, 0\n            torch.save(model.state_dict(), CFG.SAVE_PATH_FINAL)\n        else:\n            patience_counter += 1\n            if patience_counter >= CFG.PATIENCE: \n                print(\"Early stopping in Stage 2.\")\n                break\n\n    print(f\"\\nTraining Finished!\\nFinal Best QWK: {best_val_qwk_stage2:.4f}\")\n\nif __name__ == \"__main__\":\n    main()\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-13T17:01:14.463883Z","iopub.execute_input":"2025-09-13T17:01:14.464465Z","iopub.status.idle":"2025-09-13T18:38:24.851844Z","shell.execute_reply.started":"2025-09-13T17:01:14.464437Z","shell.execute_reply":"2025-09-13T18:38:24.850695Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport gc\nimport cv2\nimport numpy as np\nimport pandas as pd\nfrom tqdm.auto import tqdm\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, WeightedRandomSampler\n\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\n\nfrom sklearn.metrics import accuracy_score, cohen_kappa_score\nfrom sklearn.utils.class_weight import compute_class_weight\n\nimport timm\n\n# =============================================================================\n# CONFIGURATION FOR 4-CHANNEL (IMAGE + VESSEL MASK) INPUT\n# =============================================================================\nclass CFG:\n    # --- ENHANCEMENT: Upgraded model for more power ---\n    MODEL_NAME = 'efficientnet_b4'\n    IMG_SIZE = 384\n    BATCH_SIZE = 8\n\n    # --- DATA PATHS (Unchanged as requested) ---\n    BASE_PATH = \"/kaggle/input/aptos2019\"\n    TRAIN_CSV = os.path.join(BASE_PATH, \"train_1.csv\")\n    VAL_CSV   = os.path.join(BASE_PATH, \"valid.csv\")\n    TRAIN_DIR = os.path.join(BASE_PATH, \"train_images\", \"train_images\")\n    VAL_DIR   = os.path.join(BASE_PATH, \"val_images\", \"val_images\")\n    \n    # --- PATHS TO YOUR NEW SEGMENTED MASKS (Unchanged as requested) ---\n    SEG_BASE_PATH = \"/kaggle/working\"\n    SEG_TRAIN_DIR = os.path.join(SEG_BASE_PATH, \"train_images_predicted\")\n    SEG_VAL_DIR   = os.path.join(SEG_BASE_PATH, \"val_images_predicted\")\n\n    # --- TRAINING PIPELINE ---\n    S1_EPOCHS = 15; S1_LR = 1e-4; S1_USE_MIXUP = True\n    S2_EPOCHS = 15; S2_LR = 3e-5; S2_USE_MIXUP = False\n    \n    # --- GENERAL & SAVING ---\n    DEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n    NUM_WORKERS = 2\n    PATIENCE = 5\n    SEED = 42\n    LABEL_SMOOTHING = 0.05\n    # Updated save paths for the new model\n    SAVE_PATH_S1 = \"best_model_effnet_b4_seg_stage1.pth\"\n    SAVE_PATH_FINAL = \"best_model_effnet_b4_seg_final.pth\"\n\ndef seed_everything(seed):\n    os.environ['PYTHONHASHSEED'] = str(seed); np.random.seed(seed); torch.manual_seed(seed); torch.cuda.manual_seed(seed)\n    torch.backends.cudnn.deterministic = True; torch.backends.cudnn.benchmark = True\nseed_everything(CFG.SEED)\n\n# =============================================================================\n# PREPROCESSING & AUGMENTATIONS\n# =============================================================================\ndef preprocess_ben_graham(image, output_size):\n    try:\n        gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)\n        if gray.mean() < 15: image = cv2.resize(image, (output_size, output_size), interpolation=cv2.INTER_AREA)\n        else:\n            _, thresh = cv2.threshold(gray, 15, 255, cv2.THRESH_BINARY)\n            contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)\n            if contours:\n                largest_contour = max(contours, key=cv2.contourArea)\n                x, y, w, h = cv2.boundingRect(largest_contour)\n                image = image[y:y+h, x:x+w]\n            image = cv2.resize(image, (output_size, output_size), interpolation=cv2.INTER_AREA)\n    except Exception: image = cv2.resize(image, (output_size, output_size), interpolation=cv2.INTER_AREA)\n    b, g, r = cv2.split(image)\n    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))\n    g = clahe.apply(g)\n    return cv2.merge((b, g, r))\n\ndef get_transforms(is_train=True):\n    if is_train:\n        # --- ENHANCEMENT: Added stronger augmentations ---\n        return A.Compose([\n            A.HorizontalFlip(p=0.5),\n            A.VerticalFlip(p=0.5),\n            A.ShiftScaleRotate(shift_limit=0.05, scale_limit=0.1, rotate_limit=15, p=0.7),\n            A.GridDistortion(p=0.3), # Added\n            A.RandomBrightnessContrast(brightness_limit=0.2, contrast_limit=0.2, p=0.7),\n            A.HueSaturationValue(hue_shift_limit=10, sat_shift_limit=20, val_shift_limit=10, p=0.5),\n        ])\n    else:\n        return None\n\n# =============================================================================\n# DATASET - ENHANCED FOR ROBUSTNESS\n# =============================================================================\nclass Dataset4Channel(Dataset):\n    def __init__(self, df, img_dir, seg_dir, transform=None):\n        self.df = df.reset_index(drop=True)\n        self.img_dir = img_dir\n        self.seg_dir = seg_dir\n        self.transform = transform\n        self.post_transform = A.Compose([\n            A.Normalize(mean=[0.485, 0.456, 0.406, 0.5], std=[0.229, 0.224, 0.225, 0.5]),\n            ToTensorV2()\n        ])\n\n    def __len__(self): \n        return len(self.df)\n\n    # --- ENHANCEMENT: This version safely handles missing images or masks ---\n    def __getitem__(self, idx):\n        row = self.df.iloc[idx]\n        # Dynamically find extension for original image\n        base_img_path = os.path.join(self.img_dir, row['id_code'])\n        img_path = f\"{base_img_path}.png\" if os.path.exists(f\"{base_img_path}.png\") else f\"{base_img_path}.jpeg\"\n        \n        seg_path = os.path.join(self.seg_dir, row['id_code'] + '.png')\n        \n        # Safely read original image\n        if os.path.exists(img_path):\n            img = cv2.imread(img_path)\n            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n        else:\n            img = np.zeros((CFG.IMG_SIZE, CFG.IMG_SIZE, 3), dtype=np.uint8)\n            print(f\"Warning: Image not found at {img_path}, using black image.\")\n\n        # Safely read segmentation mask\n        if os.path.exists(seg_path):\n            mask = cv2.imread(seg_path, cv2.IMREAD_GRAYSCALE)\n        else:\n            mask = np.zeros((CFG.IMG_SIZE, CFG.IMG_SIZE), dtype=np.uint8)\n            print(f\"Warning: Mask not found at {seg_path}, using black mask.\")\n        \n        img = preprocess_ben_graham(img, CFG.IMG_SIZE)\n        mask = cv2.resize(mask, (CFG.IMG_SIZE, CFG.IMG_SIZE), interpolation=cv2.INTER_NEAREST)\n        \n        if self.transform:\n            augmented = self.transform(image=img, mask=mask)\n            img = augmented['image']\n            mask = augmented['mask']\n            \n        img_4_channel = np.dstack((img, mask))\n        img_4_channel = self.post_transform(image=img_4_channel)['image']\n        label = torch.tensor(row['diagnosis'], dtype=torch.long)\n        return img_4_channel, label\n\n# =============================================================================\n# MODEL DEFINITION\n# =============================================================================\nclass EfficientNet4ChannelOrdinal(nn.Module):\n    def __init__(self, model_name, num_classes=5, pretrained=True):\n        super().__init__()\n        self.backbone = timm.create_model(model_name, pretrained=pretrained, num_classes=0, global_pool='avg')\n        original_conv = self.backbone.conv_stem\n        original_weights = original_conv.weight.clone()\n        new_conv = nn.Conv2d(4, original_conv.out_channels, \n                             kernel_size=original_conv.kernel_size, stride=original_conv.stride, \n                             padding=original_conv.padding, bias=(original_conv.bias is not None))\n        with torch.no_grad():\n            new_conv.weight[:, :3] = original_weights\n            new_conv.weight[:, 3] = original_weights.mean(dim=1)\n        self.backbone.conv_stem = new_conv\n        feature_dim = self.backbone.num_features\n        self.classifier = nn.Sequential(nn.Dropout(0.5), nn.Linear(feature_dim, 256), nn.ReLU(), nn.Dropout(0.3), nn.Linear(256, num_classes - 1))\n    def forward(self, x): \n        return self.classifier(self.backbone(x))\n\n# =============================================================================\n# UTILITY FUNCTIONS (Losses, Metrics, etc.)\n# =============================================================================\n# (These advanced functions are already well-optimized, no changes needed)\nclass WeightedOrdinalFocalLoss(nn.Module):\n    def __init__(self, num_classes=5, gamma=2.0, class_weights=None, label_smoothing=0.0):\n        super().__init__(); self.num_classes, self.gamma, self.class_weights, self.label_smoothing = num_classes, gamma, class_weights, label_smoothing\n        self.bce = nn.BCEWithLogitsLoss(reduction='none')\n    def forward(self, outputs, targets):\n        ordinal_targets = torch.zeros_like(outputs)\n        for i, t in enumerate(targets):\n            if t > 0: ordinal_targets[i, :t] = 1.0\n        if self.label_smoothing > 0.0: ordinal_targets = ordinal_targets * (1.0 - self.label_smoothing) + 0.5 * self.label_smoothing\n        bce = self.bce(outputs, ordinal_targets)\n        if self.class_weights is not None:\n            weights = self.class_weights[targets].view(-1, 1).expand(-1, outputs.shape[1])\n            bce = bce * weights\n        pt = torch.exp(-bce); focal = (1 - pt) ** self.gamma * bce\n        return focal.mean()\n\nclass SmoothKappaLoss(nn.Module):\n    def __init__(self, num_classes=5, eps=1e-7):\n        super().__init__(); self.num_classes, self.eps = num_classes, eps\n        W = torch.zeros(num_classes, num_classes)\n        for i in range(num_classes):\n            for j in range(num_classes): W[i,j] = ((i - j)**2) / ((num_classes - 1)**2)\n        self.register_buffer(\"W\", W)\n    def forward(self, outputs, targets):\n        device = outputs.device; B = outputs.size(0); probs = torch.sigmoid(outputs)\n        class_probs = torch.zeros(B, self.num_classes, device=device)\n        class_probs[:, 0] = 1 - probs[:, 0]\n        for k in range(1, self.num_classes-1): class_probs[:, k] = probs[:, k-1] - probs[:, k]\n        class_probs[:, -1] = probs[:, -1]\n        class_probs = torch.clamp(class_probs, min=self.eps, max=1.0)\n        one_hot = F.one_hot(targets, num_classes=self.num_classes).float().to(device)\n        conf_mat = torch.matmul(one_hot.T, class_probs); hist_true = one_hot.sum(dim=0); hist_pred = class_probs.sum(dim=0)\n        expected = torch.outer(hist_true, hist_pred)\n        W = self.W.to(device); obs = torch.sum(W * conf_mat); exp = torch.sum(W * expected)\n        kappa = 1.0 - (B * obs) / (exp + self.eps)\n        return 1.0 - kappa\n\ndef mixup_data(x, y, alpha=0.4):\n    if alpha > 0: lam = np.random.beta(alpha, alpha)\n    else: lam = 1\n    batch_size = x.size()[0]; index = torch.randperm(batch_size).to(x.device)\n    mixed_x = lam * x + (1 - lam) * x[index, :]; y_a, y_b = y, y[index]\n    return mixed_x, y_a, y_b, lam\n\ndef ordinal_to_class(outputs): return torch.sum(torch.sigmoid(outputs) > 0.5, dim=1).long()\n\ndef calculate_metrics(outputs, targets):\n    preds, targets_np = ordinal_to_class(outputs).cpu().numpy(), targets.cpu().numpy()\n    return accuracy_score(targets_np, preds), cohen_kappa_score(targets_np, preds, weights='quadratic')\n\ndef clear_memory(): gc.collect(); torch.cuda.empty_cache()\n\n# =============================================================================\n# TRAINING & VALIDATION FUNCTIONS\n# =============================================================================\ndef train_epoch(model, loader, optimizer, criterion, scaler, device, use_mixup):\n    model.train(); running_loss = 0.0; all_out, all_t = [], []\n    pbar = tqdm(loader, desc=\"Training\", leave=False)\n    for images, targets in pbar:\n        images, targets = images.to(device, non_blocking=True), targets.to(device, non_blocking=True)\n        optimizer.zero_grad(set_to_none=True)\n        if use_mixup: images, targets_a, targets_b, lam = mixup_data(images, targets)\n        with torch.cuda.amp.autocast():\n            outputs = model(images)\n            if use_mixup: loss = lam * criterion(outputs, targets_a) + (1 - lam) * criterion(outputs, targets_b)\n            else: loss = criterion(outputs, targets)\n        scaler.scale(loss).backward(); scaler.step(optimizer); scaler.update()\n        running_loss += loss.item(); all_out.append(outputs.detach()); all_t.append(targets.detach())\n        pbar.set_postfix(loss=loss.item())\n    all_out, all_t = torch.cat(all_out), torch.cat(all_t)\n    return running_loss / len(loader), *calculate_metrics(all_out, all_t)\n\ndef validate_epoch(model, loader, criterion, device):\n    model.eval(); running_loss = 0.0; all_out, all_t = [], []\n    with torch.no_grad():\n        pbar = tqdm(loader, desc=\"Validating\", leave=False)\n        for images, targets in pbar:\n            images, targets = images.to(device, non_blocking=True), targets.to(device, non_blocking=True)\n            with torch.cuda.amp.autocast():\n                outputs = model(images); loss = criterion(outputs, targets)\n            running_loss += loss.item(); all_out.append(outputs); all_t.append(targets)\n    all_out, all_t = torch.cat(all_out), torch.cat(all_t)\n    return running_loss / len(loader), *calculate_metrics(all_out, all_t)\n\n# =============================================================================\n# MAIN TRAINING PIPELINE\n# =============================================================================\ndef main():\n    print(f\"Device: {CFG.DEVICE}, Model: {CFG.MODEL_NAME} (4-Channel), Image Size: {CFG.IMG_SIZE}\")\n    train_df, val_df = pd.read_csv(CFG.TRAIN_CSV), pd.read_csv(CFG.VAL_CSV)\n    \n    train_ds = Dataset4Channel(train_df, CFG.TRAIN_DIR, CFG.SEG_TRAIN_DIR, transform=get_transforms(is_train=True))\n    val_ds   = Dataset4Channel(val_df, CFG.VAL_DIR, CFG.SEG_VAL_DIR, transform=get_transforms(is_train=False))\n\n    class_weights = compute_class_weight('balanced', classes=np.unique(train_df['diagnosis']), y=train_df['diagnosis'])\n    sample_weights = np.array([class_weights[int(l)] for l in train_df['diagnosis']])\n    sampler = WeightedRandomSampler(sample_weights, num_samples=len(sample_weights), replacement=True)\n    train_loader = DataLoader(train_ds, batch_size=CFG.BATCH_SIZE, sampler=sampler, num_workers=CFG.NUM_WORKERS, pin_memory=True)\n    val_loader = DataLoader(val_ds, batch_size=CFG.BATCH_SIZE*2, shuffle=False, num_workers=CFG.NUM_WORKERS, pin_memory=True)\n    \n    model = EfficientNet4ChannelOrdinal(CFG.MODEL_NAME).to(CFG.DEVICE)\n    class_weights_loss = torch.tensor(class_weights, dtype=torch.float).to(CFG.DEVICE)\n    focal_loss = WeightedOrdinalFocalLoss(num_classes=5, gamma=2.0, class_weights=class_weights_loss, label_smoothing=CFG.LABEL_SMOOTHING)\n    kappa_loss = SmoothKappaLoss(num_classes=5)\n    def hybrid_loss(outputs, targets): return 0.7 * kappa_loss(outputs, targets) + 0.3 * focal_loss(outputs, targets)\n    \n    scaler = torch.cuda.amp.GradScaler()\n\n    # --- STAGE 1 ---\n    print(\"\\n\" + \"=\"*50 + \"\\n     STARTING STAGE 1 (4-Channel)\\n\" + \"=\"*50)\n    opt = optim.AdamW(model.parameters(), lr=CFG.S1_LR, weight_decay=1e-4)\n    # --- ENHANCEMENT: Using improved scheduler ---\n    sched = optim.lr_scheduler.CosineAnnealingWarmRestarts(opt, T_0=10, T_mult=1, eta_min=1e-6)\n    best_val_qwk, patience_counter = -1, 0\n\n    for epoch in range(CFG.S1_EPOCHS):\n        clear_memory(); print(f\"\\nEpoch {epoch+1}/{CFG.S1_EPOCHS}\")\n        train_loss, train_acc, train_qwk = train_epoch(model, train_loader, opt, focal_loss, scaler, CFG.DEVICE, CFG.S1_USE_MIXUP)\n        val_loss, val_acc, val_qwk = validate_epoch(model, val_loader, focal_loss, CFG.DEVICE)\n        sched.step()\n        print(f\"Train -> Loss:{train_loss:.4f} Acc:{train_acc:.4f} QWK:{train_qwk:.4f}\")\n        print(f\"Valid -> Loss:{val_loss:.4f} Acc:{val_acc:.4f} QWK:{val_qwk:.4f}\")\n        if val_qwk > best_val_qwk:\n            print(f\"Val QWK improved from {best_val_qwk:.4f} to {val_qwk:.4f}. Saving model...\")\n            best_val_qwk, patience_counter = val_qwk, 0\n            torch.save(model.state_dict(), CFG.SAVE_PATH_S1)\n        else:\n            patience_counter += 1\n            if patience_counter >= CFG.PATIENCE: print(\"Early stopping in Stage 1.\"); break\n    \n    # --- STAGE 2 ---\n    print(\"\\n\" + \"=\"*50 + \"\\n     STARTING STAGE 2 (4-Channel)\\n\" + \"=\"*50)\n    if os.path.exists(CFG.SAVE_PATH_S1): model.load_state_dict(torch.load(CFG.SAVE_PATH_S1))\n    else: print(\"No Stage 1 model was saved. Continuing with the current model.\")\n\n    opt = optim.AdamW(model.parameters(), lr=CFG.S2_LR, weight_decay=1e-5)\n    # --- ENHANCEMENT: Using improved scheduler ---\n    sched = optim.lr_scheduler.CosineAnnealingWarmRestarts(opt, T_0=10, T_mult=1, eta_min=3e-6)\n    best_val_qwk_stage2, patience_counter = best_val_qwk, 0\n\n    for epoch in range(CFG.S2_EPOCHS):\n        clear_memory(); print(f\"\\nEpoch {epoch+1}/{CFG.S2_EPOCHS}\")\n        train_loss, train_acc, train_qwk = train_epoch(model, train_loader, opt, hybrid_loss, scaler, CFG.DEVICE, CFG.S2_USE_MIXUP)\n        val_loss, val_acc, val_qwk = validate_epoch(model, val_loader, hybrid_loss, CFG.DEVICE)\n        sched.step()\n        print(f\"Train -> Loss:{train_loss:.4f} Acc:{train_acc:.4f} QWK:{train_qwk:.4f}\")\n        print(f\"Valid -> Loss:{val_loss:.4f} Acc:{val_acc:.4f} QWK:{val_qwk:.4f}\")\n        if val_qwk > best_val_qwk_stage2:\n            print(f\"Val QWK improved from {best_val_qwk_stage2:.4f} to {val_qwk:.4f}. Saving final model...\")\n            best_val_qwk_stage2, patience_counter = val_qwk, 0\n            torch.save(model.state_dict(), CFG.SAVE_PATH_FINAL)\n        else:\n            patience_counter += 1\n            if patience_counter >= CFG.PATIENCE: print(\"Early stopping in Stage 2.\"); break\n\n    print(f\"\\nTraining Finished!\\nFinal Best QWK: {best_val_qwk_stage2:.4f}\")\n\nif __name__ == \"__main__\":\n    main()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-13T19:53:25.001665Z","iopub.execute_input":"2025-09-13T19:53:25.001958Z","iopub.status.idle":"2025-09-13T21:13:23.714786Z","shell.execute_reply.started":"2025-09-13T19:53:25.001937Z","shell.execute_reply":"2025-09-13T21:13:23.713895Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport gc\nimport cv2\nimport numpy as np\nimport pandas as pd\nfrom tqdm.auto import tqdm\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, WeightedRandomSampler\n\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\n\nfrom sklearn.metrics import accuracy_score, cohen_kappa_score\nfrom sklearn.utils.class_weight import compute_class_weight\n\nimport timm\n\n# =============================================================================\n# CONFIGURATION FOR 4-CHANNEL (IMAGE + VESSEL MASK) INPUT\n# =============================================================================\nclass CFG:\n    # --- ENHANCEMENT: Upgraded model for significantly more power ---\n    MODEL_NAME = 'efficientnet_b5'\n    IMG_SIZE = 384\n    BATCH_SIZE = 8 # Note: You may need to lower this to 4 if you get memory errors with B5\n\n    # --- DATA PATHS (Unchanged as requested) ---\n    BASE_PATH = \"/kaggle/input/aptos2019\"\n    TRAIN_CSV = os.path.join(BASE_PATH, \"train_1.csv\")\n    VAL_CSV   = os.path.join(BASE_PATH, \"valid.csv\")\n    TRAIN_DIR = os.path.join(BASE_PATH, \"train_images\", \"train_images\")\n    VAL_DIR   = os.path.join(BASE_PATH, \"val_images\", \"val_images\")\n    \n    # --- PATHS TO YOUR NEW SEGMENTED MASKS (Unchanged as requested) ---\n    SEG_BASE_PATH = \"/kaggle/working\"\n    SEG_TRAIN_DIR = os.path.join(SEG_BASE_PATH, \"train_images_predicted\")\n    SEG_VAL_DIR   = os.path.join(SEG_BASE_PATH, \"val_images_predicted\")\n\n    # --- ENHANCEMENT: Increased training duration for the larger model ---\n    S1_EPOCHS = 20; S1_LR = 1e-4; S1_USE_MIXUP = True\n    S2_EPOCHS = 20; S2_LR = 3e-5; S2_USE_MIXUP = False\n    \n    # --- GENERAL & SAVING ---\n    DEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n    NUM_WORKERS = 2\n    PATIENCE = 5\n    SEED = 42\n    LABEL_SMOOTHING = 0.05\n    # Updated save paths for the new model\n    SAVE_PATH_S1 = \"best_model_effnet_b5_seg_stage1.pth\"\n    SAVE_PATH_FINAL = \"best_model_effnet_b5_seg_final.pth\"\n\ndef seed_everything(seed):\n    os.environ['PYTHONHASHSEED'] = str(seed); np.random.seed(seed); torch.manual_seed(seed); torch.cuda.manual_seed(seed)\n    torch.backends.cudnn.deterministic = True; torch.backends.cudnn.benchmark = True\nseed_everything(CFG.SEED)\n\n# =============================================================================\n# PREPROCESSING & AUGMENTATIONS\n# =============================================================================\ndef preprocess_ben_graham(image, output_size):\n    try:\n        gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)\n        if gray.mean() < 15: image = cv2.resize(image, (output_size, output_size), interpolation=cv2.INTER_AREA)\n        else:\n            _, thresh = cv2.threshold(gray, 15, 255, cv2.THRESH_BINARY)\n            contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)\n            if contours:\n                largest_contour = max(contours, key=cv2.contourArea)\n                x, y, w, h = cv2.boundingRect(largest_contour)\n                image = image[y:y+h, x:x+w]\n            image = cv2.resize(image, (output_size, output_size), interpolation=cv2.INTER_AREA)\n    except Exception: image = cv2.resize(image, (output_size, output_size), interpolation=cv2.INTER_AREA)\n    b, g, r = cv2.split(image)\n    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))\n    g = clahe.apply(g)\n    return cv2.merge((b, g, r))\n\ndef get_transforms(is_train=True):\n    if is_train:\n        # --- ENHANCEMENT: Added stronger augmentations ---\n        return A.Compose([\n            A.HorizontalFlip(p=0.5),\n            A.VerticalFlip(p=0.5),\n            A.ShiftScaleRotate(shift_limit=0.05, scale_limit=0.1, rotate_limit=15, p=0.7),\n            A.GridDistortion(p=0.3),\n            A.ElasticTransform(p=0.3, alpha=120, sigma=120 * 0.05),\n            A.RandomBrightnessContrast(brightness_limit=0.2, contrast_limit=0.2, p=0.7),\n            A.HueSaturationValue(hue_shift_limit=10, sat_shift_limit=20, val_shift_limit=10, p=0.5),\n        ])\n    else:\n        return None\n\n# =============================================================================\n# DATASET - ENHANCED FOR ROBUSTNESS\n# =============================================================================\nclass Dataset4Channel(Dataset):\n    def __init__(self, df, img_dir, seg_dir, transform=None):\n        self.df = df.reset_index(drop=True)\n        self.img_dir = img_dir\n        self.seg_dir = seg_dir\n        self.transform = transform\n        self.post_transform = A.Compose([\n            A.Normalize(mean=[0.485, 0.456, 0.406, 0.5], std=[0.229, 0.224, 0.225, 0.5]),\n            ToTensorV2()\n        ])\n\n    def __len__(self): \n        return len(self.df)\n\n    # --- ENHANCEMENT: This version safely handles missing images or masks ---\n    def __getitem__(self, idx):\n        row = self.df.iloc[idx]\n        img_path = os.path.join(self.img_dir, row['id_code'] + '.png')\n        seg_path = os.path.join(self.seg_dir, row['id_code'] + '.png')\n        \n        if os.path.exists(img_path):\n            img = cv2.imread(img_path)\n            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n        else:\n            img = np.zeros((CFG.IMG_SIZE, CFG.IMG_SIZE, 3), dtype=np.uint8)\n            print(f\"Warning: Image not found at {img_path}, using black image.\")\n\n        if os.path.exists(seg_path):\n            mask = cv2.imread(seg_path, cv2.IMREAD_GRAYSCALE)\n        else:\n            mask = np.zeros((CFG.IMG_SIZE, CFG.IMG_SIZE), dtype=np.uint8)\n            print(f\"Warning: Mask not found at {seg_path}, using black mask.\")\n        \n        img = preprocess_ben_graham(img, CFG.IMG_SIZE)\n        mask = cv2.resize(mask, (CFG.IMG_SIZE, CFG.IMG_SIZE), interpolation=cv2.INTER_NEAREST)\n        \n        if self.transform:\n            augmented = self.transform(image=img, mask=mask)\n            img = augmented['image']\n            mask = augmented['mask']\n            \n        img_4_channel = np.dstack((img, mask))\n        img_4_channel = self.post_transform(image=img_4_channel)['image']\n        label = torch.tensor(row['diagnosis'], dtype=torch.long)\n        return img_4_channel, label\n\n# =============================================================================\n# MODEL DEFINITION\n# =============================================================================\nclass EfficientNet4ChannelOrdinal(nn.Module):\n    def __init__(self, model_name, num_classes=5, pretrained=True):\n        super().__init__()\n        self.backbone = timm.create_model(model_name, pretrained=pretrained, num_classes=0, global_pool='avg')\n        original_conv = self.backbone.conv_stem\n        original_weights = original_conv.weight.clone()\n        new_conv = nn.Conv2d(4, original_conv.out_channels, \n                             kernel_size=original_conv.kernel_size, stride=original_conv.stride, \n                             padding=original_conv.padding, bias=(original_conv.bias is not None))\n        with torch.no_grad():\n            new_conv.weight[:, :3] = original_weights\n            new_conv.weight[:, 3] = original_weights.mean(dim=1)\n        self.backbone.conv_stem = new_conv\n        feature_dim = self.backbone.num_features\n        self.classifier = nn.Sequential(nn.Dropout(0.5), nn.Linear(feature_dim, 256), nn.ReLU(), nn.Dropout(0.3), nn.Linear(256, num_classes - 1))\n    def forward(self, x): \n        return self.classifier(self.backbone(x))\n\n# =============================================================================\n# UTILITY FUNCTIONS\n# =============================================================================\nclass WeightedOrdinalFocalLoss(nn.Module):\n    def __init__(self, num_classes=5, gamma=2.0, class_weights=None, label_smoothing=0.0):\n        super().__init__(); self.num_classes, self.gamma, self.class_weights, self.label_smoothing = num_classes, gamma, class_weights, label_smoothing\n        self.bce = nn.BCEWithLogitsLoss(reduction='none')\n    def forward(self, outputs, targets):\n        ordinal_targets = torch.zeros_like(outputs)\n        for i, t in enumerate(targets):\n            if t > 0: ordinal_targets[i, :t] = 1.0\n        if self.label_smoothing > 0.0: ordinal_targets = ordinal_targets * (1.0 - self.label_smoothing) + 0.5 * self.label_smoothing\n        bce = self.bce(outputs, ordinal_targets)\n        if self.class_weights is not None:\n            weights = self.class_weights[targets].view(-1, 1).expand(-1, outputs.shape[1])\n            bce = bce * weights\n        pt = torch.exp(-bce); focal = (1 - pt) ** self.gamma * bce\n        return focal.mean()\n\nclass SmoothKappaLoss(nn.Module):\n    def __init__(self, num_classes=5, eps=1e-7):\n        super().__init__(); self.num_classes, self.eps = num_classes, eps\n        W = torch.zeros(num_classes, num_classes)\n        for i in range(num_classes):\n            for j in range(num_classes): W[i,j] = ((i - j)**2) / ((num_classes - 1)**2)\n        self.register_buffer(\"W\", W)\n    def forward(self, outputs, targets):\n        device=outputs.device; B=outputs.size(0); probs=torch.sigmoid(outputs)\n        class_probs=torch.zeros(B,self.num_classes,device=device)\n        class_probs[:,0]=1-probs[:,0]\n        for k in range(1,self.num_classes-1): class_probs[:,k]=probs[:,k-1]-probs[:,k]\n        class_probs[:,-1]=probs[:,-1]\n        class_probs=torch.clamp(class_probs,min=self.eps,max=1.0)\n        one_hot=F.one_hot(targets,num_classes=self.num_classes).float().to(device)\n        conf_mat=torch.matmul(one_hot.T,class_probs); hist_true=one_hot.sum(dim=0); hist_pred=class_probs.sum(dim=0)\n        expected=torch.outer(hist_true,hist_pred)\n        W=self.W.to(device); obs=torch.sum(W*conf_mat); exp=torch.sum(W*expected)\n        kappa=1.0-(B*obs)/(exp+self.eps)\n        return 1.0 - kappa\n\ndef mixup_data(x, y, alpha=0.4):\n    if alpha > 0: lam = np.random.beta(alpha, alpha)\n    else: lam = 1\n    batch_size = x.size()[0]; index = torch.randperm(batch_size).to(x.device)\n    mixed_x = lam * x + (1 - lam) * x[index, :]; y_a, y_b = y, y[index]\n    return mixed_x, y_a, y_b, lam\n\ndef ordinal_to_class(outputs): return torch.sum(torch.sigmoid(outputs) > 0.5, dim=1).long()\n\ndef calculate_metrics(outputs, targets):\n    preds, targets_np = ordinal_to_class(outputs).cpu().numpy(), targets.cpu().numpy()\n    return accuracy_score(targets_np, preds), cohen_kappa_score(targets_np, preds, weights='quadratic')\n\ndef clear_memory(): gc.collect(); torch.cuda.empty_cache()\n\n# =============================================================================\n# TRAINING & VALIDATION FUNCTIONS\n# =============================================================================\ndef train_epoch(model, loader, optimizer, criterion, scaler, device, use_mixup):\n    model.train(); running_loss = 0.0; all_out, all_t = [], []\n    pbar = tqdm(loader, desc=\"Training\", leave=False)\n    for images, targets in pbar:\n        images, targets = images.to(device, non_blocking=True), targets.to(device, non_blocking=True)\n        optimizer.zero_grad(set_to_none=True)\n        if use_mixup: images, targets_a, targets_b, lam = mixup_data(images, targets)\n        with torch.cuda.amp.autocast():\n            outputs = model(images)\n            if use_mixup: loss = lam * criterion(outputs, targets_a) + (1 - lam) * criterion(outputs, targets_b)\n            else: loss = criterion(outputs, targets)\n        scaler.scale(loss).backward(); scaler.step(optimizer); scaler.update()\n        running_loss += loss.item(); all_out.append(outputs.detach()); all_t.append(targets.detach())\n        pbar.set_postfix(loss=loss.item())\n    all_out, all_t = torch.cat(all_out), torch.cat(all_t)\n    return running_loss / len(loader), *calculate_metrics(all_out, all_t)\n\ndef validate_epoch(model, loader, criterion, device):\n    model.eval(); running_loss = 0.0; all_out, all_t = [], []\n    with torch.no_grad():\n        pbar = tqdm(loader, desc=\"Validating\", leave=False)\n        for images, targets in pbar:\n            images, targets = images.to(device, non_blocking=True), targets.to(device, non_blocking=True)\n            with torch.cuda.amp.autocast():\n                outputs = model(images); loss = criterion(outputs, targets)\n            running_loss += loss.item(); all_out.append(outputs); all_t.append(targets)\n    all_out, all_t = torch.cat(all_out), torch.cat(all_t)\n    return running_loss / len(loader), *calculate_metrics(all_out, all_t)\n\n# =============================================================================\n# MAIN TRAINING PIPELINE\n# =============================================================================\ndef main():\n    print(f\"Device: {CFG.DEVICE}, Model: {CFG.MODEL_NAME} (4-Channel), Image Size: {CFG.IMG_SIZE}\")\n    try:\n        train_df, val_df = pd.read_csv(CFG.TRAIN_CSV), pd.read_csv(CFG.VAL_CSV)\n    except FileNotFoundError:\n        print(f\"CRITICAL ERROR: Could not find {CFG.TRAIN_CSV} or {CFG.VAL_CSV}.\")\n        print(\"Please ensure you have run the data splitting cell first, or that the paths are correct.\")\n        return\n\n    train_ds = Dataset4Channel(train_df, CFG.TRAIN_DIR, CFG.SEG_TRAIN_DIR, transform=get_transforms(is_train=True))\n    val_ds   = Dataset4Channel(val_df, CFG.VAL_DIR, CFG.SEG_VAL_DIR, transform=get_transforms(is_train=False))\n\n    class_weights = compute_class_weight('balanced', classes=np.unique(train_df['diagnosis']), y=train_df['diagnosis'])\n    sample_weights = np.array([class_weights[int(l)] for l in train_df['diagnosis']])\n    sampler = WeightedRandomSampler(sample_weights, num_samples=len(sample_weights), replacement=True)\n    train_loader = DataLoader(train_ds, batch_size=CFG.BATCH_SIZE, sampler=sampler, num_workers=CFG.NUM_WORKERS, pin_memory=True)\n    val_loader = DataLoader(val_ds, batch_size=CFG.BATCH_SIZE*2, shuffle=False, num_workers=CFG.NUM_WORKERS, pin_memory=True)\n    \n    model = EfficientNet4ChannelOrdinal(CFG.MODEL_NAME).to(CFG.DEVICE)\n    class_weights_loss = torch.tensor(class_weights, dtype=torch.float).to(CFG.DEVICE)\n    focal_loss = WeightedOrdinalFocalLoss(num_classes=5, gamma=2.0, class_weights=class_weights_loss, label_smoothing=CFG.LABEL_SMOOTHING)\n    kappa_loss = SmoothKappaLoss(num_classes=5)\n    def hybrid_loss(outputs, targets): return 0.7 * kappa_loss(outputs, targets) + 0.3 * focal_loss(outputs, targets)\n    \n    scaler = torch.cuda.amp.GradScaler()\n\n    # --- STAGE 1 ---\n    print(\"\\n\" + \"=\"*50 + \"\\n     STARTING STAGE 1 (4-Channel)\\n\" + \"=\"*50)\n    opt = optim.AdamW(model.parameters(), lr=CFG.S1_LR, weight_decay=1e-4)\n    # --- ENHANCEMENT: Using improved scheduler ---\n    sched = optim.lr_scheduler.CosineAnnealingWarmRestarts(opt, T_0=10, T_mult=1, eta_min=1e-6)\n    best_val_qwk, patience_counter = -1, 0\n\n    for epoch in range(CFG.S1_EPOCHS):\n        clear_memory(); print(f\"\\nEpoch {epoch+1}/{CFG.S1_EPOCHS}\")\n        train_loss, train_acc, train_qwk = train_epoch(model, train_loader, opt, focal_loss, scaler, CFG.DEVICE, CFG.S1_USE_MIXUP)\n        val_loss, val_acc, val_qwk = validate_epoch(model, val_loader, focal_loss, CFG.DEVICE)\n        sched.step()\n        print(f\"Train -> Loss:{train_loss:.4f} Acc:{train_acc:.4f} QWK:{train_qwk:.4f}\")\n        print(f\"Valid -> Loss:{val_loss:.4f} Acc:{val_acc:.4f} QWK:{val_qwk:.4f}\")\n        if val_qwk > best_val_qwk:\n            print(f\"Val QWK improved from {best_val_qwk:.4f} to {val_qwk:.4f}. Saving model...\")\n            best_val_qwk, patience_counter = val_qwk, 0\n            torch.save(model.state_dict(), CFG.SAVE_PATH_S1)\n        else:\n            patience_counter += 1\n            if patience_counter >= CFG.PATIENCE: print(\"Early stopping in Stage 1.\"); break\n    \n    # --- STAGE 2 ---\n    print(\"\\n\" + \"=\"*50 + \"\\n     STARTING STAGE 2 (4-Channel)\\n\" + \"=\"*50)\n    if os.path.exists(CFG.SAVE_PATH_S1): model.load_state_dict(torch.load(CFG.SAVE_PATH_S1))\n    else: print(\"No Stage 1 model was saved. Continuing with the current model.\")\n\n    opt = optim.AdamW(model.parameters(), lr=CFG.S2_LR, weight_decay=1e-5)\n    # --- ENHANCEMENT: Using improved scheduler ---\n    sched = optim.lr_scheduler.CosineAnnealingWarmRestarts(opt, T_0=10, T_mult=1, eta_min=3e-6)\n    best_val_qwk_stage2, patience_counter = best_val_qwk, 0\n\n    for epoch in range(CFG.S2_EPOCHS):\n        clear_memory(); print(f\"\\nEpoch {epoch+1}/{CFG.S2_EPOCHS}\")\n        train_loss, train_acc, train_qwk = train_epoch(model, train_loader, opt, hybrid_loss, scaler, CFG.DEVICE, CFG.S2_USE_MIXUP)\n        val_loss, val_acc, val_qwk = validate_epoch(model, val_loader, hybrid_loss, CFG.DEVICE)\n        sched.step()\n        print(f\"Train -> Loss:{train_loss:.4f} Acc:{train_acc:.4f} QWK:{train_qwk:.4f}\")\n        print(f\"Valid -> Loss:{val_loss:.4f} Acc:{val_acc:.4f} QWK:{val_qwk:.4f}\")\n        if val_qwk > best_val_qwk_stage2:\n            print(f\"Val QWK improved from {best_val_qwk_stage2:.4f} to {val_qwk:.4f}. Saving final model...\")\n            best_val_qwk_stage2, patience_counter = val_qwk, 0\n            torch.save(model.state_dict(), CFG.SAVE_PATH_FINAL)\n        else:\n            patience_counter += 1\n            if patience_counter >= CFG.PATIENCE: print(\"Early stopping in Stage 2.\"); break\n\n    print(f\"\\nTraining Finished!\\nFinal Best QWK: {best_val_qwk_stage2:.4f}\")\n\nif __name__ == \"__main__\":\n    main()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-13T21:32:42.047640Z","iopub.execute_input":"2025-09-13T21:32:42.047946Z","iopub.status.idle":"2025-09-13T23:52:53.822585Z","shell.execute_reply.started":"2025-09-13T21:32:42.047924Z","shell.execute_reply":"2025-09-13T23:52:53.821448Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nimport numpy as np\nimport pandas as pd\nimport cv2\nimport os\nimport timm\nfrom tqdm.auto import tqdm\nfrom sklearn.metrics import cohen_kappa_score, accuracy_score, classification_report, confusion_matrix\nimport seaborn as sns\nimport matplotlib.pyplot as plt\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\nfrom torch.utils.data import Dataset, DataLoader\nfrom scipy.optimize import minimize\n\n# =============================================================================\n# CONFIGURATION FOR 4-CHANNEL EFFICIENTNET-B5\n# =============================================================================\nclass CFG:\n    # --- MODEL CONFIG ---\n    MODEL_NAME = 'efficientnet_b5'\n    IMG_SIZE = 384\n    \n    # --- PATHS ---\n    BASE_PATH = \"/kaggle/input/aptos2019\"\n    VAL_CSV = os.path.join(BASE_PATH, \"valid.csv\")\n    VAL_DIR = os.path.join(BASE_PATH, \"val_images\", \"val_images\")\n    TEST_CSV = os.path.join(BASE_PATH, \"test.csv\")\n    TEST_DIR = os.path.join(BASE_PATH, \"test_images\", \"test_images\")\n    \n    # --- PATHS TO YOUR SEGMENTED MASKS ---\n    SEG_BASE_PATH = \"/kaggle/input/segmentaion-dataset\"\n    SEG_VAL_DIR   = os.path.join(SEG_BASE_PATH, \"segmented_outputs_val/segmented_outputs_val/\")\n    SEG_TEST_DIR  = os.path.join(SEG_BASE_PATH, \"segmented_outputs_test/segmented_outputs_test/\") # Assuming this is your test mask path\n    \n    # Path to your saved 4-channel model\n    MODEL_PATH = \"/kaggle/working/best_model_effnet_b5_seg_final.pth\"\n    \n    # --- INFERENCE CONFIG ---\n    BATCH_SIZE = 8 \n    DEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n    NUM_WORKERS = 2\n\n# =============================================================================\n# REUSED CLASSES & PREPROCESSING\n# =============================================================================\nclass EfficientNet4ChannelOrdinal(nn.Module):\n    def __init__(self, model_name, num_classes=5, pretrained=False):\n        super().__init__()\n        self.backbone = timm.create_model(model_name, pretrained=pretrained, num_classes=0, global_pool='avg')\n        original_conv = self.backbone.conv_stem; original_weights = original_conv.weight.clone()\n        new_conv = nn.Conv2d(4, original_conv.out_channels, \n                             kernel_size=original_conv.kernel_size, stride=original_conv.stride, \n                             padding=original_conv.padding, bias=(original_conv.bias is not None))\n        with torch.no_grad():\n            new_conv.weight[:, :3] = original_weights\n            new_conv.weight[:, 3] = original_weights.mean(dim=1)\n        self.backbone.conv_stem = new_conv\n        feature_dim = self.backbone.num_features\n        self.classifier = nn.Sequential(nn.Dropout(0.5), nn.Linear(feature_dim, 256), nn.ReLU(), nn.Dropout(0.3), nn.Linear(256, num_classes - 1))\n    def forward(self, x): return self.classifier(self.backbone(x))\n\ndef preprocess_ben_graham(image, output_size):\n    try:\n        gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)\n        if gray.mean() < 15: image = cv2.resize(image, (output_size, output_size), interpolation=cv2.INTER_AREA)\n        else:\n            _, thresh = cv2.threshold(gray, 15, 255, cv2.THRESH_BINARY)\n            contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)\n            if contours:\n                largest_contour = max(contours, key=cv2.contourArea); x, y, w, h = cv2.boundingRect(largest_contour)\n                image = image[y:y+h, x:x+w]\n            image = cv2.resize(image, (output_size, output_size), interpolation=cv2.INTER_AREA)\n    except Exception: image = cv2.resize(image, (output_size, output_size), interpolation=cv2.INTER_AREA)\n    b, g, r = cv2.split(image); clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8)); g = clahe.apply(g)\n    return cv2.merge((b, g, r))\n\nclass Dataset4Channel(Dataset):\n    def __init__(self, df, img_dir, seg_dir, transform=None):\n        self.df, self.img_dir, self.seg_dir, self.transform = df.reset_index(drop=True), img_dir, seg_dir, transform\n        self.post_transform = A.Compose([\n            A.Normalize(mean=[0.485, 0.456, 0.406, 0.5], std=[0.229, 0.224, 0.225, 0.5]), ToTensorV2()])\n    def __len__(self): return len(self.df)\n    def __getitem__(self, idx):\n        row = self.df.iloc[idx]; img_path = os.path.join(self.img_dir, row['id_code'] + '.png'); seg_path = os.path.join(self.seg_dir, row['id_code'] + '.png')\n        img = cv2.imread(img_path); img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n        mask = cv2.imread(seg_path, cv2.IMREAD_GRAYSCALE)\n        img = preprocess_ben_graham(img, CFG.IMG_SIZE)\n        mask = cv2.resize(mask, (CFG.IMG_SIZE, CFG.IMG_SIZE), interpolation=cv2.INTER_NEAREST)\n        if self.transform:\n            augmented = self.transform(image=img, mask=mask); img = augmented['image']; mask = augmented['mask']\n        img_4_channel = np.dstack((img, mask))\n        img_4_channel = self.post_transform(image=img_4_channel)['image']\n        label = torch.tensor(row['diagnosis'], dtype=torch.long)\n        return img_4_channel, label\n\n# =============================================================================\n# OPTIMIZATION FUNCTIONS\n# =============================================================================\ndef ordinal_to_class_with_thresholds(outputs, thresholds):\n    probs = torch.sigmoid(outputs).cpu().numpy()\n    preds = np.sum(probs > thresholds, axis=1)\n    return preds\n\ndef kappa_objective(thresholds, outputs, targets):\n    preds = ordinal_to_class_with_thresholds(outputs, thresholds)\n    return -cohen_kappa_score(targets, preds, weights=\"quadratic\")\n\ndef find_best_thresholds(outputs, targets):\n    print(\"Finding optimal thresholds...\")\n    outputs = outputs.detach(); targets = targets.cpu().numpy()\n    init_thresh = np.array([0.5, 0.5, 0.5, 0.5]); bounds = [(0.1, 0.9)] * len(init_thresh)\n    res = minimize(kappa_objective, init_thresh, args=(outputs, targets), method=\"Powell\", bounds=bounds)\n    best_thresholds = res.x\n    print(f\"Optimal thresholds found: {np.round(best_thresholds, 4)}\")\n    return best_thresholds\n\n# =============================================================================\n# MAIN SCRIPT\n# =============================================================================\ndef run_optimization_and_test():\n    # --- Step 0: Load Model ---\n    model = EfficientNet4ChannelOrdinal(CFG.MODEL_NAME, pretrained=False).to(CFG.DEVICE)\n    model.load_state_dict(torch.load(CFG.MODEL_PATH, map_location=CFG.DEVICE))\n    model.eval()\n    print(f\"Model loaded successfully from {CFG.MODEL_PATH}\")\n\n    # --- Step 1: Get Raw Answers on the Validation Set ---\n    print(\"\\n--- Step 1: Evaluating on Validation Set to find thresholds ---\")\n    val_df = pd.read_csv(CFG.VAL_CSV)\n    val_dataset = Dataset4Channel(val_df, CFG.VAL_DIR, CFG.SEG_VAL_DIR, transform=None) # No augs for val\n    val_loader = DataLoader(val_dataset, batch_size=CFG.BATCH_SIZE, shuffle=False, num_workers=CFG.NUM_WORKERS)\n    \n    val_outputs_list, val_labels_list = [], []\n    with torch.no_grad():\n        for images, labels in tqdm(val_loader, desc=\"Getting Validation Outputs\"):\n            images = images.to(CFG.DEVICE); outputs = model(images)\n            val_outputs_list.append(outputs.cpu()); val_labels_list.append(labels)\n    val_outputs = torch.cat(val_outputs_list); val_labels = torch.cat(val_labels_list)\n    \n    # --- Step 2: Find the Perfect \"Grading Scale\" ---\n    print(\"\\n--- Step 2: Optimizing Thresholds ---\")\n    best_thresholds = find_best_thresholds(val_outputs, val_labels)\n\n    # --- Step 3: Use the Grading Scale on the Test Set ---\n    print(\"\\n--- Step 3: Evaluating on Test Set with new thresholds ---\")\n    test_df = pd.read_csv(CFG.TEST_CSV)\n    test_dataset = Dataset4Channel(test_df, CFG.TEST_DIR, CFG.SEG_TEST_DIR, transform=None) # No augs for test\n    test_loader = DataLoader(test_dataset, batch_size=CFG.BATCH_SIZE, shuffle=False, num_workers=CFG.NUM_WORKERS)\n\n    test_outputs_list, test_labels_list = [], []\n    with torch.no_grad():\n        for images, labels in tqdm(test_loader, desc=\"Getting Test Outputs\"):\n            images = images.to(CFG.DEVICE); outputs = model(images)\n            test_outputs_list.append(outputs.cpu()); test_labels_list.append(labels)\n    test_outputs = torch.cat(test_outputs_list); test_labels = torch.cat(test_labels_list).numpy()\n    \n    # --- FINAL RESULTS ---\n    print(\"\\n\" + \"=\"*50)\n    print(\"      FINAL RESULTS COMPARISON (EffNet-B5 + Segmentation)\")\n    print(\"=\"*50)\n\n    preds_old = torch.sum(torch.sigmoid(test_outputs) > 0.5, dim=1).numpy()\n    qwk_old = cohen_kappa_score(test_labels, preds_old, weights='quadratic')\n    acc_old = accuracy_score(test_labels, preds_old)\n    print(f\"\\nOriginal Score (Threshold = 0.5):\"); print(f\"  QWK: {qwk_old:.4f}\"); print(f\"  Accuracy: {acc_old*100:.2f}%\")\n\n    preds_new = ordinal_to_class_with_thresholds(test_outputs, best_thresholds)\n    qwk_new = cohen_kappa_score(test_labels, preds_new, weights='quadratic')\n    acc_new = accuracy_score(test_labels, preds_new)\n    print(f\"\\nPolished Score (Optimized Thresholds):\"); print(f\"  QWK: {qwk_new:.4f}\"); print(f\"  Accuracy: {acc_new*100:.2f}%\")\n    print(\"=\"*50)\n    \n    print(\"\\n--- Polished Classification Report ---\"); print(classification_report(test_labels, preds_new, target_names=[f\"Class {i}\" for i in range(5)]))\n    \n    print(\"\\n--- Polished Confusion Matrix ---\")\n    cm = confusion_matrix(test_labels, preds_new)\n    plt.figure(figsize=(8, 6)); sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', xticklabels=range(5), yticklabels=range(5))\n    plt.xlabel(\"Predicted Label (Optimized)\"); plt.ylabel(\"True Label\"); plt.title(\"Confusion Matrix (Optimized Thresholds)\")\n    plt.show()\n    \nrun_optimization_and_test()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-14T00:12:49.759089Z","iopub.execute_input":"2025-09-14T00:12:49.759435Z","iopub.status.idle":"2025-09-14T00:14:09.293409Z","shell.execute_reply.started":"2025-09-14T00:12:49.759395Z","shell.execute_reply":"2025-09-14T00:14:09.292753Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}