upload 3 files

src/dataset.py

@@ -0,0 +1,65 @@
+from PIL import Image
+from torch.utils.data import Dataset
+from torchvision import transforms
+import os 
+
+class ClassifierDataset(Dataset):
+    def __init__(self, root_dir, transform=None):
+        self.root_dir = root_dir
+        self.transform = transform
+
+        self.classes = ['0', '1']
+        self.class_to_idx = {cls: idx for idx, cls in enumerate(self.classes)}
+
+        self.samples = self._make_dataset()
+
+    def _make_dataset(self):
+        samples = []
+        for class_name in self.classes:
+            class_dir = os.path.join(self.root_dir, class_name)
+            for img_name in os.listdir(class_dir):
+                img_path = os.path.join(class_dir, img_name)
+                samples.append((img_path, self.class_to_idx[class_name]))
+        return samples
+
+    def __len__(self):
+        return len(self.samples)
+
+    def __getitem__(self, idx):
+        img_path, label = self.samples[idx]
+        img = Image.open(img_path).convert('L')  # Convert to grayscale
+        if self.transform:
+            img = self.transform(img)
+        return img, label
+    
+
+
+class CustomDataset(Dataset):
+    def __init__(self, root_dir, train_N, train_P, img_res):
+        self.root_dir = root_dir
+        self.train_N = train_N
+        self.train_P = train_P
+        self.img_res = img_res
+        self.transforms = transforms.Compose([
+            transforms.Resize(img_res),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.5], std=[0.5])  # Assuming grayscale images
+        ])
+
+    def __len__(self):
+        return min(len(os.listdir(os.path.join(self.root_dir, self.train_N))),
+                   len(os.listdir(os.path.join(self.root_dir, self.train_P))))
+
+    def __getitem__(self, idx):
+        normal_path = os.path.join(self.root_dir, self.train_N, os.listdir(os.path.join(self.root_dir, self.train_N))[idx])
+        pneumo_path = os.path.join(self.root_dir, self.train_P, os.listdir(os.path.join(self.root_dir, self.train_P))[idx])
+        
+        normal_img = Image.open(normal_path).convert("L")  # Load as grayscale
+        pneumo_img = Image.open(pneumo_path).convert("L")  # Load as grayscale
+        
+        normal_img = self.transforms(normal_img)
+        pneumo_img = self.transforms(pneumo_img)
+        
+        return normal_img, pneumo_img
+
+

src/generate_images.py

@@ -0,0 +1,96 @@
+import os
+import torch
+from PIL import Image
+import numpy as np
+from torch.utils.data import DataLoader, Dataset
+from torchvision import transforms
+from tqdm import tqdm
+
+from src.models import ResUNetGenerator
+
+# Custom Dataset
+class ImageDataset(Dataset):
+    def __init__(self, image_paths, transform=None):
+        self.image_paths = image_paths
+        self.transform = transform
+
+    def __len__(self):
+        return len(self.image_paths)
+
+    def __getitem__(self, idx):
+        img_path = self.image_paths[idx]
+        image = Image.open(img_path).convert('L')
+        if self.transform:
+            image = self.transform(image)
+        return image, img_path
+
+# Function to save image
+def save_image(tensor, path):
+    if tensor.is_cuda:
+        tensor = tensor.cpu()
+
+    array = tensor.permute(1, 2, 0).detach().numpy()
+    array = (array * 0.5 + 0.5) * 255
+    array = array.astype(np.uint8)
+    if array.shape[2] == 1:
+        array = array.squeeze(2)
+        image = Image.fromarray(array, mode='L')
+    else:
+        image = Image.fromarray(array)
+    image.save(path)
+
+# Function to load model
+def load_model(checkpoint_path, model_class, device):
+    model = model_class().to(device)
+    model.load_state_dict(torch.load(checkpoint_path, map_location=torch.device('cpu')))
+    model.eval()
+    return model
+
+def generate_images(image_folder, g_NP_checkpoint, g_PN_checkpoint, output_dir='data/translated_images', batch_size=16):
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    
+    # Load models
+    g_NP = load_model(g_NP_checkpoint, lambda: ResUNetGenerator(gf=32, channels=1), device)
+    g_PN = load_model(g_PN_checkpoint, lambda: ResUNetGenerator(gf=32, channels=1), device)
+    
+    # Create output directories
+    os.makedirs(os.path.join(output_dir, '0'), exist_ok=True)
+    os.makedirs(os.path.join(output_dir, '1'), exist_ok=True)
+
+    # Collect image paths
+    image_paths_0 = [os.path.join(image_folder, '0', fname) for fname in os.listdir(os.path.join(image_folder, '0')) if fname.endswith(('.png', '.jpg', '.jpeg'))]
+    image_paths_1 = [os.path.join(image_folder, '1', fname) for fname in os.listdir(os.path.join(image_folder, '1')) if fname.endswith(('.png', '.jpg', '.jpeg'))]
+
+    # Prepare dataset and dataloader
+    transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize(mean=[0.485], std=[0.229])])
+    dataset_0 = ImageDataset(image_paths_0, transform)
+    dataset_1 = ImageDataset(image_paths_1, transform)
+    dataloader_0 = DataLoader(dataset_0, batch_size=batch_size, shuffle=False)
+    dataloader_1 = DataLoader(dataset_1, batch_size=batch_size, shuffle=False)
+    
+    # Process images from negative (0) to positive (1)
+    with torch.no_grad():
+        for batch, paths in tqdm(dataloader_0, desc="Converting N to P: "):
+            batch = batch.to(device)
+            translated_images = g_NP(batch)
+            translated_images = g_PN(translated_images)
+            for img, path in zip(translated_images, paths):
+                save_path = os.path.join(output_dir, '1', os.path.basename(path))
+                save_image(img, save_path)
+        
+        # Process images from positive (1) to negative (0)
+        for batch, paths in tqdm(dataloader_1, desc="Converting P to N: "):
+            batch = batch.to(device)
+            translated_images = g_PN(batch)
+            translated_images = g_NP(translated_images)
+            for img, path in zip(translated_images, paths):
+                save_path = os.path.join(output_dir, '0', os.path.basename(path))
+                save_image(img, save_path)
+
+if __name__ == '__main__':
+    image_folder = r'data\rsna-pneumonia-dataset\train'
+    g_NP_checkpoint = 'models\g_NP_best.ckpt'
+    g_PN_checkpoint = 'models\g_PN_best.ckpt'
+    
+    
+    generate_images(image_folder, g_NP_checkpoint, g_PN_checkpoint)

