{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceType":"competition","sourceId":14774,"databundleVersionId":875431},{"sourceType":"competition","sourceId":4104,"databundleVersionId":46661}],"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\nimport cv2\nimport time\nimport torch\nimport shutil\nimport pandas as pd\nimport numpy as np\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport matplotlib.pyplot as plt\nimport seaborn as sns\nfrom torch.utils.data import Dataset, DataLoader, Subset\nfrom torchvision import transforms\nfrom sklearn.metrics import cohen_kappa_score, confusion_matrix, accuracy_score, precision_score, recall_score, f1_score\nfrom sklearn.model_selection import train_test_split\nfrom tqdm import tqdm\n\n# --- Project Configuration ---\nCONFIG = {\n    \"image_resolution\": 1024,\n    \"batch_size\": 1, \n    \"epochs\": 50,\n    \"lr\": 1e-4,\n    \"embed_dim\": 256,\n    \"n_layers\": 12,\n    \"ckpt_dir\": \"/kaggle/working/checkpoints\",\n    \"aptos_path\": \"/kaggle/input/competitions/aptos2019-blindness-detection\",\n    \"dr_2015_path\": \"/kaggle/input/competitions/diabetic-retinopathy-detection\"\n}\n\ndef init_environment():\n    if os.path.exists('/kaggle/working'):\n        for f in os.listdir('/kaggle/working'):\n            p = os.path.join('/kaggle/working', f)\n            try:\n                if os.path.isfile(p) or os.path.is_link(p): os.unlink(p)\n                elif os.path.isdir(p): shutil.rmtree(p)\n            except: pass\n    os.makedirs(CONFIG[\"ckpt_dir\"], exist_ok=True)\n    print(f\"[LOG] GPU Support: {torch.cuda.get_device_name(0)}\")\n\nclass SelectiveScanCell(nn.Module):\n    def __init__(self, dim, d_state=16, expand=2):\n        super().__init__()\n        self.d_inner = dim * expand\n        self.in_proj = nn.Linear(dim, self.d_inner * 2)\n        self.conv1d = nn.Conv1d(self.d_inner, self.d_inner, kernel_size=4, groups=self.d_inner, padding=3)\n        self.x_proj = nn.Linear(self.d_inner, d_state * 2 + self.d_inner)\n        self.dt_proj = nn.Linear(self.d_inner, self.d_inner)\n        self.A_log = nn.Parameter(torch.log(torch.arange(1, d_state + 1).float().repeat(self.d_inner, 1)))\n        self.D = nn.Parameter(torch.ones(self.d_inner))\n        self.out_proj = nn.Linear(self.d_inner, dim)\n\n    def forward(self, x):\n        b, l, d = x.shape\n        z = self.in_proj(x)\n        x, res = z.chunk(2, dim=-1)\n        x = x.transpose(1, 2)\n        x = self.conv1d(x)[:, :, :l]\n        x = F.silu(x).transpose(1, 2)\n        ssm_data = self.x_proj(x)\n        dt, B, C = torch.split(ssm_data, [self.d_inner, 16, 16], dim=-1)\n        dt = F.softplus(self.dt_proj(x))\n        A = -torch.exp(self.A_log)\n        y = x * torch.sigmoid(dt * A.mean(-1).view(1, 1, -1)) + res * torch.sigmoid(x)\n        return self.out_proj(y)\n\nclass ParallelVimLayer(nn.Module):\n    def __init__(self, dim):\n        super().__init__()\n        self.n_branches = 4\n        self.branch_dim = dim // self.n_branches\n        self.branches = nn.ModuleList([SelectiveScanCell(self.branch_dim) for _ in range(self.n_branches)])\n        self.norm = nn.LayerNorm(dim)\n        self.fusion = nn.Linear(dim, dim)\n\n    def forward(self, x):\n        h = self.norm(x)\n        chunks = torch.chunk(h, self.n_branches, dim=-1)\n        outs = [branch(chunk) for branch, chunk in zip(self.branches, chunks)]\n        return x + self.fusion(torch.cat(outs, dim=-1))\n\nclass VesselAwareBlock(nn.Module):\n    def __init__(self, dim):\n        super().__init__()\n        self.vessel_prior = nn.Sequential(\n            nn.Conv2d(dim, dim, kernel_size=3, padding=1, groups=dim),\n            nn.Sigmoid()\n        )\n        self.mamba = SelectiveScanCell(dim)\n        self.norm = nn.LayerNorm(dim)\n\n    def forward(self, x, spatial_shape):\n        b, l, d = x.shape\n        h, w = spatial_shape\n        feat_map = x.transpose(1, 2).view(b, d, h, w)\n        prior = self.vessel_prior(feat_map).view(b, d, l).transpose(1, 2)\n        out = self.mamba(self.norm(x))\n        return x + (out * prior)\n\nclass VimDRNetwork(nn.Module):\n    def __init__(self, embed_dim=256, depth=12):\n        super().__init__()\n        self.patch_embed = nn.Conv2d(3, embed_dim, kernel_size=16, stride=16)\n        self.layers = nn.ModuleList()\n        for i in range(depth):\n            if i % 2 == 0:\n                self.layers.append(ParallelVimLayer(embed_dim))\n            else:\n                self.layers.append(VesselAwareBlock(embed_dim))\n        self.reg_head = nn.Linear(embed_dim, 1)\n        self.pool = nn.AdaptiveAvgPool1d(1)\n\n    def forward(self, x):\n        b, c, h, w = x.shape\n        x = self.patch_embed(x)\n        spatial_shape = (x.shape[2], x.shape[3])\n        x = x.flatten(2).transpose(1, 2)\n        for layer in self.layers:\n            if isinstance(layer, VesselAwareBlock): x = layer(x, spatial_shape)\n            else: x = layer(x)\n        feat = self.pool(x.transpose(1, 2)).flatten(1)\n        return self.reg_head(feat)\n\nclass UnifiedDRDataset(Dataset):\n    def __init__(self, dataframe, train=True):\n        self.df = dataframe\n        self.tf = transforms.Compose([\n            transforms.ToPILImage(),\n            transforms.RandomHorizontalFlip() if train else transforms.Lambda(lambda x: x),\n            transforms.ToTensor(),\n            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n        ])\n\n    def __len__(self): return len(self.df)\n\n    def _ben_graham_filter(self, path):\n        img = cv2.imread(path)\n        if img is None: return np.zeros((1024, 1024, 3), np.uint8)\n        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n        img = cv2.resize(img, (CONFIG[\"image_resolution\"], CONFIG[\"image_resolution\"]))\n        return cv2.addWeighted(img, 4, cv2.GaussianBlur(img, (0,0), 10), -4, 128)\n\n    def __getitem__(self, i):\n        row = self.df.iloc[i]\n        image = self._ben_graham_filter(row['path'])\n        return self.tf(image), torch.tensor(row['diagnosis'], dtype=torch.float32)\n\ndef generate_report_visuals(history, y_true, y_pred, params, energy):\n    epochs = range(1, len(history['train_loss']) + 1)\n    plt.figure(figsize=(18, 5))\n    \n    plt.subplot(1, 3, 1)\n    plt.plot(epochs, history['train_loss'], color='blue')\n    plt.title('Loss Convergence')\n    plt.subplot(1, 3, 2)\n    plt.plot(epochs, history['val_acc'], color='orange')\n    plt.title('Accuracy Progress')\n    plt.subplot(1, 3, 3)\n    plt.plot(epochs, history['val_kappa'], color='green')\n    plt.title('Kappa Optimization')\n    plt.savefig('/kaggle/working/performance_curves.png')\n    plt.show()\n\n    cm = confusion_matrix(y_true, y_pred)\n    plt.figure(figsize=(8, 6))\n    sns.heatmap(cm, annot=True, fmt='d', cmap='RdPu')\n    plt.savefig('/kaggle/working/confusion_matrix.png')\n    plt.show()\n\n    print(\"\\n\" + \"=\"*40)\n    print(\"FINAL RESEARCH EVALUATION\")\n    print(\"=\"*40)\n    print(f\"Model Parameters: {params:.2f} M\")\n    print(f\"Energy Footprint: {energy:.2f} Wh\")\n    print(f\"Quadratic Kappa: {cohen_kappa_score(y_true, y_pred, weights='quadratic'):.4f}\")\n    print(f\"Precision Score: {precision_score(y_true, y_pred, average='weighted'):.4f}\")\n    print(f\"Recall (Sensitivity): {recall_score(y_true, y_pred, average='weighted'):.4f}\")\n    print(f\"F1-Score: {f1_score(y_true, y_pred, average='weighted'):.4f}\")\n\ndef run_experiment():\n    init_environment()\n    device = torch.device(\"cuda\")\n    \n    # Dataset 1: APTOS 2019\n    aptos_csv = f'{CONFIG[\"aptos_path\"]}/train.csv'\n    if not os.path.exists(aptos_csv):\n        print(f\"[ERROR] APTOS CSV not found at {aptos_csv}. Please verify dataset addition.\")\n        return\n        \n    aptos_df = pd.read_csv(aptos_csv)\n    aptos_df['path'] = aptos_df['id_code'].apply(lambda x: f'{CONFIG[\"aptos_path\"]}/train_images/{x}.png')\n    \n    # Dataset 2: 2015 DR Detection\n    dr2015_csv = f'{CONFIG[\"dr_2015_path\"]}/trainLabels.csv'\n    if not os.path.exists(dr2015_csv):\n        print(f\"[ERROR] DR 2015 CSV not found at {dr2015_csv}. Skipping Dataset 2.\")\n        full_df = aptos_df[['path', 'diagnosis']]\n    else:\n        dr2015_df = pd.read_csv(dr2015_csv)\n        dr2015_df['path'] = dr2015_df['image'].apply(lambda x: f'{CONFIG[\"dr_2015_path\"]}/train/{x}.jpeg')\n        dr2015_df.rename(columns={'level': 'diagnosis'}, inplace=True)\n        full_df = pd.concat([aptos_df[['path', 'diagnosis']], dr2015_df[['path', 'diagnosis']]], ignore_index=True)\n    \n    tr_idx, val_idx = train_test_split(full_df.index, test_size=0.15, random_state=42, stratify=full_df['diagnosis'])\n    \n    train_loader = DataLoader(Subset(UnifiedDRDataset(full_df, True), tr_idx), batch_size=CONFIG[\"batch_size\"], shuffle=True)\n    val_loader = DataLoader(Subset(UnifiedDRDataset(full_df, False), val_idx), batch_size=CONFIG[\"batch_size\"])\n    \n    model = VimDRNetwork(embed_dim=CONFIG[\"embed_dim\"], depth=CONFIG[\"n_layers\"]).to(device)\n    params = sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6\n    \n    optimizer = torch.optim.AdamW(model.parameters(), lr=CONFIG[\"lr\"])\n    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=CONFIG[\"epochs\"])\n    criterion = nn.MSELoss()\n    \n    history = {'train_loss': [], 'val_acc': [], 'val_kappa': []}\n    start_time = time.time()\n    \n    for epoch in range(1, CONFIG[\"epochs\"] + 1):\n        model.train()\n        train_mse = 0\n        pbar = tqdm(train_loader, desc=f\"Cycle {epoch}/{CONFIG['epochs']}\")\n        for img, lbl in pbar:\n            img, lbl = img.to(device), lbl.to(device).view(-1, 1)\n            optimizer.zero_grad()\n            out = model(img)\n            loss = criterion(out, lbl)\n            loss.backward()\n            optimizer.step()\n            train_mse += loss.item()\n            pbar.set_postfix({'MSE': f\"{loss.item():.4f}\"})\n        \n        scheduler.step()\n        \n        model.eval()\n        p, t = [], []\n        with torch.no_grad():\n            for img, lbl in val_loader:\n                img, lbl = img.to(device), lbl.to(device).view(-1, 1)\n                preds = model(img)\n                p.extend(torch.clamp(torch.round(preds), 0, 4).cpu().numpy())\n                t.extend(lbl.cpu().numpy())\n        \n        y_true, y_pred = np.array(t).astype(int), np.array(p).astype(int)\n        acc = accuracy_score(y_true, y_pred)\n        kappa = cohen_kappa_score(y_true, y_pred, weights='quadratic')\n        \n        history['train_loss'].append(train_mse/len(train_loader))\n        history['val_acc'].append(acc)\n        history['val_kappa'].append(kappa)\n        \n        print(f\"accuracy: {acc:.4f} - loss: {train_mse/len(train_loader):.4f} - val_accuracy: {acc:.4f} - learning_rate: {optimizer.param_groups[0]['lr']} - Epoch: {epoch}\")\n        torch.save(model.state_dict(), f\"{CONFIG['ckpt_dir']}/vim_dr_epoch_{epoch}.pth\")\n\n    energy = ((time.time() - start_time) / 3600) * 75 \n    generate_report_visuals(history, y_true, y_pred, params, energy)\n\nif __name__ == \"__main__\":\n    run_experiment()","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"outputs":[],"execution_count":null}]}