Merge pull request #3 from schwartzlab-methods/copilot/code-review-for-spaghetti-api

Fix critical bug and improve code quality for SPAGHETTI API

Author: YinniKun

spaghetti/__init__.py

@@ -0,0 +1,18 @@
+"""
+SPAGHETTI - Structural Phase Adaptation via Generative Histological Enhancement
+and Texture-preserving Translation Integration
+
+A PyTorch implementation for phase-contrast microscopy image transformation.
+"""
+
+from spaghetti.inferences import Spaghetti
+from spaghetti.dataset import TrainingDataset
+from spaghetti.train import train_spaghetti
+from spaghetti import utils
+
+__all__ = [
+    "Spaghetti",
+    "TrainingDataset",
+    "train_spaghetti",
+    "utils",
+]

spaghetti/_spaghetti_modules.py

@@ -1,41 +1,43 @@
-'''
+"""
 Modules for the SPAGHETTI model to translate microscopy images to H&E images
-'''
+"""
 import torch.nn as nn
 import torch.nn.functional as F
 import torch
 
+
 class ResidualBlock(nn.Module):
     def __init__(self, in_channels):
         super(ResidualBlock, self).__init__()
         self.block = nn.Sequential(
-            nn.ReflectionPad2d(1), # padding, keep the image size constant after next conv2d
+            nn.ReflectionPad2d(1),  # padding, keep the image size constant after next conv2d
             nn.Conv2d(in_channels, in_channels, 3),
             nn.InstanceNorm2d(in_channels),
             nn.ReLU(inplace=True),
             nn.ReflectionPad2d(1),
             nn.Conv2d(in_channels, in_channels, 3),
             nn.InstanceNorm2d(in_channels)
         )
-    
+
     def forward(self, x):
         return x + self.block(x)
-    
+
+
 class GeneratorResNet(nn.Module):
     def __init__(self, in_channels, num_residual_blocks=9):
         super(GeneratorResNet, self).__init__()
-        
+
         # Inital Convolution  3*256*256 -> 64*256*256
-        out_channels=64
+        out_channels = 64
         self.conv = nn.Sequential(
-            nn.ReflectionPad2d(in_channels), # padding, keep the image size constant after next conv2d
+            nn.ReflectionPad2d(in_channels),  # padding, keep the image size constant after next conv2d
             nn.Conv2d(in_channels, out_channels, 2*in_channels+1),
             nn.InstanceNorm2d(out_channels),
             nn.ReLU(inplace=True),
         )
-        
+
         channels = out_channels
-        
+
         # Downsampling   64*256*256 -> 128*128*128 -> 256*64*64
         self.down = []
         for _ in range(2):
@@ -47,31 +49,31 @@ def __init__(self, in_channels, num_residual_blocks=9):
             ]
             channels = out_channels
         self.down = nn.Sequential(*self.down)
-        
+
         # Transformation (ResNet)  256*64*64
         self.trans = [ResidualBlock(channels) for _ in range(num_residual_blocks)]
         self.trans = nn.Sequential(*self.trans)
-        
+
         # Upsampling  256*64*64 -> 128*128*128 -> 64*256*256
         self.up = []
         for _ in range(2):
             out_channels = channels // 2
             self.up += [
-                nn.Upsample(scale_factor=2), # bilinear interpolation
+                nn.Upsample(scale_factor=2),  # bilinear interpolation
                 nn.Conv2d(channels, out_channels, 3, stride=1, padding=1),
                 nn.InstanceNorm2d(out_channels),
                 nn.ReLU(inplace=True),
             ]
             channels = out_channels
         self.up = nn.Sequential(*self.up)
-        
+
         # Out layer  64*256*256 -> 3*256*256
         self.out = nn.Sequential(
             nn.ReflectionPad2d(in_channels),
             nn.Conv2d(channels, in_channels, 2*in_channels+1),
             nn.Tanh()
         )
-    
+
     def forward(self, x):
         x = self.conv(x)
         x = self.down(x)
@@ -80,42 +82,44 @@ def forward(self, x):
         x = self.out(x)
         return x
 
+
 class Discriminator(nn.Module):
     def __init__(self, in_channels):
         super(Discriminator, self).__init__()
-        
+
         self.model = nn.Sequential(
-            *self.block(in_channels, 64, normalize=False), # 3*256*256 -> 64*128*128 
+            *self.block(in_channels, 64, normalize=False),  # 3*256*256 -> 64*128*128
             *self.block(64, 128),  # 64*128*128 -> 128*64*64
-            *self.block(128, 256), # 128*64*64 -> 256*32*32
-            *self.block(256, 512), # 256*32*32 -> 512*16*16
-            
-            nn.ZeroPad2d((1,0,1,0)), # padding left and top   512*16*16 -> 512*17*17
-            nn.Conv2d(512, 1, 4, padding=1) # 512*17*17 -> 1*16*16
+            *self.block(128, 256),  # 128*64*64 -> 256*32*32
+            *self.block(256, 512),  # 256*32*32 -> 512*16*16
+
+            nn.ZeroPad2d((1, 0, 1, 0)),  # padding left and top   512*16*16 -> 512*17*17
+            nn.Conv2d(512, 1, 4, padding=1)  # 512*17*17 -> 1*16*16
         )
-        
+
         self.scale_factor = 16
-    
+
     @staticmethod
     def block(in_channels, out_channels, normalize=True):
         layers = [nn.Conv2d(in_channels, out_channels, 4, stride=2, padding=1)]
         if normalize:
             layers.append(nn.InstanceNorm2d(out_channels))
         layers.append(nn.LeakyReLU(0.2, inplace=True))
-        
+
         return layers
-        
+
     def forward(self, x):
         return self.model(x)
 
+
 class SSIMLoss(nn.Module):
     def __init__(self, window_size=11, size_average=True):
         super(SSIMLoss, self).__init__()
         self.window_size = window_size
         self.size_average = size_average
         self.channel = 1
         self.create_window(window_size)
-        
+
     def create_window(self, window_size, channel=1):
         # Create a Gaussian window (filter) with specified size and channel
         sigma = 1.5
@@ -131,14 +135,14 @@ def forward(self, img1, img2):
         if img1.size(1) != self.channel:
             self.channel = img1.size(1)
             self.create_window(self.window_size, self.channel)
-        
+
         # Move window to the same device as the images
         window = self.window.to(img1.device)
-        
+
         # Compute SSIM components
         mu1 = F.conv2d(img1, window, padding=self.window_size//2, groups=self.channel)
         mu2 = F.conv2d(img2, window, padding=self.window_size//2, groups=self.channel)
-        
+
         mu1_sq = mu1.pow(2)
         mu2_sq = mu2.pow(2)
         mu1_mu2 = mu1 * mu2
@@ -153,7 +157,7 @@ def forward(self, img1, img2):
 
         # SSIM calculation
         ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2))
-        
+
         if self.size_average:
             return torch.clamp((1 - ssim_map.mean()) / 2, 0, 1)  # SSIM loss as (1 - SSIM) / 2
         else:

spaghetti/cli_inference.py

@@ -1,19 +1,20 @@
-'''
+"""
 Entry point for the CLI inference of SPAGHETTI.
