{"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":44890,"databundleVersionId":4852232,"sourceType":"competition"},{"sourceId":7866129,"sourceType":"datasetVersion","datasetId":4614938},{"sourceId":13063258,"sourceType":"datasetVersion","datasetId":8272564},{"sourceId":670308,"sourceType":"modelInstanceVersion","isSourceIdPinned":true,"modelInstanceId":507650,"modelId":522348}],"dockerImageVersionId":31193,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# Cell 1: Import libraries and setup\nimport warnings\nwarnings.filterwarnings('ignore')\nimport pandas as pd\nimport numpy as np\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nimport torchvision.transforms as transforms\nimport timm\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\nimport cv2\nimport os\nimport matplotlib.pyplot as plt\nimport seaborn as sns\nfrom sklearn.model_selection import train_test_split, StratifiedKFold\nfrom sklearn.metrics import cohen_kappa_score, accuracy_score, confusion_matrix, classification_report\nfrom tqdm import tqdm\nimport random","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-12-02T19:56:43.288827Z","iopub.execute_input":"2025-12-02T19:56:43.289460Z","iopub.status.idle":"2025-12-02T19:56:43.295033Z","shell.execute_reply.started":"2025-12-02T19:56:43.289436Z","shell.execute_reply":"2025-12-02T19:56:43.294240Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Set paths\nTRAIN_PATH = '/kaggle/input/diabetic-retinopathy-classification-3/kaggle_retina/train'\nTEST_PATH = '/kaggle/input/diabetic-retinopathy-classification-3/kaggle_retina/test'\nTRAIN_CSV = '/kaggle/input/diabetic-retinopathy-classification-3/kaggle_retina/train.csv'\nSUBMISSION_CSV = '/kaggle/input/diabetic-retinopathy-classification-3/kaggle_retina/submission.csv'","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-02T19:56:49.253521Z","iopub.execute_input":"2025-12-02T19:56:49.253814Z","iopub.status.idle":"2025-12-02T19:56:49.258002Z","shell.execute_reply.started":"2025-12-02T19:56:49.253793Z","shell.execute_reply":"2025-12-02T19:56:49.257186Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 1. Load Data\ntrain_df = pd.read_csv(TRAIN_CSV)\nsubmission_df = pd.read_csv(SUBMISSION_CSV)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-02T19:56:54.365095Z","iopub.execute_input":"2025-12-02T19:56:54.365416Z","iopub.status.idle":"2025-12-02T19:56:54.396086Z","shell.execute_reply.started":"2025-12-02T19:56:54.365395Z","shell.execute_reply":"2025-12-02T19:56:54.395214Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(f\"Training set: {len(train_df)} images\")\nprint(f\"Test set: {len(submission_df)} images\")\nprint(\"\\nSample of training data:\")\nprint(train_df.head())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-02T18:49:17.515750Z","iopub.execute_input":"2025-12-02T18:49:17.516296Z","iopub.status.idle":"2025-12-02T18:49:17.535137Z","shell.execute_reply.started":"2025-12-02T18:49:17.516271Z","shell.execute_reply":"2025-12-02T18:49:17.534575Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 2. Class Distribution\nclass_counts = train_df['diagnosis'].value_counts().sort_index()\nnum_classes = len(class_counts)\nprint(f\"\\nNumber of classes: {num_classes}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-02T18:49:21.379864Z","iopub.execute_input":"2025-12-02T18:49:21.380416Z","iopub.status.idle":"2025-12-02T18:49:21.397729Z","shell.execute_reply.started":"2025-12-02T18:49:21.380394Z","shell.execute_reply":"2025-12-02T18:49:21.396959Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import matplotlib.pyplot as plt\nimport numpy as np\nimport cv2\nimport os\nfrom matplotlib.gridspec import GridSpec\n\n# Create figure\nfig = plt.figure(figsize=(15, 8))\ngs = GridSpec(2, 5, height_ratios=[3, 1])\n\n# Define class colors\ncolors = ['#3498db', '#2ecc71', '#f1c40f', '#e74c3c', '#9b59b6']\n\n# For each class, display an example image and stats\nfor i, class_id in enumerate(sorted(class_counts.index)):\n    # Select a random image from this class\n    class_samples = train_df[train_df['diagnosis'] == class_id]\n    if len(class_samples) > 0:\n        random_sample = class_samples.sample(1).iloc[0]\n        # Using id_code column and TRAIN_PATH variable\n        img_filename = f\"{random_sample['id_code']}.png\"  # Adjust file extension if needed\n        img_path = os.path.join(TRAIN_PATH, img_filename)\n\n        # Read the image\n        try:\n            img = cv2.imread(img_path)\n            if img is not None:\n                img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n            else:\n                print(f\"Warning: Could not load image at {img_path}\")\n                # Create a placeholder if image cannot be loaded\n                img = np.ones((100, 100, 3), dtype=np.uint8) * 200\n                cv2.putText(img, f\"Class {class_id}\", (15, 50), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1)\n        except Exception as e:\n            print(f\"Error loading image {img_path}: {e}\")\n            # Create a placeholder if there's an error\n            img = np.ones((100, 100, 3), dtype=np.uint8) * 200\n            cv2.putText(img, f\"Class {class_id}\", (15, 50), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1)\n\n    else:\n        # Create a placeholder if no images for this class\n        img = np.ones((100, 100, 3), dtype=np.uint8) * 200\n        cv2.putText(img, f\"No images\", (15, 50), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1)\n\n    # Display the image\n    ax = fig.add_subplot(gs[0, i])\n    ax.imshow(img)\n    ax.set_title(f\"Class {class_id}\", fontsize=12)\n    ax.axis('off')\n\n    # Add count information below the image\n    count = class_counts.get(class_id, 0)\n    percentage = count / train_df.shape[0] * 100\n    ax.text(0.5, -0.15, f\"{count} images\\n({percentage:.2f}%)\",\n            horizontalalignment='center', verticalalignment='center',\n            transform=ax.transAxes, fontsize=11)\n\n# Create a horizontal bar chart for total distribution\nax_bar = fig.add_subplot(gs[1, :])\n\n# Calculate positions for stacked bars\nleft = 0\nfor i, class_id in enumerate(sorted(class_counts.index)):\n    count = class_counts.get(class_id, 0)\n    width = count / train_df.shape[0]\n    ax_bar.barh(0, width, left=left, color=colors[i % len(colors)], edgecolor='white')\n\n    # Add text in the middle of each segment if wide enough\n    if width > 0.05:  # Only add text if segment is wide enough\n        ax_bar.text(left + width/2, 0, f\"{class_id}: {width*100:.1f}%\",\n                   ha='center', va='center', color='white', fontweight='bold')\n\n    left += width\n\nax_bar.set_xlim(0, 1)\nax_bar.set_yticks([])\nax_bar.set_title(\"Class Distribution\", fontsize=12)\nax_bar.set_xlabel(\"Proportion of Dataset\")\n\nplt.tight_layout()\nplt.show()\n\n# Optional: Let's print the path of one example image to verify we're looking in the right place\nif len(class_samples) > 0:\n    print(f\"Example image path: {img_path}\")\n    print(f\"Does file exist? {os.path.exists(img_path)}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-02T18:49:29.040540Z","iopub.execute_input":"2025-12-02T18:49:29.041301Z","iopub.status.idle":"2025-12-02T18:49:29.973996Z","shell.execute_reply.started":"2025-12-02T18:49:29.041276Z","shell.execute_reply":"2025-12-02T18:49:29.973280Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from PIL import Image\n\n# 4. Display Random Images from Each Class\nprint(\"\\nDisplaying random samples from each class...\")\n\n# Sample 10 images from each class\nfor class_id in range(num_classes):\n    class_images = train_df[train_df['diagnosis'] == class_id]['id_code'].tolist()\n    samples = random.sample(class_images, min(10, len(class_images)))\n\n    plt.figure(figsize=(20, 8))\n    plt.suptitle(f'Class {class_id} Samples', fontsize=16)\n\n    for i, img_id in enumerate(samples):\n        img_path = os.path.join(TRAIN_PATH, f\"{img_id}.png\")\n        img = Image.open(img_path)\n\n        plt.subplot(2, 5, i+1)\n        plt.imshow(img)\n        plt.axis('off')\n        plt.title(f\"{img_id[:6]}\")\n\n    plt.tight_layout(rect=[0, 0, 1, 0.95])\n    plt.show()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-02T18:49:36.177352Z","iopub.execute_input":"2025-12-02T18:49:36.178216Z","iopub.status.idle":"2025-12-02T18:49:45.617153Z","shell.execute_reply.started":"2025-12-02T18:49:36.178180Z","shell.execute_reply":"2025-12-02T18:49:45.616270Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nfrom PIL import Image, ImageFile # Import ImageFile to handle truncated images\nimport hashlib\nimport pandas as pd\nfrom collections import defaultdict\n\n# Allow Pillow to load truncated images (can be useful for partially corrupted files)\nImageFile.LOAD_TRUNCATED_IMAGES = True\n\ndef get_image_files(directory):\n    \"\"\"\n    Helper function to get a list of image files in a given directory.\n    Filters for common image extensions.\n    \"\"\"\n    image_files = []\n    for filename in os.listdir(directory):\n        filepath = os.path.join(directory, filename)\n        if os.path.isfile(filepath) and filename.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif')):\n            image_files.append(filepath)\n    return image_files\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-02T18:50:08.681243Z","iopub.execute_input":"2025-12-02T18:50:08.681583Z","iopub.status.idle":"2025-12-02T18:50:08.687232Z","shell.execute_reply.started":"2025-12-02T18:50:08.681556Z","shell.execute_reply":"2025-12-02T18:50:08.686437Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def find_corrupted_images(image_dir):\n    \"\"\"\n    Identifies corrupted or unreadable images in a given directory.\n    Uses Image.verify() to check file integrity.\n    \"\"\"\n    print(f\"\\n--- Checking for Corrupted Images in: {image_dir} ---\")\n    corrupted_images = []\n    total_files_checked = 0\n\n    image_files = get_image_files(image_dir)\n    if not image_files:\n        print(f\"No image files found in {image_dir}.\")\n        return []\n\n    for filepath in image_files:\n        total_files_checked += 1\n        try:\n            with Image.open(filepath) as img:\n                # The .verify() method attempts to read the image data fully,\n                # raising an exception if it encounters issues.\n                img.verify()\n        except Exception as e:\n            corrupted_images.append((os.path.basename(filepath), str(e)))\n        finally:\n            # Important: Close the image file after verify to release resources\n            # if img.verify() doesn't implicitly close it or if it errors out before closing.\n            # However, 'with Image.open()' handles closing.\n            pass\n\n    if not corrupted_images:\n        print(f\"No corrupted images found among {total_files_checked} image files.\")\n    else:\n        print(f\"Found {len(corrupted_images)} corrupted images out of {total_files_checked} image files:\")\n        for img_name, error_msg in corrupted_images:\n            print(f\"  {img_name} - Error: {error_msg}\")\n    return corrupted_images\n\n\n# --- Main Execution ---\nif __name__ == \"__main__\":\n    print(\"Starting EDA for Diabetic Retinopathy Datasets...\")\n\n    # Analyze Training Data\n    corrupted_train_images = find_corrupted_images(TRAIN_PATH)\n\n    print(\"\\n\" + \"=\"*50)\n\n    # Analyze Test Data\n    corrupted_test_images = find_corrupted_images(TEST_PATH)\n\n    print(\"\\nEDA Complete!\")\n    print(\"\\nSummary of findings:\")\n    print(f\"  Corrupted training images: {len(corrupted_train_images)}\")\n    print(f\"  Corrupted test images: {len(corrupted_test_images)}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-02T18:50:11.761401Z","iopub.execute_input":"2025-12-02T18:50:11.761722Z","iopub.status.idle":"2025-12-02T18:50:52.758946Z","shell.execute_reply.started":"2025-12-02T18:50:11.761698Z","shell.execute_reply":"2025-12-02T18:50:52.758100Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport numpy as np\nimport cv2\nimport matplotlib.pyplot as plt\nfrom tqdm import tqdm\nimport random\nimport pandas as pd\n\n# Paths\nPROCESSED_TRAIN_PATH = '/kaggle/input/diabetic-retinopathy-classification-3/kaggle_retina/train'\nTRAIN_CSV = '/kaggle/input/diabetic-retinopathy-classification-3/kaggle_retina/train.csv'\n\n# Use available path\nIMAGE_PATH = PROCESSED_TRAIN_PATH if os.path.exists(PROCESSED_TRAIN_PATH) else ORIGINAL_TRAIN_PATH\n\n# Load CSV and build label dictionary\ntrain_df = pd.read_csv(TRAIN_CSV)\nlabel_dict = dict(zip(train_df['id_code'], train_df['diagnosis']))\n\n# Function to apply various enhancements to retinal images - improved version\ndef enhance_retinal_image(img, method):\n    \"\"\"\n    Apply different enhancement methods to retinal images to highlight features\n    \"\"\"\n    # Ensure we're working with RGB images\n    if len(img.shape) == 2:\n        img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)\n    elif img.shape[2] == 4:  # RGBA\n        img = cv2.cvtColor(img, cv2.COLOR_RGBA2RGB)\n\n    # Create a mask for non-black areas (retina region)\n    mask = (img.sum(axis=2) > 0).astype(np.uint8) * 255\n\n    if method == 'original':\n        return img\n\n    elif method == 'clahe':\n        # Convert to LAB color space for CLAHE\n        lab = cv2.cvtColor(img, cv2.COLOR_RGB2LAB)\n        # Apply CLAHE to L channel\n        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))\n        lab[..., 0] = clahe.apply(lab[..., 0])\n        # Convert back to RGB\n        enhanced = cv2.cvtColor(lab, cv2.COLOR_LAB2RGB)\n        return enhanced\n\n    elif method == 'green_channel':\n        # Extract green channel (best for vessel visualization)\n        green = img.copy()\n        green[:, :, 0] = 0  # Remove red\n        green[:, :, 2] = 0  # Remove blue\n        return green\n\n    elif method == 'edge_detection':\n        # Convert to grayscale\n        gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)\n        # Apply median blur to reduce noise while preserving edges\n        blurred = cv2.medianBlur(gray, 5)\n        # Apply adaptive thresholding to highlight vessels\n        thresh = cv2.adaptiveThreshold(\n            blurred, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,\n            cv2.THRESH_BINARY_INV, 11, 2)\n        # Apply mask to keep only the retina region\n        thresh = cv2.bitwise_and(thresh, thresh, mask=mask)\n        # Convert back to RGB for visualization\n        edges_rgb = cv2.cvtColor(thresh, cv2.COLOR_GRAY2RGB)\n        return edges_rgb\n\n\n    elif method == 'heatmap':\n        # Convert to grayscale\n        gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)\n        # Apply CLAHE for contrast enhancement\n        clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8, 8))\n        enhanced = clahe.apply(gray)\n        # Apply colormap for heatmap visualization\n        heatmap = cv2.applyColorMap(enhanced, cv2.COLORMAP_JET)\n        # Convert from BGR to RGB\n        heatmap = cv2.cvtColor(heatmap, cv2.COLOR_BGR2RGB)\n        # Apply mask to keep only the retina region\n        heatmap_masked = cv2.bitwise_and(heatmap, heatmap, mask=mask)\n        return heatmap_masked\n\n    # Default case\n    return img\n\n# Visualize one image with enhancements and save each output\ndef visualize_enhancements(image_path, label=None, save_output=False, save_dir='/kaggle/working/method_train'):\n    img = cv2.imread(image_path)\n    if img is None:\n        print(f\"Error: Could not read image at {image_path}\")\n        return\n\n    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n    methods = [\n        'original', 'clahe', 'green_channel', 'edge_detection', 'heatmap'\n    ]\n\n    fig, axes = plt.subplots(1, 5, figsize=(20, 10))\n    axes = axes.flatten()\n\n    img_id = os.path.splitext(os.path.basename(image_path))[0]\n\n    for i, method in enumerate(methods):\n        enhanced = enhance_retinal_image(img, method)\n        title = method.replace('_', ' ').title()\n        if method == 'original':\n            title += f\"\\nLabel: {label}\"\n        axes[i].imshow(enhanced)\n        axes[i].set_title(title)\n        axes[i].axis('off')\n\n        # Save each enhanced image into method folder\n        if save_output:\n            method_save_path = os.path.join(save_dir, method)\n            os.makedirs(method_save_path, exist_ok=True)\n            out_img = cv2.cvtColor(enhanced, cv2.COLOR_RGB2BGR)\n            out_path = os.path.join(method_save_path, f\"{img_id}.png\")\n            cv2.imwrite(out_path, out_img)\n\n    plt.tight_layout()\n    plt.show()\n\n# Process and save ALL images, not just samples\ndef visualize_all_enhancements(image_paths, save_output=True):\n    for i, img_path in enumerate(tqdm(image_paths, desc=\"Processing Images\")):\n        img_id = os.path.splitext(os.path.basename(img_path))[0]\n        label = label_dict.get(img_id, \"Unknown\")\n        visualize_enhancements(img_path, label, save_output=save_output)\n\n# Get image paths\nimage_paths = [os.path.join(IMAGE_PATH, filename)\n               for filename in os.listdir(IMAGE_PATH)\n               if filename.lower().endswith(('.png', '.jpg', '.jpeg'))]\n\nif image_paths:\n    random_image_path = random.choice(image_paths)\n    img_id = os.path.splitext(os.path.basename(random_image_path))[0]\n    label = label_dict.get(img_id, \"Unknown\")\n\n    print(f\"Showing enhancements for a single image: {img_id}, Label: {label}\")\n    visualize_enhancements(random_image_path, label=label, save_output=True)\nelse:\n    print(f\"No images found in {IMAGE_PATH}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-02T18:51:18.320749Z","iopub.execute_input":"2025-12-02T18:51:18.321041Z","iopub.status.idle":"2025-12-02T18:51:18.970777Z","shell.execute_reply.started":"2025-12-02T18:51:18.321018Z","shell.execute_reply":"2025-12-02T18:51:18.969816Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport random\nimport cv2\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\n\n# Paths\nPROCESSED_TRAIN_PATH = '/kaggle/input/diabetic-retinopathy-classification-3/kaggle_retina/train'\nTRAIN_CSV = '/kaggle/input/diabetic-retinopathy-classification-3/kaggle_retina/train.csv'\n\n# Load image paths and labels\nimage_paths = [os.path.join(PROCESSED_TRAIN_PATH, filename)\n               for filename in os.listdir(PROCESSED_TRAIN_PATH)\n               if filename.lower().endswith(('.png', '.jpg', '.jpeg'))]\n\ntrain_df = pd.read_csv(TRAIN_CSV)\nlabel_dict = dict(zip(train_df['id_code'], train_df['diagnosis']))\n\ndef clahe_enhance_with_resize(img, target_size=386):\n    # Ensure RGB\n    if len(img.shape) == 2:\n        img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)\n    elif img.shape[2] == 4:\n        img = cv2.cvtColor(img, cv2.COLOR_RGBA2RGB)\n\n    # --- Resize with aspect ratio preserved and padding ---\n    h, w = img.shape[:2]\n    scale = target_size / max(h, w)\n    new_h, new_w = int(h * scale), int(w * scale)\n    resized = cv2.resize(img, (new_w, new_h))\n\n    top = (target_size - new_h) // 2\n    bottom = target_size - new_h - top\n    left = (target_size - new_w) // 2\n    right = target_size - new_w - left\n\n    padded = cv2.copyMakeBorder(resized, top, bottom, left, right,\n                                borderType=cv2.BORDER_CONSTANT, value=[0, 0, 0])\n\n    # --- Apply CLAHE ---\n    lab = cv2.cvtColor(padded, cv2.COLOR_RGB2LAB)\n    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))\n    lab[..., 0] = clahe.apply(lab[..., 0])\n    enhanced = cv2.cvtColor(lab, cv2.COLOR_LAB2RGB)\n    return enhanced\n\n\n# Visualize 8 random CLAHE-enhanced images\ndef visualize_8_random_clahe():\n    selected = random.sample(image_paths, 8)\n\n    fig, axes = plt.subplots(2, 4, figsize=(20, 10))\n    axes = axes.flatten()\n\n    for i, img_path in enumerate(selected):\n        img_id = os.path.splitext(os.path.basename(img_path))[0]\n        label = label_dict.get(img_id, \"Unknown\")\n        img = cv2.imread(img_path)\n        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n        clahe_img = clahe_enhance_with_resize(img)\n        axes[i].imshow(clahe_img)\n        axes[i].set_title(f\"Label: {label}\")\n        axes[i].axis('off')\n\n    plt.tight_layout()\n    plt.show()\n\n# Run\nvisualize_8_random_clahe()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-02T18:51:32.222478Z","iopub.execute_input":"2025-12-02T18:51:32.223258Z","iopub.status.idle":"2025-12-02T18:51:33.689435Z","shell.execute_reply.started":"2025-12-02T18:51:32.223226Z","shell.execute_reply":"2025-12-02T18:51:33.688258Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pandas as pd\nSAMPLE_SUBMISSION_PATH = '/kaggle/input/dr-trainlabels/trainLabels.csv'\ntry:\n    df = pd.read_csv(SAMPLE_SUBMISSION_PATH, header=None)\n    print(f\"Original columns (indices): {df.columns.tolist()}\")\n    new_column_names = {\n        0: 'id_code',\n        1: 'diagnosis'\n    }\n    \n    df.rename(columns=new_column_names, inplace=True)\n    output_path = 'submission_renamed.csv'\n    df.to_csv(output_path, index=False)\n    print(f\"Renamed columns: {df.columns.tolist()}\")\n    print(f\"File successfully saved to: {output_path}\")\n\nexcept FileNotFoundError:\n    print(f\"Error: File not found at {SAMPLE_SUBMISSION_PATH}. Please verify the path.\")\n\nexcept Exception as e:\n    print(f\"An error occurred: {e}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-02T19:59:08.313661Z","iopub.execute_input":"2025-12-02T19:59:08.314444Z","iopub.status.idle":"2025-12-02T19:59:08.391512Z","shell.execute_reply.started":"2025-12-02T19:59:08.314419Z","shell.execute_reply":"2025-12-02T19:59:08.390855Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.utils.data import Dataset, DataLoader\nimport pandas as pd\nimport cv2\nimport numpy as np\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\nfrom sklearn.model_selection import StratifiedKFold\nfrom sklearn.metrics import cohen_kappa_score\nfrom sklearn.utils.class_weight import compute_class_weight\nimport timm\nimport os\nimport joblib\nimport warnings\n\nwarnings.filterwarnings(\"ignore\")\n\n# --- Configuration ---\nTRAIN_CSV = '/kaggle/input/diabetic-retinopathy-classification-3/kaggle_retina/train.csv'\nTRAIN_DIR = '/kaggle/input/diabetic-retinopathy-classification-3/kaggle_retina/train'\n\nTEST_DIR = '/kaggle/input/diabetic-retinopathy-train-unzipped/train'\nSAMPLE_SUBMISSION_PATH = '/kaggle/working/submission_renamed.csv'\nOUTPUT_DIR = \"output_results_v3_vit_focal_threehead\" # Changed output directory to reflect changes\n\n# Model and training parameters\nIMG_SIZE = 384 # ViT standard input size\nBATCH_SIZE = 16\nN_EPOCHS_M1 = 15 # Binary classification is easier, may require fewer epochs\nN_EPOCHS_M2 = 30\nN_FOLDS = 5\n# Switched to a Vision Transformer backbone\nENCODER_MODEL_NAME = 'vit_base_patch16_384'\n# Adjusted Learning Rates for ViT and new pipeline\nLR_M1_HEAD = 3e-4\nLR_M1_BACKBONE = 3e-5\nLR_M2_HEAD = 3e-4 # Using this for all three heads in Model 2\nLR_M2_BACKBONE = 3e-5\nREG_LOSS_WEIGHT = 0.1 # This weight will now only apply to the regression head loss in Model 1\nORD_LOSS_WEIGHT = 0.5 # New weight for ordinal loss in Model 2\nREG_LOSS_WEIGHT_M2 = 0.5 # New weight for regression loss in Model 2\nDEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\n# Create output directories\nos.makedirs(OUTPUT_DIR, exist_ok=True)\n# MODEL_SAVE_DIR = os.path.join(OUTPUT_DIR, \"saved_models\")\n# os.makedirs(MODEL_SAVE_DIR, exist_ok=True)\n\nMODEL_SAVE_DIR='/kaggle/input/vit/pytorch/default/1/output_results_v3_vit_focal_threehead/saved_models'\n\n# --- Preprocessing Functions ---\ndef apply_clahe(image):\n    lab = cv2.cvtColor(image, cv2.COLOR_RGB2LAB)\n    l, a, b = cv2.split(lab)\n    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))\n    cl = clahe.apply(l)\n    limg = cv2.merge((cl, a, b))\n    return cv2.cvtColor(limg, cv2.COLOR_LAB2RGB)\n\ndef apply_unsharp_mask(image, sigma=1.0, strength=1.5):\n    blurred = cv2.GaussianBlur(image, (0, 0), sigma)\n    return cv2.addWeighted(image, 1.0 + strength, blurred, -strength, 0)\n\n# --- Dataset Class ---\nclass DRDataset(Dataset):\n    def __init__(self, df_data, img_dir, transforms=None, is_test=False, clf_target_col=None, reg_target_col=None, ordinal_mode=False):\n        self.df = df_data.reset_index(drop=True)\n        self.img_dir = img_dir\n        self.transforms = transforms\n        self.is_test = is_test\n        self.clf_target_col = clf_target_col\n        self.reg_target_col = reg_target_col\n        self.ordinal_mode = ordinal_mode\n\n    def __len__(self):\n        return len(self.df)\n\n    def __getitem__(self, idx):\n        img_id = self.df.loc[idx, 'id_code']\n        img_name = f\"{img_id}.jpeg\"\n        img_path = os.path.join(self.img_dir, img_name)\n        image = cv2.imread(img_path)\n\n        if image is None:\n            print(f\"FATAL: Image not found at {img_path}. Check data paths.\")\n            image = np.zeros((IMG_SIZE, IMG_SIZE, 3), dtype=np.uint8)\n        else:\n            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n\n        image = apply_clahe(image)\n        image = apply_unsharp_mask(image)\n\n        if self.transforms:\n            image = self.transforms(image=image)['image']\n\n        if self.is_test:\n            return image, img_id\n        else:\n            label_clf = self.df.loc[idx, self.clf_target_col]\n            label_reg = self.df.loc[idx, self.reg_target_col]\n\n            # For Model 2 (ordinal_mode=True), return all three targets\n            if self.ordinal_mode:\n                return image, (torch.tensor(label_clf, dtype=torch.long),\n                               torch.tensor(label_reg, dtype=torch.float).unsqueeze(0),\n                               torch.tensor(label_clf, dtype=torch.long)) # Use clf label also for ordinal target generation\n            else: # For Model 1 (binary classification)\n                return image, (torch.tensor(label_clf, dtype=torch.long),\n                               torch.tensor(label_reg, dtype=torch.float).unsqueeze(0))\n\n# --- Augmentations ---\ndef get_train_transforms(img_size):\n    return A.Compose([\n        A.Resize(img_size, img_size),\n        A.HorizontalFlip(p=0.5),\n        A.VerticalFlip(p=0.5),\n        A.RandomRotate90(p=0.5),\n        A.RandomBrightnessContrast(p=0.5),\n        A.GaussNoise(p=0.2),\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(img_size):\n    return A.Compose([\n        A.Resize(img_size, img_size),\n        A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),\n        ToTensorV2()\n    ])\n\n# --- Model Architectures ---\nclass MultiTaskModel(nn.Module): # For Model 1 (Binary Classification)\n    def __init__(self, encoder_name, num_classes, pretrained=True):\n        super().__init__()\n        self.encoder = timm.create_model(encoder_name, pretrained=pretrained)\n        n_features = self.encoder.head.in_features\n        self.encoder.head = nn.Identity() # Remove the original head\n\n        self.classification_head = nn.Linear(n_features, num_classes)\n        self.regression_head = nn.Linear(n_features, 1)\n\n    def forward(self, x):\n        features = self.encoder(x)\n        clf_output = self.classification_head(features)\n        reg_output = self.regression_head(features)\n        return clf_output, reg_output\n\nclass ThreeHeadModel(nn.Module): # For Model 2 (4-Class Severity Grading)\n    def __init__(self, encoder_name, num_classes, pretrained=True):\n        super().__init__()\n        self.encoder = timm.create_model(encoder_name, pretrained=pretrained, num_classes=0) # num_classes=0 for feature extraction\n\n        # Determine encoder_output_shape dynamically\n        # Create a dummy input with the correct image size\n        dummy_input_size = (1, 3, IMG_SIZE, IMG_SIZE)\n        dummy_input = torch.randn(dummy_input_size)\n\n        with torch.no_grad(): # Ensure no gradients are computed during this dummy pass\n            encoder_output = self.encoder(dummy_input)\n        encoder_output_shape = encoder_output.shape[-1] # Assumes features are in the last dimension\n\n        intermediate_features = 512\n        self.shared_dense = nn.Sequential(\n            nn.Linear(encoder_output_shape, intermediate_features),\n            nn.BatchNorm1d(intermediate_features), # As per Fig 5 showing BatchNorm1D\n            nn.ReLU()\n        )\n        self.classification_head = nn.Linear(intermediate_features, num_classes)\n        self.regression_head = nn.Linear(intermediate_features, 1)\n        self.ordinal_head = nn.Linear(intermediate_features, num_classes - 1) # N_CLASSES - 1 for ordinal output\n\n    def forward(self, x):\n        features = self.encoder(x)\n        shared_out = self.shared_dense(features)\n        clf_out = self.classification_head(shared_out)\n        reg_out = self.regression_head(shared_out)\n        ord_out = self.ordinal_head(shared_out)\n        return clf_out, reg_out, ord_out\n\n# --- Loss Functions ---\n# Focal Loss Implementation\nclass FocalLoss(nn.Module):\n    def __init__(self, alpha=None, gamma=2.0, reduction='mean'):\n        super(FocalLoss, self).__init__()\n        self.alpha = alpha\n        self.gamma = gamma\n        self.reduction = reduction\n\n    def forward(self, inputs, targets):\n        p = torch.softmax(inputs, dim=1)\n        p_t = p.gather(1, targets.unsqueeze(1)).squeeze(1)\n        log_p_t = torch.log(p_t.clamp(min=1e-8)) # Clamp to avoid log(0)\n        focal_term = (1 - p_t) ** self.gamma\n\n        if self.alpha is not None:\n            alpha_t = self.alpha.gather(0, targets)\n            loss = -alpha_t * focal_term * log_p_t\n        else:\n            loss = -focal_term * log_p_t\n\n        if self.reduction == 'mean':\n            return loss.mean()\n        elif self.reduction == 'sum':\n            return loss.sum()\n        else: # 'none'\n            return loss\n\n# Helper for ordinal targets\ndef get_ordinal_targets(labels, num_classes):\n    ordinal_targets = torch.zeros(labels.size(0), num_classes - 1, device=labels.device)\n    for i, label in enumerate(labels):\n        if label > 0: # Labels are 0,1,2,3 for 4 classes, so if label is > 0, it means it's a DR grade\n            ordinal_targets[i, :label.item()] = 1.0 # If label is 1, targets[i,0]=1. If label is 2, targets[i,0:2]=1 etc.\n    return ordinal_targets\n\n# --- Training and Validation ---\n# Modified train_epoch_mtl to handle either 2 outputs (Model 1) or 3 outputs (Model 2)\ndef train_epoch_mtl(loader, model, optimizer, crit_clf, crit_reg, crit_ord, reg_weight, ord_weight, device, model_type=\"M1\"):\n    model.train()\n    total_loss = 0\n    for data, targets in loader:\n        data = data.to(device)\n        optimizer.zero_grad()\n\n        if model_type == \"M1\": # For Binary Classification\n            targets_clf, targets_reg = targets[0].to(device), targets[1].to(device)\n            clf_out, reg_out = model(data)\n            loss_clf = crit_clf(clf_out, targets_clf)\n            loss_reg = crit_reg(reg_out, targets_reg)\n            loss = loss_clf + reg_weight * loss_reg\n        elif model_type == \"M2\": # For 4-Class Severity Grading with ThreeHeadModel\n            targets_clf, targets_reg, targets_ord_raw = targets[0].to(device), targets[1].to(device), targets[2].to(device)\n            clf_out, reg_out, ord_out = model(data)\n\n            loss_clf = crit_clf(clf_out, targets_clf) # Focal Loss\n            loss_reg = crit_reg(reg_out, targets_reg) # MSE Loss\n\n            # Ordinal loss\n            # For get_ordinal_targets, we pass the raw classification target (0,1,2,3)\n            # and N_CLASSES for Model 2 (which is 4)\n            ord_targets = get_ordinal_targets(targets_ord_raw, 4) # Hardcoding 4 for M2 classes\n            loss_ord = crit_ord(ord_out, ord_targets) # BCEWithLogitsLoss\n\n            loss = loss_clf + reg_weight * loss_reg + ord_weight * loss_ord\n        else:\n            raise ValueError(\"Invalid model_type. Must be 'M1' or 'M2'.\")\n\n        loss.backward()\n        optimizer.step()\n        total_loss += loss.item()\n    return total_loss / len(loader)\n\n# Modified validate_epoch_mtl to handle either 2 outputs (Model 1) or 3 outputs (Model 2)\n# And to always use the classification output for Kappa calculation.\ndef validate_epoch_mtl(loader, model, device, num_classes, model_type=\"M1\"):\n    model.eval()\n    all_preds, all_targets = [], []\n    with torch.no_grad():\n        for data, targets in loader:\n            data = data.to(device)\n            if model_type == \"M1\":\n                targets_clf = targets[0].to(device)\n                clf_out, _ = model(data)\n            elif model_type == \"M2\":\n                targets_clf = targets[0].to(device)\n                clf_out, _, _ = model(data)\n            else:\n                raise ValueError(\"Invalid model_type. Must be 'M1' or 'M2'.\")\n\n            predictions = torch.argmax(clf_out, dim=1)\n            all_preds.extend(predictions.cpu().numpy())\n            all_targets.extend(targets_clf.cpu().numpy()) # Use .cpu().numpy() for metrics\n\n    # Corrected: Changed 'simple' to None for binary kappa calculation, and 'quadratic' for multi-class\n    return cohen_kappa_score(all_targets, all_preds, weights='quadratic' if num_classes > 2 else None)\n\n# --- Prediction with TTA ---\ndef predict_with_tta(model, images, device, model_type=\"M1\"):\n    model.eval()\n    with torch.no_grad():\n        if model_type == \"M1\":\n            clf_out_orig, _ = model(images)\n        elif model_type == \"M2\":\n            clf_out_orig, _, _ = model(images) # Only need clf_out\n        else:\n            raise ValueError(\"Invalid model_type. Must be 'M1' or 'M2'.\")\n\n        probs_orig = torch.softmax(clf_out_orig, dim=1)\n\n        images_hflip = torch.flip(images, dims=[3])\n        if model_type == \"M1\":\n            clf_out_hflip, _ = model(images_hflip)\n        elif model_type == \"M2\":\n            clf_out_hflip, _, _ = model(images_hflip) # Only need clf_out\n\n        probs_hflip = torch.softmax(clf_out_hflip, dim=1)\n        return (probs_orig + probs_hflip) / 2.0\n\n# --- Main Script ---\ndef run_training():\n    print(f\"Using device: {DEVICE}\")\n    try:\n        main_df = pd.read_csv(TRAIN_CSV)\n        print(f\"Successfully loaded {TRAIN_CSV} with {len(main_df)} records.\")\n    except FileNotFoundError:\n        print(f\"FATAL ERROR: Training CSV not found at {TRAIN_CSV}.\")\n        return\n\n    # # --- MODEL 1: Binary DR vs No DR ---\n    # print(\"\\n\" + \"=\"*50 + \"\\nStage 1: Training Model 1 (DR vs. No DR)\\n\" + \"=\"*50)\n    # main_df['model1_label'] = (main_df['diagnosis'] > 0).astype(int) # 0 for No DR, 1 for DR\n    # skf = StratifiedKFold(n_splits=N_FOLDS, shuffle=True, random_state=42)\n\n    # for fold, (train_idx, val_idx) in enumerate(skf.split(main_df['id_code'], main_df['model1_label'])):\n    #     print(f\"\\n--- Fold {fold+1}/{N_FOLDS} for Model 1 ---\")\n    #     model1 = MultiTaskModel(ENCODER_MODEL_NAME, num_classes=2).to(DEVICE)\n\n    #     train_loader = DataLoader(DRDataset(main_df.iloc[train_idx], TRAIN_DIR, get_train_transforms(IMG_SIZE), clf_target_col='model1_label', reg_target_col='model1_label'), batch_size=BATCH_SIZE, shuffle=True, num_workers=2)\n    #     val_loader = DataLoader(DRDataset(main_df.iloc[val_idx], TRAIN_DIR, get_valid_transforms(IMG_SIZE), clf_target_col='model1_label', reg_target_col='model1_label'), batch_size=BATCH_SIZE, shuffle=False, num_workers=2)\n\n    #     # Optimizer for Model 1\n    #     optimizer = optim.AdamW([\n    #         {'params': model1.encoder.parameters(), 'lr': LR_M1_BACKBONE},\n    #         {'params': model1.classification_head.parameters(), 'lr': LR_M1_HEAD},\n        #     {'params': model1.regression_head.parameters(), 'lr': LR_M1_HEAD}\n        # ])\n        # scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=N_EPOCHS_M1)\n        # crit_clf_m1, crit_reg_m1 = nn.CrossEntropyLoss(), nn.MSELoss() # Original losses for Model 1\n\n        # best_kappa = -1.0\n        # for epoch in range(N_EPOCHS_M1):\n        #     train_loss = train_epoch_mtl(train_loader, model1, optimizer, crit_clf_m1, crit_reg_m1, None, REG_LOSS_WEIGHT, 0, DEVICE, model_type=\"M1\")\n        #     val_kappa = validate_epoch_mtl(val_loader, model1, DEVICE, 2, model_type=\"M1\")\n        #     scheduler.step()\n        #     print(f\"Epoch {epoch+1}/{N_EPOCHS_M1} | Train Loss: {train_loss:.4f} | Val Kappa: {val_kappa:.4f}\")\n        #     if val_kappa > best_kappa:\n        #         best_kappa = val_kappa\n        #         torch.save(model1.state_dict(), os.path.join(MODEL_SAVE_DIR, f\"best_model_1_fold_{fold}.pth\"))\n        #         print(f\"Saved Model 1 for fold {fold+1} with Kappa: {best_kappa:.4f}\")\n        # print(f\"Best Kappa for Model 1, Fold {fold+1}: {best_kappa:.4f}\")\n\n    # # --- MODEL 2: 4-Class Severity Grading ---\n    # print(\"\\n\" + \"=\"*50 + \"\\nStage 2: Training Model 2 (Severity Grades 1-4)\\n\" + \"=\"*50)\n    # dr_df = main_df[main_df['diagnosis'] > 0].copy()\n    # dr_df['model2_label'] = dr_df['diagnosis'] - 1 # Map 1,2,3,4 to 0,1,2,3 for 4 classes\n\n    # # Calculate class weights for Focal Loss for Model 2\n    # class_counts_m2 = dr_df['model2_label'].value_counts().sort_index()\n    # total_samples_m2 = len(dr_df)\n\n    # # Ensure all 4 classes (0, 1, 2, 3) are represented in class_counts_m2 for weight calculation\n    # # Fill missing classes with 0 count to avoid KeyError during reindexing\n    # all_classes_m2 = np.arange(4) # Classes 0, 1, 2, 3\n    # class_counts_m2 = class_counts_m2.reindex(all_classes_m2, fill_value=0)\n\n    # # Handle division by zero if a class has 0 samples (though unlikely with StratifiedKFold)\n    # # Add a small epsilon to avoid NaN in weights if a class count is 0\n    # weights_raw_m2 = total_samples_m2 / (4 * class_counts_m2.values + 1e-8)\n    # class_weights_focal_m2_tensor = torch.tensor(weights_raw_m2, dtype=torch.float).to(DEVICE)\n    # print(f\"Class weights for Model 2 (Focal Loss): {class_weights_focal_m2_tensor}\")\n\n    # crit_clf_m2 = FocalLoss(alpha=class_weights_focal_m2_tensor, gamma=2.0) # Focal Loss for Model 2\n    # crit_reg_m2 = nn.MSELoss() # MSE for regression head in Model 2\n    # crit_ord_m2 = nn.BCEWithLogitsLoss() # BCEWithLogitsLoss for ordinal head in Model 2\n\n    # skf_model2 = StratifiedKFold(n_splits=N_FOLDS, shuffle=True, random_state=42)\n    # for fold, (train_idx, val_idx) in enumerate(skf_model2.split(dr_df['id_code'], dr_df['model2_label'])):\n    #     print(f\"\\n--- Fold {fold+1}/{N_FOLDS} for Model 2 ---\")\n\n    #     # Use ThreeHeadModel for Model 2\n    #     model2 = ThreeHeadModel(ENCODER_MODEL_NAME, num_classes=4).to(DEVICE) # num_classes is 4 (0-3)\n\n    #     train_loader = DataLoader(DRDataset(dr_df.iloc[train_idx], TRAIN_DIR, get_train_transforms(IMG_SIZE), clf_target_col='model2_label', reg_target_col='model2_label', ordinal_mode=True), batch_size=max(1, BATCH_SIZE//2), shuffle=True, num_workers=2)\n    #     val_loader = DataLoader(DRDataset(dr_df.iloc[val_idx], TRAIN_DIR, get_valid_transforms(IMG_SIZE), clf_target_col='model2_label', reg_target_col='model2_label', ordinal_mode=True), batch_size=max(1, BATCH_SIZE//2), shuffle=False, num_workers=2)\n\n    #     # Optimizer for Model 2, optimizing all three heads\n    #     optimizer_m2 = optim.AdamW([\n    #         {'params': model2.encoder.parameters(), 'lr': LR_M2_BACKBONE},\n    #         {'params': model2.shared_dense.parameters(), 'lr': LR_M2_HEAD},\n    #         {'params': model2.classification_head.parameters(), 'lr': LR_M2_HEAD},\n    #         {'params': model2.regression_head.parameters(), 'lr': LR_M2_HEAD},\n    #         {'params': model2.ordinal_head.parameters(), 'lr': LR_M2_HEAD}\n    #     ])\n    #     scheduler_m2 = optim.lr_scheduler.CosineAnnealingLR(optimizer_m2, T_max=N_EPOCHS_M2)\n\n        # best_kappa = -1.0\n        # for epoch in range(N_EPOCHS_M2):\n        #     train_loss = train_epoch_mtl(train_loader, model2, optimizer_m2, crit_clf_m2, crit_reg_m2, crit_ord_m2, REG_LOSS_WEIGHT_M2, ORD_LOSS_WEIGHT, DEVICE, model_type=\"M2\")\n        #     val_kappa = validate_epoch_mtl(val_loader, model2, DEVICE, 4, model_type=\"M2\") # Pass num_classes=4 for quadratic kappa\n        #     scheduler_m2.step()\n        #     print(f\"Epoch {epoch+1}/{N_EPOCHS_M2} | Train Loss: {train_loss:.4f} | Val Kappa: {val_kappa:.4f}\")\n        #     if val_kappa > best_kappa:\n        #         best_kappa = val_kappa\n        #         torch.save(model2.state_dict(), os.path.join(MODEL_SAVE_DIR, f\"best_model_2_fold_{fold}.pth\"))\n        #         print(f\"Saved Model 2 for fold {fold+1} with Kappa: {best_kappa:.4f}\")\n        # print(f\"Best Kappa for Model 2, Fold {fold+1}: {best_kappa:.4f}\")\n        \n# --- Prediction on Test Set ---\n    print(\"\\n\" + \"=\"*50 + \"\\nStage 3: Generating Test Predictions\\n\" + \"=\"*50)\n    try:\n        sub_df = pd.read_csv(SAMPLE_SUBMISSION_PATH)\n    except FileNotFoundError:\n        print(f\"Warning: Sample submission file not found at {SAMPLE_SUBMISSION_PATH}.\")\n        sub_df = pd.DataFrame()\n    if not sub_df.empty:\n        test_dataset = DRDataset(sub_df, TEST_DIR, get_valid_transforms(IMG_SIZE), is_test=True)\n        test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False)\n\n        models1, models2 = [], []\n        for fold in range(N_FOLDS):\n            # Load Model 1 (MultiTaskModel)\n            m1 = MultiTaskModel(ENCODER_MODEL_NAME, 2).to(DEVICE)\n            m1_path = os.path.join(MODEL_SAVE_DIR, f\"best_model_1_fold_{fold}.pth\")\n            if os.path.exists(m1_path):\n                m1.load_state_dict(torch.load(m1_path)); models1.append(m1.eval())\n            else:\n                print(f\"Warning: Model 1 for fold {fold} not found at {m1_path}. Skipping.\")\n\n            # Load Model 2 (ThreeHeadModel)\n            m2 = ThreeHeadModel(ENCODER_MODEL_NAME, 4).to(DEVICE) # num_classes=4 for Model 2\n            m2_path = os.path.join(MODEL_SAVE_DIR, f\"best_model_2_fold_{fold}.pth\")\n            if os.path.exists(m2_path):\n                m2.load_state_dict(torch.load(m2_path)); models2.append(m2.eval())\n            else:\n                print(f\"Warning: Model 2 for fold {fold} not found at {m2_path}. Skipping.\")\n\n\n        if not models1 or not models2:\n            print(\"ERROR: Not enough models loaded for prediction. Ensure training completed successfully.\")\n            return\n\n        final_predictions = []\n        with torch.no_grad():\n            for images, img_ids in test_loader:\n                images = images.to(DEVICE)\n\n                # Stage 1: Get soft-voted binary predictions from Model 1\n                all_fold_probs1 = torch.stack([predict_with_tta(m, images, DEVICE, model_type=\"M1\") for m in models1], dim=1)\n                avg_probs1 = torch.mean(all_fold_probs1, dim=1)\n                preds1 = torch.argmax(avg_probs1, dim=1)\n\n                final_batch_preds = torch.zeros_like(preds1) # Final predictions will be 0-4\n                # Identify images with DR present (predicted as 1 by Model 1)\n                dr_indices = (preds1 == 1).nonzero(as_tuple=True)[0]\n\n                if len(dr_indices) > 0:\n                    images_for_m2 = images[dr_indices]\n                    # Stage 2: Get soft-voted 4-class predictions for DR images from Model 2\n                    all_fold_probs2 = torch.stack([predict_with_tta(m, images_for_m2, DEVICE, model_type=\"M2\") for m in models2], dim=1)\n                    avg_probs2 = torch.mean(all_fold_probs2, dim=1)\n                    preds2 = torch.argmax(avg_probs2, dim=1)\n\n                    # Map predictions back to original labels (0,1,2,3 -> 1,2,3,4)\n                    for i, original_idx in enumerate(dr_indices):\n                        final_batch_preds[original_idx] = preds2[i] + 1 # Add 1 to get original diagnosis labels (1-4)\n\n                final_predictions.extend(final_batch_preds.cpu().numpy())\n\n        sub_df['diagnosis'] = final_predictions\n        sub_df.to_csv(os.path.join(OUTPUT_DIR, 'submission.csv'), index=False)\n        print(f\"Submission file saved to: {os.path.join(OUTPUT_DIR, 'submission.csv')}\")\n\nif __name__ == '__main__':\n    run_training()\n   ","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-02T20:16:03.155671Z","iopub.execute_input":"2025-12-02T20:16:03.156225Z"}},"outputs":[],"execution_count":null}]}