Upload 10 files

dataset.py

@@ -0,0 +1,150 @@
+import base64
+import io
+import zlib
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+import torchvision.transforms.v2 as transforms
+from typing import Optional, Tuple
+
+def decode_array(encoded_base64_str):
+    decoded = base64.b64decode(encoded_base64_str)
+    decompressed = zlib.decompress(decoded)
+    return np.load(io.BytesIO(decompressed))
+
+def encode_array(array):
+    bytes_io = io.BytesIO()
+    np.save(bytes_io, array, allow_pickle=False)
+    compressed = zlib.compress(bytes_io.getvalue(), level=9)
+    return base64.b64encode(compressed).decode('utf-8')
+
+class BaseMicrographDataset(Dataset):
+    def __init__(self, df, window_size: int):
+        self.df = df
+        self.window_size = window_size
+
+    def __len__(self) -> int:
+        return len(self.df)
+
+    def load_and_normalize_image(self, encoded_image: str) -> torch.Tensor:
+        image = decode_array(encoded_image).astype(np.float32)
+        image = (image - image.min()) / (image.max() - image.min())
+        if len(image.shape) == 2:
+            image = image[np.newaxis, ...]
+        return torch.from_numpy(image)
+
+    def load_mask(self, encoded_mask: str) -> torch.Tensor:
+        mask = decode_array(encoded_mask).astype(np.float32)
+        if len(mask.shape) == 2:
+            mask = mask[np.newaxis, ...]
+        return torch.from_numpy(mask)
+
+    def pad_to_min_size(self, image: torch.Tensor, min_size: int) -> Tuple[torch.Tensor, Tuple[int, int]]:
+        _, h, w = image.shape
+        pad_h = max(0, min_size - h)
+        pad_w = max(0, min_size - w)
+        padded = torch.nn.functional.pad(image, (0, pad_w, 0, pad_h), mode="reflect")
+        return padded, (pad_h, pad_w)
+
+class TrainMicrographDataset(BaseMicrographDataset):
+    """Dataset for training with random augmentations"""
+
+    def __init__(self, df, window_size: int):
+        super().__init__(df, window_size)
+
+        # Define training-specific transforms
+        self.shared_transform = transforms.Compose([
+            transforms.RandomCrop(window_size),
+            transforms.RandomVerticalFlip(),
+            transforms.RandomHorizontalFlip()
+        ])
+        self.image_only_transforms = transforms.Compose([
+            transforms.GaussianBlur(7, sigma=(0.1, 2.))
+        ])
+
+    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]:
+        row = self.df.iloc[idx]
+
+        # Load and preprocess image
+        image = self.load_and_normalize_image(row['image'])
+        image, _ = self.pad_to_min_size(image, self.window_size)
+        image = self.image_only_transforms(image)
+
+        # Load and preprocess mask
+        mask = self.load_mask(row['mask'])
+        mask, _ = self.pad_to_min_size(mask, self.window_size)
+
+        # Apply shared transforms to both image and mask
+        stacked = torch.cat([image, mask], dim=0)
+        stacked = self.shared_transform(stacked)
+        image, mask = torch.split(stacked, [1, 1], dim=0)
+
+        return image, mask
+
+
+class ValidationMicrographDataset(BaseMicrographDataset):
+    """Dataset for validation using corner crops. This is a good idea because the regions of interest can be
+        at the edges of the image"""
+
+    def __init__(self, df, window_size: int):
+        super().__init__(df, window_size)
+        # Define 5 fixed crops: 4 corners + center
+        self.n_crops = 5
+
+    def __len__(self) -> int:
+        return len(self.df) * self.n_crops
+
+    def get_crop_coordinates(self, image_shape: Tuple[int, int], crop_idx: int) -> Tuple[int, int]:
+        """Get coordinates for specific crop index"""
+        h, w = image_shape
+
+        if crop_idx == 4:  # Center crop
+            h_start = (h - self.window_size) // 2
+            w_start = (w - self.window_size) // 2
+        else:
+            h_start = 0 if crop_idx < 2 else h - self.window_size
+            w_start = 0 if crop_idx % 2 == 0 else w - self.window_size
+
+        return h_start, w_start
+
+    def crop_tensors(self, image: torch.Tensor, mask: torch.Tensor,
+                     h_start: int, w_start: int) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Extract a crop from both image and mask"""
+        h_end = h_start + self.window_size
+        w_end = w_start + self.window_size
+
+        return (
+            image[:, h_start:h_end, w_start:w_end],
+            mask[:, h_start:h_end, w_start:w_end]
+        )
+
+    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]:
+        image_idx = idx // self.n_crops
+        crop_idx = idx % self.n_crops
+        row = self.df.iloc[image_idx]
+
+        # Load and preprocess image and mask
+        image = self.load_and_normalize_image(row['image'])
+        image, _ = self.pad_to_min_size(image, self.window_size)
+
+        mask = self.load_mask(row['mask'])
+        mask, _ = self.pad_to_min_size(mask, self.window_size)
+
+        # Get specific corner/center crop
+        h_start, w_start = self.get_crop_coordinates(image.shape[1:], crop_idx)
+        image, mask = self.crop_tensors(image, mask, h_start, w_start)
+
+        return image, mask
+
+
+class InferenceMicrographDataset(BaseMicrographDataset):
+    """Dataset for inference without any augmentations"""
+
+    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, str, Tuple[int, int]]:
+        row = self.df.iloc[idx]
+
+        # Load and preprocess image
+        image = self.load_and_normalize_image(row['image'])
+        image, padding = self.pad_to_min_size(image, self.window_size)
+
+        return image, row['Id'], padding

final_checkpoint.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:24ceee5d5db945c0d25ecfde13508b40f039842ab864aafa714f048ccc17a881
+size 1016916

infer.py

@@ -0,0 +1,76 @@
+#!/usr/bin/env python3
+import os
+import pandas as pd
+import torch
+from model import MicrographCleaner
+from dataset import InferenceMicrographDataset, decode_array
+from inference_utils import sliding_window_inference
+import matplotlib.pyplot as plt
+import tqdm
+
+
+def main():
+    # Create predictions directory if it doesn't exist
+    os.makedirs('predictions', exist_ok=True)
+
+    # Parameters
+    WINDOW_SIZE = 512
+    THRESHOLD = 0.5
+    OVERLAP = 0.5
+
+    # Load model
+    model = MicrographCleaner.load_from_checkpoint('final_checkpoint.ckpt', map_location='cpu')
+    model.eval()
+
+    # Load test data
+    test_df = pd.read_csv('test.csv')
+    test_dataset = InferenceMicrographDataset(test_df, window_size=WINDOW_SIZE)
+
+    # Process each image
+    unique_ids = set()
+    model.eval()
+    with torch.inference_mode():
+        for idx in tqdm.tqdm(range(len(test_dataset))):
+            image, image_id, (pad_h, pad_w) = test_dataset[idx]
+
+            # Skip if already processed
+            if image_id in unique_ids:
+                continue
+            unique_ids.add(image_id)
+
+            # Perform inference
+            pred = sliding_window_inference(
+                model,
+                image,
+                window_size=WINDOW_SIZE,
+                overlap=OVERLAP
+            )
+
+            # Remove padding if necessary
+            if pad_h > 0:
+                pred = pred[..., :-pad_h, :]
+            if pad_w > 0:
+                pred = pred[..., :-pad_w]
+
+            # Convert to binary mask
+            pred_mask = (pred > THRESHOLD).cpu().numpy()[0]
+
+            # Create visualization
+            orig_image = decode_array(test_df.iloc[idx]['image'])
+
+            fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 5))
+            ax1.imshow(orig_image, cmap='gray')
+            ax1.set_title('Original Image')
+            ax1.axis('off')
+
+            ax2.imshow(pred_mask, cmap='gray')
+            ax2.set_title('Predicted Mask')
+            ax2.axis('off')
+
+            plt.tight_layout()
+            plt.savefig(f'predictions/{image_id}_prediction.png')
+            plt.close()
+
+
+if __name__ == "__main__":
+    main()

inference_utils.py

@@ -0,0 +1,54 @@
+import torch
+import numpy as np
+
+def sliding_window_inference(model, image, window_size, overlap=0.5):
+    """Perform sliding window inference on large images"""
+    model.eval()
+
+    # Get dimensions
+    _, height, width = image.shape
+    stride = int(window_size * (1 - overlap))
+
+    # Calculate number of windows needed
+    n_h = int(np.ceil((height - window_size) / stride) + 1)
+    n_w = int(np.ceil((width - window_size) / stride) + 1)
+
+    # Create empty prediction map and count map for averaging
+    pred_map = torch.zeros((1, height, width)).to(model.device)
+    count_map = torch.zeros((1, height, width)).to(model.device)
+
+    # Slide window over image
+    with torch.no_grad():
+        for i in range(n_h):
+            for j in range(n_w):
+                # Calculate window boundaries
+                h_start = min(i * stride, height - window_size)
+                w_start = min(j * stride, width - window_size)
+                h_end = h_start + window_size
+                w_end = w_start + window_size
+
+                # Extract window
+                window = image[:, h_start:h_end, w_start:w_end]
+
+                # If window is smaller than window_size, pad it
+                if window.shape[1:] != (window_size, window_size):
+                    pad_h = window_size - window.shape[1]
+                    pad_w = window_size - window.shape[2]
+                    window = torch.nn.functional.pad(window, (0, pad_w, 0, pad_h))
+
+                # Make prediction
+                window = window.unsqueeze(0)  # Add batch dimension
+                pred = model(window)
+                pred = pred.squeeze(0)  # Remove batch dimension
+
+                # If window was padded, remove padding from prediction
+                if window.shape[2] - h_end + h_start > 0 or window.shape[3] - w_end + w_start > 0:
+                    pred = pred[:, :h_end - h_start, :w_end - w_start]
+
+                # Add prediction to map
+                pred_map[:, h_start:h_end, w_start:w_end] += pred
+                count_map[:, h_start:h_end, w_start:w_end] += 1
+
+    # Average overlapping predictions
+    final_pred = pred_map / count_map
+    return final_pred.cpu()

kaggle_id.txt

@@ -0,0 +1 @@
+rsancg00

model.py

@@ -0,0 +1,59 @@
+import torch
+import torch.nn as nn
+import pytorch_lightning as pl
+
+class SimpleCNN(nn.Module):
+    def __init__(self, n_hidden_layers, n_kernels, kernel_size):
+        super().__init__()
+        self.n_hidden_layers = n_hidden_layers
+        layers = [
+            nn.Conv2d(1, n_kernels, kernel_size=kernel_size, padding='same'),
+            nn.GroupNorm(4, n_kernels),
+            nn.PReLU()
+        ]
+
+        for _ in range(self.n_hidden_layers):
+            layers.extend([
+                nn.Conv2d(n_kernels, n_kernels, kernel_size=kernel_size, padding='same'),
+                nn.GroupNorm(4, n_kernels),
+                nn.PReLU(),
+            ])
+
+        layers.extend([
+            nn.Conv2d(n_kernels, 1, kernel_size=1),
+            nn.Sigmoid()
+        ])
+
+        self.conv_layers = nn.Sequential(*layers)
+
+    def forward(self, x):
+        return self.conv_layers(x)
+
+class MicrographCleaner(pl.LightningModule):
+    def __init__(self, n_hidden_layers=12, n_kernels=16, kernel_size=5, learning_rate=0.001):
+        super().__init__()
+        self.save_hyperparameters()
+        self.model = SimpleCNN(n_hidden_layers, n_kernels, kernel_size)
+        self.lossF = nn.BCELoss()
+        self.learning_rate = learning_rate
+        self.val_imgs_to_log = []
+
+    def forward(self, x):
+        return self.model(x)
+
+    def training_step(self, batch, batch_idx):
+        images, masks = batch
+        outputs = self(images)
+        loss = self.lossF(outputs, masks)
+        self.log('train_loss', loss, on_step=True, on_epoch=True, prog_bar=True)
+        return loss
+
+    def validation_step(self, batch, batch_idx):
+        images, masks = batch
+        outputs = self(images)
+        loss = self.lossF(outputs, masks)
+        self.log('val_loss', loss, on_epoch=True, prog_bar=True)
+        return loss
+
+    def configure_optimizers(self):
+        return torch.optim.Adam(self.parameters(), lr=self.learning_rate)

report_template.md

@@ -0,0 +1,38 @@
+# Cryo-EM Image Segmentation Report
+
+## Phase 1: Manual Implementation
+
+### Approach
+[Describe your approach to solving the problem, including the model architecture, loss functions, and training strategy]
+
+### Experiments
+
+| Experiment | Description | Training Loss | Validation Loss | Public Score | Private Score |
+|------------|-------------|---------------|-----------------|--------------|---------------|
+| Baseline   | Simple CNN  | 0.XX         | 0.XX           | 0.XX        | 0.XX         |
+| Exp 1      | [Change 1]  | 0.XX         | 0.XX           | 0.XX        | 0.XX         |
+| Exp 2      | [Change 2]  | 0.XX         | 0.XX           | 0.XX        | 0.XX         |
+
+### Training Curves
+[Insert training and validation loss curves for your final solution]
+
+### Analysis
+[Analyze the results of your experiments, discussing what worked and what didn't]
+
+## Phase 2: Open Resources
+
+### Approach
+[Describe the tools and pre-implemented solutions you used]
+
+### Results
+
+| Method | Description | Public Score | Private Score |
+|--------|-------------|--------------|---------------|
+| Method 1| [Description]| 0.XX        | 0.XX         |
+| Method 2| [Description]| 0.XX        | 0.XX         |
+
+### Comparison
+[Compare the results between Phase 1 and Phase 2, discussing the benefits and drawbacks of each approach]
+
+## Conclusions
+[Summarize your findings and discuss potential future improvements]

requirements.txt

@@ -0,0 +1,8 @@
+torch>=2.0.0
+torchvision>=0.15.0
+pytorch-lightning>=2.0.0
+pandas>=1.5.0
+numpy>=1.23.0
+matplotlib>=3.5.0
+scikit-learn>=1.0.0
+tqdm>=4.65.0

train.py

@@ -0,0 +1,62 @@
+#!/usr/bin/env python3
+import os
+import pandas as pd
+import pytorch_lightning as pl
+from pytorch_lightning.callbacks import ModelCheckpoint
+from pytorch_lightning.loggers import TensorBoardLogger
+from torch.utils.data import DataLoader
+from sklearn.model_selection import train_test_split
+
+from model import MicrographCleaner
+from dataset import TrainMicrographDataset, ValidationMicrographDataset
+
+
+def main():
+    # Training parameters
+    WINDOW_SIZE = 512
+    BATCH_SIZE = 8
+    N_EPOCHS = 3 #TODO, change this to many more epochs
+
+    # Load and split data
+    train_df = pd.read_csv('train.csv')
+    train_df, val_df = train_test_split(train_df, test_size=0.2, random_state=42)
+
+    # Create datasets and dataloaders
+    train_dataset = TrainMicrographDataset(train_df, window_size=WINDOW_SIZE)
+    val_dataset = ValidationMicrographDataset(val_df, window_size=WINDOW_SIZE)
+
+    train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=4)
+    val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, num_workers=4)
+
+    # Initialize model
+    model = MicrographCleaner()
+
+    # Setup training
+    logger = TensorBoardLogger('lightning_logs', name='micrograph_cleaner')
+    checkpoint_callback = ModelCheckpoint(
+        monitor='val_loss',
+        dirpath='checkpoints',
+        filename='micrograph-{epoch:02d}-{val_loss:.2f}',
+        save_top_k=3,
+        mode='min'
+    )
+
+    # Initialize trainer
+    trainer = pl.Trainer(
+        max_epochs=N_EPOCHS,
+        accelerator='auto',
+        devices=1,
+        logger=logger,
+        callbacks=[checkpoint_callback],
+        log_every_n_steps=10
+    )
+
+    # Train model
+    trainer.fit(model, train_loader, val_loader)
+
+    # Save final checkpoint as final_checkpoint.pt
+    trainer.save_checkpoint("final_checkpoint.pt")
+
+
+if __name__ == "__main__":
+    main()