{"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":11535317,"sourceType":"datasetVersion","datasetId":7234715},{"sourceId":11592169,"sourceType":"datasetVersion","datasetId":7269146}],"dockerImageVersionId":31011,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"## Notebook adapted from the original copy\n\nKudos to https://www.kaggle.com/code/tavuksuzdurum/resnet18 for the original notebook!\n\nMy modifications:\n\n- Modified cross validation to use random patches instead of a single slice S_6. Included S_6 in training\n- Added some plots to visualize the submission","metadata":{}},{"cell_type":"code","source":"import os\nimport time\nimport random\nfrom pathlib import Path\n\nimport cv2\nimport h5py\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torchvision\nfrom torch.utils.data import DataLoader, Dataset\nimport torchvision.models as models\nimport torchvision.transforms as T\nfrom sklearn.model_selection import train_test_split\nfrom scipy.stats import spearmanr\nfrom tqdm import tqdm\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport imgaug\nimport wandb  \nimport warnings\n\nimport h5py\nimport pandas as pd\nimport numpy as np\nimport matplotlib.pyplot as plt\nimport seaborn as sns\n\nwarnings.filterwarnings(\"ignore\")\n# Original weights: https://github.com/ozanciga/self-supervised-histopathology/releases/tag/tenpercent","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-05-12T13:56:59.224009Z","iopub.execute_input":"2025-05-12T13:56:59.224343Z","iopub.status.idle":"2025-05-12T13:56:59.230077Z","shell.execute_reply.started":"2025-05-12T13:56:59.224320Z","shell.execute_reply":"2025-05-12T13:56:59.229420Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\ndef is_interactive():\n   return os.environ.get('KAGGLE_KERNEL_RUN_TYPE','') == \"Interactive\"\nprint(\"is interactive session?\", is_interactive())\nnum_epochs = 4 if is_interactive() else 16\n\ntime_limit = 60 if is_interactive() else 3600","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-12T13:56:59.231356Z","iopub.execute_input":"2025-05-12T13:56:59.231655Z","iopub.status.idle":"2025-05-12T13:56:59.253278Z","shell.execute_reply.started":"2025-05-12T13:56:59.231638Z","shell.execute_reply":"2025-05-12T13:56:59.252721Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"CONFIG = {\n    \"seed\": 42,\n    \"data_path\": \"/kaggle/input/el-hackathon-2025\",\n    \"output_dir\": \"model_outputs\",\n    \"batch_size\": 32,\n    \"num_workers\": 6,\n    \"learning_rate\": 0.003,\n    \"weight_decay\": 1e-2,\n    \"scheduler_step_size\": 5,\n    \"scheduler_gamma\": 0.1,\n    \"num_classes\": 35,\n    \"image_size\": (108, 108),\n    \"patch_size\": 54,\n    \"max_epochs\": num_epochs,\n    \"checkpoint_epochs\": [3,4,5,7,9,11,13,15], # epoch 3 performed best\n    \"use_wandb\": False,  # Set to False if you don't want to use wandb\n    \"model_type\": \"resnet18\",  \n    \"mixed_precision\": False,  # Use mixed precision training\n}\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-12T13:56:59.253929Z","iopub.execute_input":"2025-05-12T13:56:59.254113Z","iopub.status.idle":"2025-05-12T13:56:59.269600Z","shell.execute_reply.started":"2025-05-12T13:56:59.254098Z","shell.execute_reply":"2025-05-12T13:56:59.268894Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\n# Utility functions\ndef set_seed(seed=42):\n    \"\"\"Set seeds for reproducibility\"\"\"\n    random.seed(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    torch.cuda.manual_seed_all(seed)\n    torch.backends.cudnn.deterministic = True\n    torch.backends.cudnn.benchmark = False\n    imgaug.seed(seed)\n\n\ndef get_device():\n    \"\"\"Get the appropriate device\"\"\"\n    if torch.cuda.is_available():\n        return torch.device(\"cuda\")\n    elif torch.backends.mps.is_available():\n        return torch.device(\"mps\")  \n    else:\n        return torch.device(\"cpu\")\n\n\ndef spearman_rank_correlation(x, y):\n    \"\"\"Calculate Spearman rank correlation\"\"\"\n    # Handle edge cases\n    if np.all(x == x[0]) or np.all(y == y[0]):\n        return 0.0\n    return spearmanr(x, y)[0]\n\n\ndef spearman_corr(preds, targets):\n    \"\"\"Calculate mean Spearman correlation across all samples\"\"\"\n    correlations = []\n    for i in range(len(preds)):\n        corr = spearman_rank_correlation(preds[i], targets[i])\n        # Only count valid correlations\n        if not np.isnan(corr):\n            correlations.append(corr)\n    return np.mean(correlations)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-12T13:56:59.270817Z","iopub.execute_input":"2025-05-12T13:56:59.271463Z","iopub.status.idle":"2025-05-12T13:56:59.289563Z","shell.execute_reply.started":"2025-05-12T13:56:59.271446Z","shell.execute_reply":"2025-05-12T13:56:59.288919Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import random\n\nclass HackhathonDataset(Dataset):\n    \"\"\"Dataset class for the hackathon competition\"\"\"\n\n    def __init__(self, data_path, transform=None, mode=\"train\"):\n        self.data_path = data_path\n        self.materials = []\n        self.transform = transform\n\n        # Define slides for train and validation sets\n        train_slides = [\"S_1\", \"S_2\", \"S_3\", \"S_4\", \"S_5\", \"S_6\"]\n        val_slides = [\"S_1\", \"S_2\", \"S_3\", \"S_4\", \"S_5\", \"S_6\"]  # Combine S1 to S6 for validation\n        test_slide = [\"S_7\"]\n        self.mode = mode\n\n        if mode == \"train\":\n            slide_list = train_slides\n        elif mode == \"val\":\n            slide_list = val_slides\n        else:\n            slide_list = test_slide\n\n        with h5py.File(f\"{self.data_path}/elucidata_ai_challenge_data.h5\", \"r\") as h5file:\n            images_group = \"images/Train\" if mode != \"test\" else \"images/Test\"\n            spots_group = \"spots/Train\" if mode != \"test\" else \"spots/Test\"\n\n            train_images = h5file[images_group]\n            train_spots = h5file[spots_group]\n\n            for slide_name in tqdm(slide_list, desc=f\"Loading {mode} data\"):\n                if slide_name in train_images.keys():\n                    image = np.array(train_images[slide_name])\n                    spots = np.array(train_spots[slide_name])\n                    df = pd.DataFrame(spots)\n                    self._split_into_patches(image, df, CONFIG[\"patch_size\"])\n\n        print(f\"{len(self.materials)} patches initialized for {mode} set\")\n\n    def __len__(self):\n        return len(self.materials)\n\n    def __getitem__(self, idx):\n        image, stats = self.materials[idx]\n\n        if self.transform:\n            image = self.transform(image)\n\n        # Use only gene expression values (excluding x, y coordinates)\n        stats = torch.tensor(stats[2:], dtype=torch.float32)\n        return image, stats\n\n    def _split_into_patches(self, arr, df, patch_size):\n        \"\"\"Split the image into patches centered on spot coordinates\"\"\"\n        h, w, c = arr.shape\n\n        for idx in range(len(df)):\n            row = df.iloc[idx]\n            x, y = int(row[\"x\"]), int(row[\"y\"])\n\n            half_size = patch_size // 2\n\n            # Ensure patches don't go outside image boundaries\n            y_min = max(y - half_size, 0)\n            y_max = min(y + half_size, h)\n            x_min = max(x - half_size, 0)\n            x_max = min(x + half_size, w)\n\n            patch = arr[y_min:y_max, x_min:x_max, :]\n\n            # Only include complete patches\n            if patch.shape[0] == patch_size and patch.shape[1] == patch_size:\n                self.materials.append([patch, row])\n            else:\n                # Handle incomplete patches by padding\n                padded_patch = np.zeros((patch_size, patch_size, c), dtype=patch.dtype)\n                padded_patch[:patch.shape[0], :patch.shape[1], :] = patch\n                self.materials.append([padded_patch, row])\n\n    def _get_random_spots(self, slides):\n        \"\"\"Randomly sample spots from the first 6 slices\"\"\"\n        random_spots = []\n        for slide in slides:\n            with h5py.File(f\"{self.data_path}/elucidata_ai_challenge_data.h5\", \"r\") as h5file:\n                spots_group = f\"spots/Train\" if self.mode != \"test\" else \"spots/Test\"\n                train_spots = h5file[spots_group]\n                if slide in train_spots:\n                    spots = np.array(train_spots[slide])\n                    random_spots.extend(random.sample(list(spots), k=min(len(spots), 1000)))  # Adjust 'k' for desired number of spots\n        return random_spots\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-12T13:56:59.364863Z","iopub.execute_input":"2025-05-12T13:56:59.365110Z","iopub.status.idle":"2025-05-12T13:56:59.378091Z","shell.execute_reply.started":"2025-05-12T13:56:59.365084Z","shell.execute_reply":"2025-05-12T13:56:59.377543Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class DifferentiableSpearmanLoss(nn.Module):\n    \"\"\"Differentiable approximation of Spearman correlation loss\"\"\"\n\n    def __init__(self, regularization_strength=1.0):\n        super().__init__()\n        self.regularization_strength = regularization_strength\n\n    def forward(self, y_pred, y_true):\n        y_pred = y_pred.float()\n        y_true = y_true.float()\n\n        # Calculate soft ranks\n        pred_rank = self._soft_rank(y_pred)\n        true_rank = self._soft_rank(y_true)\n\n        # Normalize ranks\n        pred_rank = F.normalize(pred_rank, dim=1)\n        true_rank = F.normalize(true_rank, dim=1)\n\n        # Calculate correlation\n        spearman = torch.sum(pred_rank * true_rank, dim=1)\n        return 1 - spearman.mean()\n\n    def _soft_rank(self, x, regularization_strength=None):\n        if regularization_strength is None:\n            regularization_strength = self.regularization_strength\n\n        x = x.unsqueeze(-1)  # [batch, n, 1]\n        diff = x - x.transpose(-1, -2)  # [batch, n, n]\n        P = torch.sigmoid(-regularization_strength * diff)  # pairwise comparisons\n        ranks = P.sum(dim=-1)  # approximate ranks\n        return ranks\n\n\nclass CombinedLoss(nn.Module):\n    \"\"\"Combined loss function using L1 and Spearman correlation\"\"\"\n\n    def __init__(self, alpha=0.5, regularization_strength=1.0):\n        super().__init__()\n        self.l1 = nn.L1Loss()\n        self.spearman = DifferentiableSpearmanLoss(regularization_strength)\n        self.alpha = alpha\n\n    def forward(self, y_pred, y_true):\n        l1_loss = self.l1(y_pred, y_true)\n        spearman_loss = self.spearman(y_pred, y_true)\n        return l1_loss + self.alpha * spearman_loss\n\n\ndef create_model(model_type, num_classes):\n    \"\"\"Create and initialize the model\"\"\"\n    if model_type == \"resnet18\":\n        #model = models.wide_resnet101_2(weights=models.Wide_ResNet101_2_Weights.IMAGENET1K_V1)\n        #model.fc = nn.Linear(model.fc.in_features, num_classes)\n        MODEL_PATH = '/kaggle/input/ckpt-file/tenpercent_resnet18.ckpt'\n        RETURN_PREACTIVATION = False  # return features from the model, if false return classification logits\n        NUM_CLASSES = 35  # only used if RETURN_PREACTIVATION = False\n\n        def load_model_weights(model, weights):\n\n            model_dict = model.state_dict()\n            weights = {k: v for k, v in weights.items() if k in model_dict}\n            if weights == {}:\n                print('No weight could be loaded..')\n            model_dict.update(weights)\n            model.load_state_dict(model_dict)\n\n            return model\n\n        model = models.resnet18(pretrained=False)\n\n        state = torch.load(MODEL_PATH, map_location='cuda', weights_only=False)\n\n        state_dict = state['state_dict']\n        for key in list(state_dict.keys()):\n            state_dict[key.replace('model.', '').replace('resnet.', '')] = state_dict.pop(key)\n\n        model = load_model_weights(model, state_dict)\n\n        if RETURN_PREACTIVATION:\n            model.fc = torch.nn.Sequential()\n        else:\n            model.fc = torch.nn.Linear(model.fc.in_features, NUM_CLASSES)\n\n    elif model_type == \"resnet152\":\n        model = models.resnet152(weights=models.ResNet152_Weights.IMAGENET1K_V1)\n        model.fc = nn.Linear(model.fc.in_features, num_classes)\n    elif model_type == \"efficientnet_b0\":\n        model = models.efficientnet_b0(weights=models.EfficientNet_B0_Weights.IMAGENET1K_V1)\n        model.classifier = nn.Linear(model.classifier[1].in_features, num_classes)\n    elif model_type == \"efficientnet_b3\":\n        model = models.efficientnet_b7(weights=models.EfficientNet_B7_Weights.IMAGENET1K_V1)\n        model.classifier = nn.Linear(model.classifier[1].in_features, num_classes)\n    elif model_type == \"convnext_large\":\n        model = models.convnext_large(weights=models.ConvNeXt_Large_Weights.IMAGENET1K_V1)\n        in_features = model.classifier[2].in_features\n        model.classifier[2] = nn.Linear(in_features, num_classes)\n    else:\n        raise ValueError(f\"Unsupported model type: {model_type}\")\n\n    return model\n\n\n\ndef get_transforms():\n    \"\"\"Get data transformations for training and validation\"\"\"\n    train_transform = T.Compose([\n        T.ToTensor(),\n        T.Resize(CONFIG[\"image_size\"]),\n        T.RandomApply([T.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.4, hue=0.1)], p=0.8),\n        T.RandomHorizontalFlip(p=0.5),\n        T.RandomVerticalFlip(p=0.5),\n        T.RandomRotation(degrees=45),\n        T.RandomAffine(degrees=0, translate=(0.1, 0.1), scale=(0.9, 1.1)),\n    ])\n\n    val_transform = T.Compose([\n        T.ToTensor(),\n        T.Resize(CONFIG[\"image_size\"]),\n    ])\n\n    return train_transform, val_transform","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-12T13:56:59.379257Z","iopub.execute_input":"2025-05-12T13:56:59.379467Z","iopub.status.idle":"2025-05-12T13:56:59.397930Z","shell.execute_reply.started":"2025-05-12T13:56:59.379453Z","shell.execute_reply":"2025-05-12T13:56:59.397162Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def save_checkpoint(model, optimizer, scheduler, epoch, metrics, filename):\n    \"\"\"Save model checkpoint\"\"\"\n    checkpoint = {\n        'epoch': epoch,\n        'model_state_dict': model.state_dict(),\n        'optimizer_state_dict': optimizer.state_dict(),\n        'scheduler_state_dict': scheduler.state_dict(),\n        'metrics': metrics\n    }\n    torch.save(checkpoint, filename)\n\n\ndef train_one_epoch(model, dataloader, loss_fn, optimizer, device, scaler=None):\n    \"\"\"Train model for one epoch\"\"\"\n    model.train()\n    epoch_loss = 0\n    all_preds, all_labels = [], []\n\n    progress_bar = tqdm(dataloader, desc=\"Training\")\n\n    for images, labels in progress_bar:\n        images = images.to(device)\n        labels = labels.to(device)\n\n        if scaler:  # Using mixed precision\n            with torch.cuda.amp.autocast():\n                outputs = model(images)\n                loss = loss_fn(outputs, labels)\n\n            optimizer.zero_grad()\n            scaler.scale(loss).backward()\n            scaler.step(optimizer)\n            scaler.update()\n        else:\n            outputs = model(images)\n            loss = loss_fn(outputs, labels)\n\n            optimizer.zero_grad()\n            loss.backward()\n            optimizer.step()\n\n        epoch_loss += loss.item()\n\n        # Store predictions and labels for metrics calculation\n        all_preds.extend(outputs.detach().cpu().numpy())\n        all_labels.extend(labels.detach().cpu().numpy())\n\n        # Update progress bar\n        progress_bar.set_postfix({\"loss\": f\"{loss.item():.4f}\"})\n\n    # Calculate metrics\n    avg_loss = epoch_loss / len(dataloader)\n    spearman_score = spearman_corr(all_preds, all_labels)\n\n    return avg_loss, spearman_score, all_preds, all_labels\n\n\ndef validate(model, dataloader, loss_fn, device):\n    \"\"\"Validate model on validation set\"\"\"\n    model.eval()\n    val_loss = 0\n    all_preds, all_labels = [], []\n\n    with torch.no_grad():\n        for images, labels in tqdm(dataloader, desc=\"Validation\"):\n            images = images.to(device)\n            labels = labels.to(device)\n\n            outputs = model(images)\n            loss = loss_fn(outputs, labels)\n\n            val_loss += loss.item()\n\n            # Store predictions and labels for metrics calculation\n            all_preds.extend(outputs.cpu().numpy())\n            all_labels.extend(labels.cpu().numpy())\n\n    # Calculate metrics\n    avg_loss = val_loss / len(dataloader)\n    spearman_score = spearman_corr(all_preds, all_labels)\n\n    return avg_loss, spearman_score, all_preds, all_labels\n\n\ndef predict_test_set(model, data_path, device, transform):\n    \"\"\"Generate predictions for the test set\"\"\"\n    with h5py.File(f\"{data_path}/elucidata_ai_challenge_data.h5\", \"r\") as f:\n        test_spots = f[\"spots/Test\"]\n        test_images = f[\"images/Test\"]\n        sample = 'S_7'  # Test sample\n        image = np.array(test_images[sample])\n        spots = np.array(test_spots[sample])\n        x, y = spots[\"x\"], spots[\"y\"]\n\n        outputs = []\n        model.eval()\n        with torch.no_grad():\n            patch_size = CONFIG[\"patch_size\"]\n            for x_, y_ in tqdm(zip(x, y), desc=\"Generating predictions\", total=len(x)):\n                half_size = patch_size // 2\n                # Ensure indices are within bounds\n                y_min = max(y_ - half_size, 0)\n                y_max = min(y_ + half_size, image.shape[0])\n                x_min = max(x_ - half_size, 0)\n                x_max = min(x_ + half_size, image.shape[1])\n                patch = image[y_min:y_max, x_min:x_max, :]\n\n                # Handle incomplete patches\n                if patch.shape[0] != patch_size or patch.shape[1] != patch_size:\n                    padded_patch = np.zeros((patch_size, patch_size, 3), dtype=patch.dtype)\n                    padded_patch[:patch.shape[0], :patch.shape[1], :] = patch\n                    patch = padded_patch\n\n                patch_tensor = transform(patch)\n                patch_tensor = patch_tensor.to(device)\n                output = model(patch_tensor.unsqueeze(0)).cpu().numpy()\n                outputs.append(output[0])\n\n    return np.array(outputs), x, y\n\n\ndef save_submission(predictions):\n    \"\"\"Save predictions to submission file\"\"\"\n    example_df = pd.read_csv(\"/kaggle/input/sample-submission/submission (1).csv\")\n    ID = example_df[\"ID\"]\n    output_df = pd.DataFrame(predictions)\n    submission_df = pd.concat([ID, output_df], axis=1)\n    submission_df.columns = example_df.columns\n\n    output_file = \"submission.csv\"\n    submission_df.to_csv(output_file, index=False)\n    print(f\"Saved submission to {output_file}\")\n\n    return output_file\n\n\ndef plot_cell_abundance(predictions):\n    # Get the list of cell types (columns)\n    cell_types = predictions.columns\n\n    # Set the number of rows and columns for the grid of subplots\n    num_cell_types = len(cell_types)\n    cols = 5  # You can adjust this depending on the number of cell types\n    rows = (num_cell_types // cols) + (1 if num_cell_types % cols > 0 else 0)\n    \n    fig, axes = plt.subplots(rows, cols, figsize=(cols*5, rows*5))\n    \n    # Flatten axes array in case of multiple rows and columns\n    axes = axes.flatten()\n\n    # Plot abundance for each cell type\n    for i, cell_type in enumerate(cell_types):\n        axes[i].bar(np.arange(len(predictions)), predictions[cell_type], color='skyblue')\n        axes[i].set_title(f'Abundance of {cell_type}')\n        axes[i].set_xlabel('Sample Index')\n        axes[i].set_ylabel('Abundance')\n        \n    # Remove extra subplots if necessary\n    for i in range(num_cell_types, len(axes)):\n        axes[i].axis('off')\n\n    # Adjust layout to prevent overlap\n    plt.tight_layout()\n    plt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-12T13:56:59.398730Z","iopub.execute_input":"2025-05-12T13:56:59.398999Z","iopub.status.idle":"2025-05-12T13:56:59.420777Z","shell.execute_reply.started":"2025-05-12T13:56:59.398981Z","shell.execute_reply":"2025-05-12T13:56:59.420076Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def main():\n    \"\"\"Main training function\"\"\"\n    set_seed(CONFIG[\"seed\"])\n    device = get_device()\n    print(f\"Using device: {device}\")\n\n    # Create data transformations\n    train_transform, val_transform = get_transforms()\n\n    # Create datasets and dataloaders\n    train_dataset = HackhathonDataset(\n        CONFIG[\"data_path\"],\n        transform=train_transform,\n        mode=\"train\"\n    )\n    val_dataset = HackhathonDataset(\n        CONFIG[\"data_path\"],\n        transform=val_transform,\n        mode=\"val\"\n    )\n\n    train_loader = DataLoader(\n        train_dataset,\n        batch_size=CONFIG[\"batch_size\"],\n        shuffle=True,\n        num_workers=CONFIG[\"num_workers\"],\n        pin_memory=True\n    )\n\n    val_loader = DataLoader(\n        val_dataset,\n        batch_size=CONFIG[\"batch_size\"],\n        shuffle=False,\n        num_workers=CONFIG[\"num_workers\"],\n        pin_memory=True\n    )\n\n    # Create model\n    model = create_model(CONFIG[\"model_type\"], CONFIG[\"num_classes\"])\n    model = model.to(device)\n\n    # Create loss function, optimizer and scheduler\n    loss_fn = CombinedLoss(alpha=0.7)\n    optimizer = torch.optim.AdamW(\n        model.parameters(),\n        lr=CONFIG[\"learning_rate\"],\n        weight_decay=CONFIG[\"weight_decay\"]\n    )\n\n    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(\n        optimizer,\n        mode='min',\n        factor=CONFIG[\"scheduler_gamma\"],\n        patience=3,\n\n    )\n\n    # Initialize mixed precision scaler if needed\n    scaler = torch.cuda.amp.GradScaler() if CONFIG[\"mixed_precision\"] and device.type == \"cuda\" else None\n\n    # Initialize wandb if enabled\n    if CONFIG[\"use_wandb\"]:\n        wandb.init(\n            project=\"hackathon-gene-expression\",\n            config=CONFIG,\n            name=f\"{CONFIG['model_type']}_run\"\n        )\n\n        # Log model architecture\n        wandb.watch(model)\n\n    # Training loop\n    best_val_spearman = -1\n    for epoch in range(CONFIG[\"max_epochs\"]):\n        print(f\"\\nEpoch {epoch + 1}/{CONFIG['max_epochs']}\")\n\n        # Train one epoch\n        train_loss, train_spearman, _, _ = train_one_epoch(\n            model, train_loader, loss_fn, optimizer, device, scaler\n        )\n\n        # Validate\n        val_loss, val_spearman, _, _ = validate(\n            model, val_loader, loss_fn, device\n        )\n\n        # Update scheduler\n        scheduler.step(val_loss)\n\n        # Log metrics\n        metrics = {\n            \"train_loss\": train_loss,\n            \"train_spearman\": train_spearman,\n            \"val_loss\": val_loss,\n            \"val_spearman\": val_spearman,\n            \"learning_rate\": optimizer.param_groups[0]['lr']\n        }\n        for param_group in optimizer.param_groups:\n            print(f\"current learning rate: {param_group['lr']}\")\n\n        print(f\"Train Loss: {train_loss:.4f}, Train Spearman: {train_spearman:.4f}\")\n        print(f\"Val Loss: {val_loss:.4f}, Val Spearman: {val_spearman:.4f}\")\n\n        if CONFIG[\"use_wandb\"]:\n            wandb.log(metrics)\n\n        # Save best model\n        #if val_spearman > best_val_spearman:\n        #    best_val_spearman = val_spearman\n        #    save_checkpoint(\n        #        model, optimizer, scheduler, epoch, metrics,\n        #        f\"{CONFIG['output_dir']}/{CONFIG['model_type']}_best.pth\"\n        #    )\n        #    print(f\"Saved best model with validation Spearman: {val_spearman:.4f}\")\n\n        # Generate and save predictions at checkpoint epochs\n        if epoch in CONFIG[\"checkpoint_epochs\"]:\n            # Save regular checkpoint\n            #save_checkpoint(\n            #    model, optimizer, scheduler, epoch, metrics,\n            #    f\"{CONFIG['output_dir']}/{CONFIG['model_type']}_epoch_{epoch}.pth\"\n            #)\n\n            # Generate test predictions\n            test_preds, _, _ = predict_test_set(model, CONFIG[\"data_path\"], device, val_transform)\n\n            # Save submission file\n            submission_file = save_submission(test_preds)\n            #plot_cell_abundance(test_preds)\n\n            if CONFIG[\"use_wandb\"]:\n                wandb.save(submission_file)\n\n    # Clean up\n    if CONFIG[\"use_wandb\"]:\n        wandb.finish()\n\n    print(\"Training completed!\")\n    return test_preds\n\n\ntest_preds = main()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-12T13:56:59.422303Z","iopub.execute_input":"2025-05-12T13:56:59.422503Z","iopub.status.idle":"2025-05-12T13:59:10.071779Z","shell.execute_reply.started":"2025-05-12T13:56:59.422489Z","shell.execute_reply":"2025-05-12T13:59:10.070925Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!head submission.csv","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-12T13:59:10.073320Z","iopub.execute_input":"2025-05-12T13:59:10.073597Z","iopub.status.idle":"2025-05-12T13:59:10.244823Z","shell.execute_reply.started":"2025-05-12T13:59:10.073574Z","shell.execute_reply":"2025-05-12T13:59:10.243941Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"submission_df = pd.read_csv(\"submission.csv\")\nsubmission_df","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-12T13:59:10.246013Z","iopub.execute_input":"2025-05-12T13:59:10.246284Z","iopub.status.idle":"2025-05-12T13:59:10.284670Z","shell.execute_reply.started":"2025-05-12T13:59:10.246261Z","shell.execute_reply":"2025-05-12T13:59:10.284090Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import matplotlib.pyplot as plt\nplot_cell_abundance(submission_df)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-12T13:59:10.285383Z","iopub.execute_input":"2025-05-12T13:59:10.285590Z","iopub.status.idle":"2025-05-12T14:00:45.067938Z","shell.execute_reply.started":"2025-05-12T13:59:10.285574Z","shell.execute_reply":"2025-05-12T14:00:45.067131Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\ndef plot_cell_abundance_with_spots_and_image(submission_df, hdf5_path):\n    # Get the list of cell types (columns C1, C2, ..., C35)\n    cell_types = [col for col in submission_df.columns if col.startswith('C')]\n\n    # Set the number of rows and columns for the grid of subplots\n    num_cell_types = len(cell_types)\n    cols = 5  # Number of columns in the grid\n    rows = (num_cell_types // cols) + (1 if num_cell_types % cols > 0 else 0)\n\n    # Open the HDF5 file to load the image of the test slice (S_7)\n    with h5py.File(hdf5_path, 'r') as f:\n        test_slice_image = np.array(f['images/Test/S_7'])  # Load the S_7 image\n        spots_coords = np.array(f['spots/Test/S_7'])  # Load spots coordinates (x, y)\n\n    # Convert spots coordinates into a DataFrame and create an 'ID' column\n    spots_df = pd.DataFrame(spots_coords, columns=['x', 'y'] + [f'CellType_{i}' for i in range(1, 36)])\n    spots_df['ID'] = spots_df.index  # Add an 'ID' column to match submission_df\n\n    # Merge the submission dataframe with the spots dataframe on 'ID'\n    submission_df = submission_df.merge(spots_df, on='ID')\n    # Create subplots\n    fig, axes = plt.subplots(rows, cols, figsize=(cols * 5, rows * 5))\n\n    # Flatten axes for easy indexing\n    axes = axes.flatten()\n\n    # Plot the image of the test slice (S_7) in the first subplot\n    axes[0].imshow(test_slice_image)\n    axes[0].axis('off')\n    axes[0].set_title('Test Slice (S_7)')\n\n    # Plot the distribution of each cell type in the grid, along with the spots\n    for i, cell_type in enumerate(cell_types):\n        if i == 0:\n            # Skip the first subplot as it already contains the image\n            continue\n\n        # Scatter plot for spots with coloring based on cell type abundance\n        scatter = axes[i].scatter(submission_df['x'], submission_df['y'], \n                                  c=submission_df[cell_type], cmap='viridis', s=10, alpha=0.6)\n\n        # Title and labels\n        axes[i].set_title(f'{cell_type} - Spot Distribution')\n        axes[i].set_xlabel('X Coordinate')\n        axes[i].set_ylabel('Y Coordinate')\n\n        # Add color bar to show abundance levels\n        fig.colorbar(scatter, ax=axes[i], label='Abundance')\n\n    # Remove extra subplots if necessary\n    for i in range(num_cell_types + 1, len(axes)):  # +1 because the first subplot contains the image\n        axes[i].axis('off')\n\n    # Adjust layout to prevent overlap\n    plt.tight_layout()\n    plt.show()\n\n# Example usage:\n# submission_df = pd.read_csv('submission.csv')  # Assuming the data is stored in a CSV\nhdf5_path = '/kaggle/input/el-hackathon-2025/elucidata_ai_challenge_data.h5'\nplot_cell_abundance_with_spots_and_image(submission_df, hdf5_path)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-12T14:00:45.068899Z","iopub.execute_input":"2025-05-12T14:00:45.069429Z","iopub.status.idle":"2025-05-12T14:00:55.210309Z","shell.execute_reply.started":"2025-05-12T14:00:45.069406Z","shell.execute_reply":"2025-05-12T14:00:55.209520Z"}},"outputs":[],"execution_count":null}]}