src/train.py

@@ -0,0 +1,124 @@
+from torchvision import transforms
+from torch.utils.data import DataLoader
+from lightning.pytorch.loggers.wandb import WandbLogger
+from lightning.pytorch.callbacks import ModelCheckpoint
+import lightning as pl
+import wandb
+
+from src.dataset import ClassifierDataset, CustomDataset
+from src.classifier import Classifier
+from src.models import CycleGAN
+from src.config import CFG
+
+def train_classifier(image_size,
+                     batch_size,
+                     epochs,
+                     resume_ckpt_path,
+                     train_dir,
+                     val_dir,
+                     checkpoint_dir,
+                     project,
+                     job_name):
+    
+    clf_wandb_logger = WandbLogger(project=project, name=job_name, log_model="all")
+
+    transform = transforms.Compose([
+        transforms.Resize((image_size, image_size)),  # Resize image to 512x512
+        transforms.ToTensor(),          
+        transforms.Normalize(mean=[0.485], std=[0.229])  # Normalize image
+    ])
+
+    # Define dataset paths
+    # train_dir = "/kaggle/working/CycleGan-CFE/train-data/train"
+    # val_dir = "/kaggle/working/CycleGan-CFE/train-data/val"
+
+    # Create datasets
+    train_dataset = ClassifierDataset(root_dir=train_dir, transform=transform)
+    val_dataset = ClassifierDataset(root_dir=val_dir, transform=transform)
+    print("Total Training Images: ",len(train_dataset))
+    print("Total Validation Images: ",len(val_dataset))
+
+    # Create data loaders
+    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, pin_memory=True, num_workers=4)
+    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, pin_memory=True, num_workers=4)
+    # Instantiate the classifier model
+    clf = Classifier(transfer=True)
+
+    checkpoint_callback = ModelCheckpoint(
+        monitor='val_loss',
+        dirpath=checkpoint_dir,
+        filename='efficientnet_b2-epoch{epoch:02d}-val_loss{val_loss:.2f}',
+        auto_insert_metric_name=False,
+        save_weights_only=False,
+        save_top_k=3,
+        mode='min'
+    )
+    # Set up PyTorch Lightning Trainer with multiple GPUs and tqdm progress bar
+    trainer = pl.Trainer(
+        devices="auto",
+        precision="16-mixed",
+        accelerator="auto",
+        max_epochs=epochs,
+        accumulate_grad_batches=10,
+        log_every_n_steps=1,
+        check_val_every_n_epoch=1,
+        benchmark=True,
+        logger=clf_wandb_logger,
+        callbacks=[checkpoint_callback],
+    )
+
+    # Train the classifier
+    trainer.fit(clf, train_loader, val_loader, ckpt_path=resume_ckpt_path)
+    wandb.finish()
+
+
+def train_cyclegan(image_size,
+                   batch_size,
+                   epochs,
+                   classifier_path,
+                   resume_ckpt_path,
+                   train_dir,
+                   val_dir,
+                   test_dir,
+                   checkpoint_dir,
+                   project,
+                   job_name,
+):
+
+
+    testdata_dir = test_dir
+    train_N = "0"
+    train_P = "1"
+    img_res = (image_size, image_size)
+
+    test_dataset = CustomDataset(root_dir=testdata_dir, train_N=train_N, train_P=train_P, img_res=img_res)
+    test_dataloader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=4, pin_memory=True)
+
+    wandb_logger = WandbLogger(project=project, name=job_name, log_model="all")
+    print(classifier_path)
+    cyclegan = CycleGAN(train_dir=train_dir, val_dir=val_dir, test_dataloader=test_dataloader, classifier_path=classifier_path, checkpoint_dir=checkpoint_dir, gf=CFG.GAN_FILTERS, df=CFG.DIS_FILTERS)
+
+    gan_checkpoint_callback = ModelCheckpoint(dirpath=checkpoint_dir,
+                                        filename='cyclegan-epoch_{epoch}-vloss_{val_generator_loss:.2f}',
+                                        monitor='val_generator_loss',
+                                        save_top_k=3,
+                                        save_last=True,
+                                        save_weights_only=False,
+                                        verbose=True,
+                                        mode='min')
+
+
+    # Create the trainer
+    trainer = pl.Trainer(
+        accelerator="auto",
+        precision="16-mixed",
+        max_epochs=epochs,
+        log_every_n_steps=1,
+        benchmark=True,
+        devices="auto",
+        logger=wandb_logger,
+        callbacks= [gan_checkpoint_callback]
+    )
+
+    # Train the CycleGAN model
+    trainer.fit(cyclegan, ckpt_path=resume_ckpt_path)