{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.11","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":94147,"databundleVersionId":11390004,"sourceType":"competition"},{"sourceId":237184244,"sourceType":"kernelVersion"}],"dockerImageVersionId":31011,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Prediction using Spotnet\n\nCredits to https://www.kaggle.com/code/tarundirector/histology-eda-spotnet-visual-spatial-dl . Kudos for the amazing notebook!\n\nThis notebook implements a complete deep learning pipeline for the Elucidata AI Challenge 2025, where the task is to predict the cell type composition (35 classes) at specific spatial transcriptomics spots from H&E-stained tissue images.\n\n🚀 Approach Summary\nWe combine image-based deep learning with spatial reasoning using a two-stage architecture:\n\n- SpotNet (custom hybrid CNN):\n\n    - Takes an image patch + spatial coordinates as input\n\n    - Uses a modified ResNet34 to process RGB + (x, y) location maps\n\n    - Outputs an initial 35-dimensional cell type prediction\n\n- SpotGNN (graph neural network refinement):\n\n    - Models inter-spot dependencies within each slide\n\n    - Learns how neighboring spots influence cell composition\n\n    - Refines SpotNet features using a GAT-based architecture\n\n🧱 Key Components\n- Data loading from HDF5 files with per-slide image arrays and spot-level metadata\n\n- 5-channel input patches: RGB + normalized x and y coordinate maps\n\n- Fourier-encoded spatial features to help learn tissue layout\n\n-  Spot-level k-fold cross-validation across all slides (not just one held-out slide)\n\n- Model checkpointing per fold, plus optional test-time ensemble\n\n- Optional GNN refinement step using torch_geometric to leverage spot connectivity\n\n- Final ensemble submission using predictions averaged over all folds\n\n📦 Outputs\n- Best SpotNet model per fold (best_model_fold{n}.pt)\n\n- Best SpotGNN model per fold (gnn_model_fold{n}.pt)\n\n- Final ensemble predictions written to submission.csv\n\n","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2025-04-30T14:06:18.576597Z","iopub.execute_input":"2025-04-30T14:06:18.576793Z","iopub.status.idle":"2025-04-30T14:06:20.419557Z","shell.execute_reply.started":"2025-04-30T14:06:18.576749Z","shell.execute_reply":"2025-04-30T14:06:20.418788Z"}}},{"cell_type":"code","source":"!ls ../working","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-03T16:08:41.437913Z","iopub.execute_input":"2025-05-03T16:08:41.438686Z","iopub.status.idle":"2025-05-03T16:08:41.578216Z","shell.execute_reply.started":"2025-05-03T16:08:41.438659Z","shell.execute_reply":"2025-05-03T16:08:41.577512Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install torch-scatter torch-sparse torch-cluster torch-spline-conv torch-geometric \\\n    -f https://data.pyg.org/whl/torch-2.0.0+cu118.html\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-03T16:08:41.579553Z","iopub.execute_input":"2025-05-03T16:08:41.579779Z","iopub.status.idle":"2025-05-03T16:08:44.628245Z","shell.execute_reply.started":"2025-05-03T16:08:41.579759Z","shell.execute_reply":"2025-05-03T16:08:44.627512Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"nb_type = \"Train\"","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-03T16:08:44.629256Z","iopub.execute_input":"2025-05-03T16:08:44.629499Z","iopub.status.idle":"2025-05-03T16:08:44.633717Z","shell.execute_reply.started":"2025-05-03T16:08:44.629476Z","shell.execute_reply":"2025-05-03T16:08:44.633046Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# --- Core Python Libraries ---\nimport os\nimport math\n\n# --- Data Processing ---\nimport numpy as np\nimport pandas as pd\nimport h5py\n\n# --- Plotting ---\nimport matplotlib.pyplot as plt\nfrom matplotlib import cm\nfrom matplotlib.patches import Rectangle\n\n# --- Deep Learning: PyTorch & TorchVision ---\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.utils.data import Dataset, DataLoader\nfrom torchvision import transforms, models\n\n# --- Metrics ---\nfrom scipy.stats import spearmanr","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-03T16:08:44.635330Z","iopub.execute_input":"2025-05-03T16:08:44.635714Z","iopub.status.idle":"2025-05-03T16:08:44.651756Z","shell.execute_reply.started":"2025-05-03T16:08:44.635698Z","shell.execute_reply":"2025-05-03T16:08:44.651119Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from torch_geometric.nn import GATConv\nfrom sklearn.neighbors import NearestNeighbors\nfrom torch_geometric.data import Data\nimport torch.nn.functional as F\n\nclass SpotGNN(torch.nn.Module):\n    def __init__(self, in_dim=640, hidden_dim=256, out_dim=35):  # 512+128 if using ResNet34\n        super().__init__()\n        self.gat1 = GATConv(in_dim, hidden_dim, heads=4, concat=True)\n        self.gat2 = GATConv(4 * hidden_dim, out_dim, heads=1)\n\n    def forward(self, x, edge_index):\n        x = F.relu(self.gat1(x, edge_index))\n        return self.gat2(x, edge_index)\n\n\ndef extract_spotnet_features(model, features):\n    model.eval()\n    all_feats = []\n    with torch.no_grad():\n        for img5, coord2 in features:\n            img5 = img5.unsqueeze(0).to(DEVICE)\n            coord2 = coord2.unsqueeze(0).to(DEVICE)\n            f_img = model.cnn(img5)\n            f_coord = model.coord_mlp(model.coord_ff(coord2))\n            feat = torch.cat([f_img, f_coord], dim=1).squeeze(0).cpu()\n            all_feats.append(feat)\n    return torch.stack(all_feats)  # [N, feature_dim]\n\n\ndef build_features_and_graph(df_slide, images, transform, k=6):\n    features = []\n    labels = []\n    coords = []\n    slide = df_slide[\"slide\"].iloc[0]\n    img = images[slide]\n    h, w, _ = img.shape\n\n    for i, row in df_slide.iterrows():\n        x, y = row[\"x\"], row[\"y\"]\n        x_norm, y_norm = row[\"x_norm\"], row[\"y_norm\"]\n        lbl = torch.tensor(row[\"label\"], dtype=torch.float32)\n\n        half = PATCH_SIZE // 2\n        y0, y1 = max(0,y-half), min(h,y+half)\n        x0, x1 = max(0,x-half), min(w,x+half)\n        patch = img[y0:y1, x0:x1]\n        pad_y = (max(0,half-y), max(0,(y+half)-h))\n        pad_x = (max(0,half-x), max(0,(x+half)-w))\n        patch = np.pad(patch, (pad_y, pad_x, (0,0)), mode='constant', constant_values=0)\n        img_t = transform(patch)\n\n        x_map = torch.full((1, PATCH_SIZE, PATCH_SIZE), x_norm)\n        y_map = torch.full((1, PATCH_SIZE, PATCH_SIZE), y_norm)\n        img5 = torch.cat([img_t, x_map, y_map], dim=0)\n\n        coord2 = torch.tensor([x_norm, y_norm], dtype=torch.float32)\n\n        features.append((img5, coord2))\n        labels.append(lbl)\n        coords.append([x_norm, y_norm])\n\n    coords = np.array(coords)\n    nbrs = NearestNeighbors(n_neighbors=k).fit(coords)\n    edge_index = nbrs.kneighbors_graph(coords).tocoo()\n    edge_index = torch.tensor(np.vstack([edge_index.row, edge_index.col]), dtype=torch.long)\n\n    return features, torch.stack(labels), edge_index\n\n\ndef train_spotgnn_for_fold(fold, df_fold, images, transform, model_path, save_path=\"gnn_fold{}.pt\", epochs=20, k=6):\n    print(f\"\\n🔁 Training SpotGNN for fold {fold}\")\n    \n    # Load SpotNet\n    spotnet = SpotNet(num_types=NUM_CELL_TYPES).to(DEVICE)\n    spotnet.load_state_dict(torch.load(model_path.format(fold)))\n    spotnet.eval()\n\n    # Build features and graph\n    features, labels, edge_index = build_features_and_graph(df_fold, images, transform, k=k)\n    x_all = extract_spotnet_features(spotnet, features).to(DEVICE)\n    labels = labels.to(DEVICE)\n    edge_index = edge_index.to(DEVICE)\n\n    # Train GNN\n    gnn = SpotGNN(in_dim=x_all.shape[1], out_dim=NUM_CELL_TYPES).to(DEVICE)\n    optimizer = torch.optim.Adam(gnn.parameters(), lr=1e-3)\n\n    for epoch in range(1, epochs + 1):\n        gnn.train()\n        preds = gnn(x_all, edge_index)\n        loss = F.mse_loss(preds, labels)\n        optimizer.zero_grad(); loss.backward(); optimizer.step()\n        print(f\"[Fold {fold} | Epoch {epoch}] GNN Loss: {loss.item():.4f}\")\n\n    # Save GNN\n    torch.save(gnn.state_dict(), save_path.format(fold))\n    print(f\"✅ Saved GNN model for fold {fold} to {save_path.format(fold)}\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-03T16:08:44.652470Z","iopub.execute_input":"2025-05-03T16:08:44.652682Z","iopub.status.idle":"2025-05-03T16:08:44.668435Z","shell.execute_reply.started":"2025-05-03T16:08:44.652659Z","shell.execute_reply":"2025-05-03T16:08:44.667820Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def forward_graph_spotnet(model, gnn, features, edge_index):\n    \"\"\"\n    Apply SpotNet to extract features, then GNN for refinement.\n    \"\"\"\n    model.eval(); gnn.eval()\n    with torch.no_grad():\n        xs = []\n        for img5, coord2 in features:\n            img5 = img5.unsqueeze(0).to(DEVICE)\n            coord2 = coord2.unsqueeze(0).to(DEVICE)\n            feat = model.cnn(img5)\n            coord_feat = model.coord_mlp(model.coord_ff(coord2))\n            x = torch.cat([feat, coord_feat], dim=1)\n            xs.append(x.squeeze(0).cpu())\n        x_all = torch.stack(xs)\n        preds = gnn(x_all.to(DEVICE), edge_index.to(DEVICE))\n    return preds.cpu()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-03T16:08:44.669311Z","iopub.execute_input":"2025-05-03T16:08:44.669499Z","iopub.status.idle":"2025-05-03T16:08:44.688486Z","shell.execute_reply.started":"2025-05-03T16:08:44.669485Z","shell.execute_reply":"2025-05-03T16:08:44.687780Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from torch_geometric.data import Data\nfrom sklearn.neighbors import NearestNeighbors\n\ndef build_slide_graph(df_slide, images, transform, k=6):\n    \"\"\"\n    Build a graph of spots within a single slide.\n    Each node has img5, coord2, and label.\n    \"\"\"\n    features = []\n    labels = []\n    coords = []\n    \n    slide = df_slide[\"slide\"].iloc[0]\n    img = images[slide]\n    h, w, _ = img.shape\n\n    for i, row in df_slide.iterrows():\n        x, y = row[\"x\"], row[\"y\"]\n        x_norm, y_norm = row[\"x_norm\"], row[\"y_norm\"]\n        lbl = torch.tensor(row[\"label\"], dtype=torch.float32)\n\n        # extract patch\n        half = PATCH_SIZE // 2\n        y0, y1 = max(0,y-half), min(h,y+half)\n        x0, x1 = max(0,x-half), min(w,x+half)\n        patch = img[y0:y1, x0:x1]\n        pad_y = (max(0,half-y), max(0,(y+half)-h))\n        pad_x = (max(0,half-x), max(0,(x+half)-w))\n        patch = np.pad(patch, (pad_y, pad_x, (0,0)), mode='constant', constant_values=0)\n        img_t = transform(patch)\n\n        x_map = torch.full((1, PATCH_SIZE, PATCH_SIZE), x_norm)\n        y_map = torch.full((1, PATCH_SIZE, PATCH_SIZE), y_norm)\n        img5 = torch.cat([img_t, x_map, y_map], dim=0)\n\n        coord2 = torch.tensor([x_norm, y_norm], dtype=torch.float32)\n\n        features.append((img5, coord2))\n        labels.append(lbl)\n        coords.append([x_norm, y_norm])\n\n    coords = np.array(coords)\n    nbrs = NearestNeighbors(n_neighbors=k).fit(coords)\n    edge_index = nbrs.kneighbors_graph(coords).tocoo()\n    edge_index = torch.tensor(np.vstack([edge_index.row, edge_index.col]), dtype=torch.long)\n\n    return features, labels, edge_index\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-03T16:08:44.689224Z","iopub.execute_input":"2025-05-03T16:08:44.689833Z","iopub.status.idle":"2025-05-03T16:08:44.708574Z","shell.execute_reply.started":"2025-05-03T16:08:44.689817Z","shell.execute_reply":"2025-05-03T16:08:44.708035Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Trnasform images and spots into a dataframe. Used for implemeting cross-validation more easily\ndef build_spot_dataframe(images, spots, shifts):\n    rows = []\n    for slide in images:\n        img = images[slide]\n        h, w, _ = img.shape\n        dx, dy = shifts.get(slide, (0, 0))\n        for i, spot in enumerate(spots[slide]):\n            x, y = spot[\"x\"] + dx, spot[\"y\"] + dy\n            x_norm, y_norm = x / w, y / h\n            label = [spot[f\"C{j}\"] for j in range(1, NUM_CELL_TYPES + 1)]\n            rows.append({\n                \"slide\": slide,\n                \"x\": int(x),\n                \"y\": int(y),\n                \"x_norm\": x_norm,\n                \"y_norm\": y_norm,\n                \"label\": label\n            })\n    return pd.DataFrame(rows)\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-03T16:08:44.709199Z","iopub.execute_input":"2025-05-03T16:08:44.709358Z","iopub.status.idle":"2025-05-03T16:08:44.725328Z","shell.execute_reply.started":"2025-05-03T16:08:44.709347Z","shell.execute_reply":"2025-05-03T16:08:44.724753Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import warnings\nwarnings.filterwarnings(\"ignore\", category=UserWarning)\n\nimport os, h5py, math\nimport numpy as np, pandas as pd\nimport torch\nimport torch.nn as nn, torch.optim as optim\nfrom torch.utils.data import Dataset, DataLoader\nfrom torchvision import transforms, models\nfrom scipy.stats import spearmanr\nimport matplotlib.pyplot as plt\n\n# --- 1) Load data dicts ---\nfile_path = \"/kaggle/input/el-hackathon-2025/elucidata_ai_challenge_data.h5\"\nwith h5py.File(file_path, \"r\") as f:\n    train_images = {k: np.array(v) for k,v in f[\"images/Train\"].items()}\n    train_spots  = {k: np.array(v) for k,v in f[\"spots/Train\"].items()}\n    test_images  = {k: np.array(v) for k,v in f[\"images/Test\"].items()}\n    test_spots   = {k: np.array(v) for k,v in f[\"spots/Test\"].items()}\n\n# --- 2) Slide-specific pixel shifts (as before) ---\nshifts = {\n    \"S_1\":(-65,-70), \"S_2\":(-65,-70),\n    \"S_3\":(-15,0),   \"S_4\":(0,0),\n    \"S_5\":(0,0),     \"S_6\":(-30,0),\n    # S_7: test → defaults to (0,0)\n}\n\n# --- 3) Hyperparams & Transforms ---\nPATCH_SIZE     = 384\nBATCH_SIZE     = 16\nLR             = 1e-4\nEPOCHS         = 20\nNUM_CELL_TYPES = 35\nDEVICE         = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\ntrain_tfm = transforms.Compose([\n    transforms.ToPILImage(),\n    transforms.RandomHorizontalFlip(),\n    transforms.RandomVerticalFlip(),\n    transforms.RandomRotation(90),\n    transforms.ColorJitter(0.2,0.2,0.2,0.1),\n    transforms.RandomAffine(10, scale=(0.9,1.1)),\n    transforms.Resize((PATCH_SIZE, PATCH_SIZE)),\n    transforms.ToTensor(),\n    transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225]),\n])\neval_tfm = transforms.Compose([\n    transforms.ToPILImage(),\n    transforms.Resize((PATCH_SIZE, PATCH_SIZE)),\n    transforms.ToTensor(),\n    transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225]),\n])\n\n# --- 4) Dataset: returns (img5chan, coord2, label) or (img5chan, coord2, idx) ---\nclass HistologyDatasetFromDF(Dataset):\n    def __init__(self, df, images, transform, mode=\"Train\"):\n        self.df = df.reset_index(drop=True)\n        self.images = images\n        self.transform = transform\n        self.mode = mode\n\n    def __len__(self): return len(self.df)\n\n    def __getitem__(self, idx):\n        row = self.df.iloc[idx]\n        slide = row[\"slide\"]\n        x, y = row[\"x\"], row[\"y\"]\n        x_norm, y_norm = row[\"x_norm\"], row[\"y_norm\"]\n    \n        img = self.images[slide]\n        h, w, _ = img.shape\n        half = PATCH_SIZE // 2\n\n        y0, y1 = max(0, y - half), min(h, y + half)\n        x0, x1 = max(0, x - half), min(w, x + half)\n        patch = img[y0:y1, x0:x1]\n        pad_y = (max(0, half - y), max(0, (y + half) - h))\n        pad_x = (max(0, half - x), max(0, (x + half) - w))\n        patch = np.pad(patch, (pad_y, pad_x, (0, 0)), mode='constant', constant_values=0)\n        img_t = self.transform(patch)\n\n        x_map = torch.full((1, PATCH_SIZE, PATCH_SIZE), x_norm)\n        y_map = torch.full((1, PATCH_SIZE, PATCH_SIZE), y_norm)\n        img5 = torch.cat([img_t, x_map, y_map], dim=0)\n\n        coord2 = torch.tensor([x_norm, y_norm], dtype=torch.float32)\n\n        if self.mode == \"Submission\":\n            return img5, coord2, torch.tensor(idx, dtype=torch.long)  # ID for submission\n        else:\n            label = torch.tensor(row[\"label\"], dtype=torch.float32)\n            return img5, coord2, label\n\n\n\n# --- 5) Fourier MLP on 2-D coords ---\nclass FourierFeature(nn.Module):\n    def __init__(self, in_dim=2, mapping_size=64, scale=10.0):\n        super().__init__()\n        B = torch.randn(mapping_size, in_dim) * scale\n        self.register_buffer(\"B\", B)\n\n    def forward(self, x):\n        x_proj = 2*math.pi * x @ self.B.t()        # [B,64]\n        return torch.cat([x_proj.sin(), x_proj.cos()], dim=-1)  # [B,128]\n\n# --- 6) Hybrid CoordConv + CoordMLP SpotNet ---\nclass SpotNet(nn.Module):\n    def __init__(self, num_types=35, backbone=\"resnet34\"):\n        super().__init__()\n        # --- Image tower with CoordConv(5→64) ---\n        if backbone==\"resnet34\":\n            cnn = models.resnet34(weights=models.ResNet34_Weights.DEFAULT)\n        else:\n            cnn = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)\n        old1 = cnn.conv1\n        new1 = nn.Conv2d(5, old1.out_channels,\n                         kernel_size=old1.kernel_size,\n                         stride=old1.stride,\n                         padding=old1.padding,\n                         bias=old1.bias is not None)\n        with torch.no_grad():\n            new1.weight[:,:3] = old1.weight  # copy RGB\n            new1.weight[:,3:] = old1.weight[:,:2].mean(dim=1,keepdim=True)*0.0\n        cnn.conv1 = new1\n        dim = cnn.fc.in_features\n        cnn.fc = nn.Identity()\n        self.cnn = cnn\n\n        # --- Coord MLP tower ---\n        self.coord_ff  = FourierFeature(in_dim=2, mapping_size=64, scale=10.0)\n        self.coord_mlp = nn.Sequential(\n            nn.Linear(128,256), nn.ReLU(),\n            nn.Linear(256,128), nn.ReLU()\n        )\n\n        # --- Fusion head ---\n        self.head = nn.Sequential(\n            nn.Linear(dim + 128, 512), nn.ReLU(), nn.Dropout(0.5),\n            nn.Linear(512, 256),       nn.ReLU(), nn.Dropout(0.3),\n            nn.Linear(256, num_types)\n        )\n\n    def forward(self, img5, coord2):\n        f = self.cnn(img5)              # [B, dim]\n        h = self.coord_ff(coord2)       # [B,128]\n        h = self.coord_mlp(h)           # [B,128]\n        x = torch.cat([f, h], dim=1)    # [B, dim+128]\n        return self.head(x)             # [B,35]\n\n# --- 7) Loss & validation ---\ndef spearman_loss(pred, target):\n    rho=[]\n    p_np=pred.detach().cpu().numpy()\n    t_np=target.detach().cpu().numpy()\n    for p,t in zip(p_np,t_np):\n        r = spearmanr(p,t)[0]\n        if not np.isnan(r): rho.append(r)\n    return 1 - np.mean(rho)\n\ndef valid_loop(model, loader):\n    model.eval()\n    preds, trues = [], []\n    with torch.no_grad():\n        for img5, coord2, y in loader:\n            img5, coord2 = img5.to(DEVICE), coord2.to(DEVICE)\n            preds.append(model(img5, coord2).cpu())\n            trues.append(y)\n    preds = torch.cat(preds); trues = torch.cat(trues)\n    return spearman_loss(preds, trues)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-03T16:08:44.726053Z","iopub.execute_input":"2025-05-03T16:08:44.726264Z","iopub.status.idle":"2025-05-03T16:08:46.519399Z","shell.execute_reply.started":"2025-05-03T16:08:44.726244Z","shell.execute_reply":"2025-05-03T16:08:46.518832Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from sklearn.model_selection import train_test_split\n\n#df_spots = build_spot_dataframe(train_images, train_spots, shifts)\n#df_train, df_val = train_test_split(df_spots, test_size=0.15, random_state=42, stratify=df_spots[\"slide\"])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-03T16:08:46.520937Z","iopub.execute_input":"2025-05-03T16:08:46.521158Z","iopub.status.idle":"2025-05-03T16:08:46.524298Z","shell.execute_reply.started":"2025-05-03T16:08:46.521142Z","shell.execute_reply":"2025-05-03T16:08:46.523700Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#nb_type = \"Train\"\nfrom sklearn.model_selection import StratifiedKFold\ndf_folds = []\nif nb_type == \"Train\":\n    # Step 1: Build spot-level DataFrame\n    df_spots = build_spot_dataframe(train_images, train_spots, shifts)\n\n    # Step 2: Initialize k-fold splitter\n    skf = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)\n\n    for fold, (train_idx, val_idx) in enumerate(skf.split(df_spots, df_spots[\"slide\"])):\n        print(f\"\\n🔁 Fold {fold + 1}\")\n        df_train = df_spots.iloc[train_idx].reset_index(drop=True)\n        df_val   = df_spots.iloc[val_idx].reset_index(drop=True)\n        df_folds.append((df_train, df_val))\n\n\n        ds_tr = HistologyDatasetFromDF(df_train, train_images, train_tfm)\n        ds_va = HistologyDatasetFromDF(df_val, train_images, eval_tfm)\n\n        loader_tr = DataLoader(ds_tr, batch_size=BATCH_SIZE, shuffle=True,  num_workers=4)\n        loader_va = DataLoader(ds_va, batch_size=BATCH_SIZE, shuffle=False, num_workers=4)\n\n        model     = SpotNet(num_types=NUM_CELL_TYPES).to(DEVICE)\n        optimizer = optim.Adam(model.parameters(), lr=LR)\n        criterion = nn.MSELoss()\n\n        history, best_spear = {\"train_loss\":[], \"val_spearman\":[]}, -1\n        print(\"🚀 Starting training for fold\", fold + 1)\n        for epoch in range(1, EPOCHS+1):\n            # train epoch\n            model.train()\n            tloss = 0\n            for img5, coord2, y in loader_tr:\n                img5, coord2, y = img5.to(DEVICE), coord2.to(DEVICE), y.to(DEVICE)\n                p = model(img5, coord2)\n                l = criterion(p, y)\n                optimizer.zero_grad(); l.backward(); optimizer.step()\n                tloss += l.item()\n            tloss /= len(loader_tr)\n\n            # validate\n            vspear = 1 - valid_loop(model, loader_va)\n            history[\"train_loss\"].append(tloss)\n            history[\"val_spearman\"].append(vspear)\n            print(f\"[Fold {fold+1} | Epoch {epoch}] Train-Loss: {tloss:.4f}  Val-Spearman: {vspear:.4f}\")\n            if vspear > best_spear:\n                best_spear = vspear\n                torch.save(model.state_dict(), f\"best_model_fold{fold+1}.pt\")\n                print(f\"✅ Saved best model for fold {fold+1} (Spearman={best_spear:.4f})\")\n\n        # plot\n        eps = list(range(1,EPOCHS+1))\n        plt.figure(figsize=(12,4))\n        plt.suptitle(f\"Fold {fold+1}\")\n        plt.subplot(1,2,1)\n        plt.plot(eps, history[\"train_loss\"], label=\"Train Loss\")\n        plt.xlabel(\"Epoch\"); plt.ylabel(\"MSE Loss\"); plt.legend()\n        plt.subplot(1,2,2)\n        plt.plot(eps, history[\"val_spearman\"], label=\"Val Spearman\", color=\"C1\")\n        plt.xlabel(\"Epoch\"); plt.ylabel(\"Spearman\"); plt.legend()\n        plt.tight_layout(); plt.show()\n\n        print(f\"\\n✅ Fold {fold+1} complete. Best Val Spearman = {best_spear:.4f}\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-03T16:08:46.524882Z","iopub.execute_input":"2025-05-03T16:08:46.525084Z","execution_failed":"2025-05-03T16:08:46.788Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if nb_type == \"Train\":\n    # --- 10) Graph Refinement after SpotNet Training ---\n    use_gnn = True  # Set to False to skip this block\n    \n    if use_gnn:\n        for fold in range(1, 6):\n            _, df_val = df_folds[fold - 1]  # use the stored val DataFrame\n            train_spotgnn_for_fold(\n                fold=fold,\n                df_fold=df_val,\n                images=train_images,\n                transform=eval_tfm,\n                model_path=\"best_model_fold{}.pt\",\n                save_path=\"gnn_model_fold{}.pt\",\n                epochs=20\n            )\n\n\n","metadata":{"trusted":true,"execution":{"execution_failed":"2025-05-03T16:08:46.788Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"nb_type = \"Submission\"\n\ndef build_test_dataframe(images, spots, shifts, slide=\"S_7\"):\n    img = images[slide]\n    h, w, _ = img.shape\n    dx, dy = shifts.get(slide, (0, 0))\n    rows = []\n    for i, spot in enumerate(spots[slide]):\n        x = int(spot[\"x\"] + dx)\n        y = int(spot[\"y\"] + dy)\n        x_norm = x / w\n        y_norm = y / h\n        rows.append({\n            \"slide\": slide,\n            \"x\": x,\n            \"y\": y,\n            \"x_norm\": x_norm,\n            \"y_norm\": y_norm,\n            \"label\": [0.0] * NUM_CELL_TYPES  # dummy label, not used at test time\n        })\n    return pd.DataFrame(rows)\n\n# --- 11) Submission block (SpotNet + GNN ensemble) ---\nif nb_type == \"Submission\":\n    df_ts = build_test_dataframe(test_images, test_spots, shifts, slide=\"S_7\")\n    ds_ts = HistologyDatasetFromDF(df_ts, test_images, eval_tfm, mode=\"Submission\")\n    loader_ts = DataLoader(ds_ts, batch_size=1, shuffle=False, num_workers=2)\n\n\n    # Load SpotNet + GNN models\n    spotnets = []\n    gnns = []\n    for k in range(1, 6):\n        # SpotNet\n        sn = SpotNet(num_types=NUM_CELL_TYPES).to(DEVICE)\n        sn.load_state_dict(torch.load(f\"best_model_fold{k}.pt\", map_location=DEVICE))\n        sn.eval()\n        spotnets.append(sn)\n\n        # GNN\n        gnn = SpotGNN(in_dim=640).to(DEVICE)  # 512+128 for ResNet34\n        gnn.load_state_dict(torch.load(f\"gnn_model_fold{k}.pt\", map_location=DEVICE))\n        gnn.eval()\n        gnns.append(gnn)\n\n    # Predict\n    subs = []\n    with torch.no_grad():\n        for img5, coord2, idxs in loader_ts:\n            img5, coord2 = img5.to(DEVICE), coord2.to(DEVICE)\n\n            fold_preds = []\n            for k in range(5):\n                # Feature extraction from SpotNet\n                f_img = spotnets[k].cnn(img5)\n                f_coord = spotnets[k].coord_mlp(spotnets[k].coord_ff(coord2))\n                feat = torch.cat([f_img, f_coord], dim=1)  # [1, 640]\n\n                # GNN (acting like a 1-node graph)\n                pred = gnns[k](feat, edge_index=torch.tensor([[0], [0]], device=DEVICE))  # dummy edge_index\n                fold_preds.append(pred.squeeze(0).cpu())\n\n            # Ensemble average\n            final_pred = torch.stack(fold_preds).mean(dim=0).numpy()\n            subs.append([idxs.item()] + final_pred.tolist())\n\n\n    # Write submission file\n    cols = [\"ID\"] + [f\"C{j}\" for j in range(1, NUM_CELL_TYPES + 1)]\n    sub_df = pd.DataFrame(subs, columns=cols)\n    sub_df.sort_values(\"ID\", inplace=True)\n    sub_df.to_csv(\"submission.csv\", index=False)\n    print(\"✅ submission.csv written:\", sub_df.shape)\n","metadata":{"trusted":true,"execution":{"execution_failed":"2025-05-03T16:08:46.788Z"}},"outputs":[],"execution_count":null}]}