delete files

src/generate_images.py

@@ -1,96 +0,0 @@
-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)