-'''
+"""
 import os
 import argparse
 from spaghetti import inferences
 from PIL import Image
 
+
 def inference(input, output, checkpoint):
-    '''
+    """
     The inference function for the CLI inference
     args:
         input: str, the input image or directory to translate
         output: str, the output directory to save the translated image(s)
         checkpoint: str, the path to the checkpoint file
-    '''
+    """
     # check if input is a directory
     if os.path.isdir(input):
         # get all images
@@ -34,6 +35,7 @@ def inference(input, output, checkpoint):
     processed_imgs = model.pre_processing(pil_imgs, transform="default")
     model.inference(processed_imgs, names, output)
 
+
 def main():
     parser = argparse.ArgumentParser(description="CLI for translating PCM images using SPAGHETTI")
     parser.add_argument("--input", '-i', type=str, help="The input image or directory to translate")
@@ -47,5 +49,6 @@ def main():
     inference(args.input, args.output, args.checkpoint)
     print("Inference Completed. Images saved to ", args.output)
 
+
 if __name__ == "__main__":
-    main()
+    main()

spaghetti/dataset.py

@@ -1,28 +1,29 @@
-'''
+"""
 Prepare the datasets for training and inference
-'''
+"""
 import random
 from PIL import Image
 import os
 from torch.utils.data import Dataset
 
+
 class TrainingDataset(Dataset):
-    '''
+    """
     The dataset class for the SPAGHETTI model training
     args:
         path_1: list of strings, the paths to the images in domain 1
         path_2: list of strings, the paths to the images in domain 2
         transform_1: the transformation for domain 1 images
         transform_2: the transformation for domain 2 images
         num_sample: int, optional, the number of images to sample from each domain
-    '''
-    def __init__(self, path_1: list[str], path_2: list[str], 
+    """
+    def __init__(self, path_1: list[str], path_2: list[str],
                  transform_1, transform_2, num_sample=None):
         random.seed(42)
         # domain 1
         domain1_paths = []
         for each in path_1:
-            domain1_paths.extend([os.path.join(each, x) for x in os.listdir(each) 
+            domain1_paths.extend([os.path.join(each, x) for x in os.listdir(each)
                                   if x.endswith((".png", ".jpg", ".jpeg", ".tiff", ".tif"))])
         if num_sample:
             try:
@@ -31,7 +32,7 @@ def __init__(self, path_1: list[str], path_2: list[str],
                 self.domain1_images = random.choices(domain1_paths, k=num_sample)
         else:
             self.domain1_images = domain1_paths
-        
+
         # domain 2
         domain2_paths = []
         for each in path_2:
@@ -60,9 +61,9 @@ def __getitem__(self, index):
         domain1_img_path = self.domain1_images[index % self.domain1_len]
         domain2_img_path = self.domain2_images[index % self.domain2_len]
         domain1_img = Image.open(domain1_img_path).convert("RGB")
-        domain2_img = Image.open(domain2_img_path).convert("RGB") 
+        domain2_img = Image.open(domain2_img_path).convert("RGB")
 
         domain1_img = self.transform_1(domain1_img)
         domain2_img = self.transform_2(domain2_img)
 
-        return domain1_img, domain2_img
+        return domain1_img, domain2_img