s4583222 - Stable Diffusion model (Wrong branch fix)

recognition/s4583222_Solution/dataset.py

@@ -0,0 +1,48 @@
+import os
+import torch
+import torchvision
+from PIL import Image
+from torch.utils.data import Dataset
+import numpy as np
+from torch.utils.data import DataLoader
+
+training_images = '/home/Student/s4583222/COMP3710/Images/Train' #Directory of the training brain images
+IMG_SIZE = 128 #Size to rescale the images to
+
+class BrainDataset(Dataset):
+    """
+    Class to preprocess the Oasis Brain Dataset.
+    We do not want to include any of the brain masks
+    """
+    def __init__(self, image_directory, transform = None):
+        self.image_directory = image_directory
+        self.transform = transform
+        self.images = os.listdir(image_directory)
+
+    def __len__(self):
+        return len(self.images)
+
+    def __getitem__(self, index):
+        image_path = os.path.join(self.image_directory, self.images[index])
+        image = Image.open(image_path).convert("RGB")#Open image as RGB
+
+        if self.transform is not None:
+            image = self.transform(image)#Perform transforms on the image
+
+        return image
+
+def get_data_loaders(training_images, batch, workers, pin_mem):
+    """
+    Given the training image path, batch size, number of workers, and pin memory status.
+    This function creates the transforms for each image, transforms the image, and then passes it to DataLoader.
+    The function then returns the data loader for the training images.
+    """
+    transforms = torchvision.transforms.Compose([
+        torchvision.transforms.Resize((IMG_SIZE,IMG_SIZE)),#Resize the images
+        torchvision.transforms.RandomHorizontalFlip(),#Randomly flip image
+        torchvision.transforms.ToTensor(),#Convert to a tensor
+        torchvision.transforms.Lambda(lambda t: (t * 2) - 1)#Normalise
+    ])
+    training_dataset = BrainDataset(image_directory=training_images, transform=transforms)#Transform image
+    training_loader = DataLoader(training_dataset, batch_size=batch, num_workers=workers, pin_memory=pin_mem, drop_last=True)#Pass transformed images to DataLoader
+    return training_loader #Return the data loader for the training images.

recognition/s4583222_Solution/modules.py

@@ -0,0 +1,135 @@
+import torch
+from torch import nn
+from tqdm import tqdm
+import torch.nn.functional as F
+import math
+
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+IMG_SIZE = 256
+
+
+class PositionEmbedding(nn.Module):
+    """
+    Takes a tensor of size (batch size, 1) and transforms it into a tensor of size (batch size, time_emb)
+    Acknowledgment From:   
+    https://colab.research.google.com/github/huggingface/notebooks/blob/main/examples/annotated_diffusion.ipynb#scrollTo=592aa765
+    """
+    def __init__(self, time_emb):
+        super().__init__()
+        self.time_emb = time_emb
+        self.relu = nn.ReLU()
+        self.linear = nn.Linear(time_emb, time_emb)
+
+    def forward(self, time):
+        # time = time.to(DEVICE)
+        device = time.device
+        embed = math.log(10000) / ((self.time_emb // 2) - 1)
+        embed = torch.exp(torch.arange((self.time_emb // 2), device=device) * - embed)
+        embed = time[:, None] * embed[None, :]
+        position_encode = torch.cat((embed.sin(), embed.cos()), dim = -1)
+        position_encode = self.linear(position_encode)
+        position_encode = self.relu(position_encode)
+        return position_encode
+
+class Block(nn.Module):
+    """
+    This is a residual block of the UNET. Both Contraction and Expansion path uses this block.
+    """
+    def __init__(self, in_chan, out_chan, time_emb):
+        super().__init__()    
+        self.relu = nn.ReLU()
+        self.time = nn.Linear(time_emb, out_chan)
+        self.conv_1 = nn.Conv2d(in_chan, out_chan, kernel_size=3, padding=1)  
+        self.conv_2 = nn.Conv2d(out_chan, out_chan, kernel_size=3, padding=1)
+        self.bnorm1 = nn.BatchNorm2d(out_chan)
+        self.bnorm2 = nn.BatchNorm2d(out_chan)               
+
+    def forward(self, x, t):
+        out = self.conv_1(x)
+        out = self.bnorm1(out)
+        out = self.relu(out)
+        time_embedding = self.time(t)
+        time_embedding = self.relu(time_embedding)
+        time_embedding = time_embedding[(...,) + (None,) * 2]
+        out = out + time_embedding
+        out = self.conv_2(out)
+        out = self.bnorm2(out)
+        out = self.relu(out)
+        return out
+
+
+class UNETModel(nn.Module):
+    """
+    Modified UNET model with Sinusoidal Position Embedding
+    """
+    def __init__(self, in_chan = 3, out_chan = 3, device = DEVICE):
+        super().__init__()
+        self.device = device
+        time_dimension = 32
+        #Network 64 128 256 512 1024
+
+        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
+        self.time_embedding = PositionEmbedding(time_dimension)
+
+        #Contraction Path
+        self.first_level = Block(in_chan, 64, time_dimension)
+        self.second_level = Block(64, 128, time_dimension)
+        self.third_level = Block(128, 256, time_dimension)
+        self.fourth_level = Block(256, 512, time_dimension)
+
+        #Bottle neck (floor)
+        self.bottle_neck = Block(512, 1024, time_dimension)
+
+        #Expansion path
+        self.fifth_level_up = nn.ConvTranspose2d(1024, 512, kernel_size=2, stride=2)       
+        self.fifth_level = Block(1024, 512, time_dimension)
+        self.sixth_level_up = nn.ConvTranspose2d(512, 256, kernel_size=2, stride=2)
+        self.sixth_level = Block(512, 256, time_dimension)
+        self.seventh_level_up = nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2)
+        self.seventh_level = Block(256, 128, time_dimension)
+        self.eigth_level_up = nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2)
+        self.eigth_level = Block(128, 64, time_dimension)
+
+        #Final Conv
+        self.last_conv = nn.Conv2d(64, out_chan, kernel_size=1)
+
+    def forward(self, x, t):
+        t = self.time_embedding(t)
+
+        x1 = self.first_level(x, t)#First residual block
+        x1_skip = x1#Skip connection
+
+        x2 = self.pool(x1)#Maxpool
+        x2 = self.second_level(x2, t)#Second residual block
+        x2_skip = x2
+
+        x3 = self.pool(x2)
+        x3 = self.third_level(x3, t)#Third residual block
+        x3_skip = x3
+
+        x4 = self.pool(x3)
+        x4 = self.fourth_level(x4, t)#Fourth residual block
+        x4_skip = x4
+
+        x_bottle = self.pool(x4)
+        x_bottle = self.bottle_neck(x_bottle, t)#Double conv with time embedding
+
+        x5 = self.fifth_level_up(x_bottle)#Upsample        
+        x5 = torch.cat([x4_skip, x5], dim = 1)#Concat with skip connection
+        x5 = self.fifth_level(x5, t)#Fifth residual block
+
+        x6 = self.sixth_level_up(x5) 
+        x6 = torch.cat([x3_skip, x6], dim = 1)
+        x6 = self.sixth_level(x6, t)#Sixth residual block
+
+        x7 = self.seventh_level_up(x6) 
+        x7 = torch.cat([x2_skip, x7], dim = 1)
+        x7 = self.seventh_level(x7, t)#Seventh residual block
+
+        x8 = self.eigth_level_up(x7) 
+        x8 = torch.cat([x1_skip, x8], dim = 1)
+        x8 = self.eigth_level(x8, t)#Eigth residual block 
+
+        output_x = self.last_conv(x8)#Final convolution, output image is the same dimensions as input
+
+        return output_x 

recognition/s4583222_Solution/predict.py

@@ -0,0 +1,119 @@
+import os
+import torch
+import torchvision
+import torch.nn as nn
+from matplotlib import pyplot as plt
+from tqdm import tqdm
+from torch import optim
+from torchvision import transforms 
+from modules import UNETModel
+import numpy as np
+from PIL import Image
+import torch.nn.functional as F
+
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+BATCH_SIZE = 64 #If we drop this to 1, then len(training_loader) = 11,328
+IMG_HEIGHT = 128 #Actual is 256
+IMG_WIDTH = 128 #Actual is 256
+TRAINING_DIR = '/home/Student/s4583222/COMP3710/Images/Train'
+
+NOISE_STEP = 600 #How many timesteps we want to destroy an image
+
+BETA = torch.linspace(0.0001, 0.02, NOISE_STEP) #Linear beta schedule
+ALPHA = 1. - BETA
+ALPHA_HAT = torch.cumprod(ALPHA, axis=0) #alpha hat is the cumulative product of alpha
+ALPHA_HAT_PREV = F.pad(ALPHA_HAT[:-1], (1,0), value=1.0) #Previous alpha hat
+SQRT_ALPHA_REC = torch.sqrt(1.0/ALPHA) #Square root of the reciprical of alpha
+PROS_VAR = BETA * (1. - ALPHA_HAT_PREV)/(1. - ALPHA_HAT) #Calculation for the posterior
+
+def index_list(vals, t, x_shape):
+    """
+    Returns the time index for a batch
+    Acknowledgment From:   
+    https://colab.research.google.com/github/huggingface/notebooks/blob/main/examples/annotated_diffusion.ipynb#scrollTo=592aa765
+    """
+    batch_size = t.shape[0]
+    out = vals.gather(-1,t.cpu())
+    out = out.reshape(batch_size, *((1,)*(len(x_shape)-1))).to(t.device)
+    return out
+
+def show_image(image, epoch, i, disc = "Blank"):
+    """
+    Takes in an image tensor and converts it to an image. The image will be coloured so it will save it and open the image again to
+    convert it to a grayscale image. This image is them saved
+    Acknowledgment From:   
+    https://colab.research.google.com/drive/1sjy9odlSSy0RBVgMTgP7s99NXsqglsUL?usp=sharing#scrollTo=Rj17psVw7Shg
+    """    
+    reverse_transforms = transforms.Compose([
+        transforms.Lambda(lambda t: (t + 1) / 2), #Renormalise tensor
+        transforms.Lambda(lambda t: t.permute(1, 2, 0)), #Move chanels to the back
+        transforms.Lambda(lambda t: t * 255.),#Get pixels back to colour
+        transforms.Lambda(lambda t: t.numpy().astype(np.uint8)),#Convert to numpy array
+        transforms.ToPILImage(),#Convert to image
+    ])
+
+    if len(image.shape) == 4:# Take first image of batch
+        image = image[0, :, :, :] 
+
+    plt.imshow(reverse_transforms(image))
+    plt.savefig(f"{epoch}_{i}_{disc}_colour") #Saves coloured brain image
+    im = Image.open(f"{epoch}_{i}_{disc}_colour.png").convert("L") #Opens coloured brain image as gray scale
+    plt.imshow(im, cmap='gray')
+    plt.savefig(f"{epoch}_{i}_{disc}_gray")#Saves grayscale image
+
+
+
+
+@torch.no_grad() #Required so the weights of previous images does not effect the image we are generating
+def sample_image_at_timestep(rand_noise, t, model):
+    """
+    Returns the time index for a batch
+    Acknowledgment From:   
+    https://colab.research.google.com/github/huggingface/notebooks/blob/main/examples/annotated_diffusion.ipynb#scrollTo=592aa765
+    """    
+    beta_time = index_list(BETA, t, rand_noise.shape)
+    sqrt_one_minus_alpha_hat_time = index_list(torch.sqrt(1. - ALPHA_HAT), t, rand_noise.shape)
+    sqrt_alpha_reciprical_time = index_list(SQRT_ALPHA_REC, t, rand_noise.shape)
+    prosterior_var_time = index_list(PROS_VAR, t, rand_noise.shape)
+
+    predicted_noise = model(rand_noise, t)
+    mean = sqrt_alpha_reciprical_time * (rand_noise - beta_time * predicted_noise / sqrt_one_minus_alpha_hat_time)
+    if t == 0:
+        return mean
+    else:
+        noise = torch.randn_like(rand_noise)
+        return mean + torch.sqrt(prosterior_var_time) * noise  
+
+@torch.no_grad() #Required so the weights of previous images does not effect the image we are generating
+def sample_image(epoch, i, model):
+    """
+    This function is looped through and it loops through timesteps, slowly denoising an image of random noise.
+    Acknowledgment From:   
+    https://colab.research.google.com/drive/1sjy9odlSSy0RBVgMTgP7s99NXsqglsUL?usp=sharing#scrollTo=Rj17psVw7Shg
+    """
+    noise = torch.randn((1,3, IMG_HEIGHT, IMG_WIDTH), device=DEVICE)#Random noise   
+    plt.figure(figsize=(15,15))
+    plt.axis('off')
+    for ind in range(0, NOISE_STEP)[::-1]:#Reversed from 600 to 0
+        t = torch.full((1,), ind, device=DEVICE, dtype=torch.long)
+        noise = sample_image_at_timestep(noise, t, model)#Denoised image
+    show_image(noise.detach().cpu(), epoch, i, f"{ind}_backward")#After looped through all timesteps show image    
+
+def predict():
+    """
+    This function loads a previously saved model and then loops a function gen_images number of times.
+    It will product gen_images number of images.
+    Prints Done! to .out when finished
+    """
+    model = UNETModel().to(DEVICE)
+    model.load_state_dict(torch.load("Model_2.pt"))
+    gen_images = 1 #Generate x number images
+    images_generated = 0
+    iteration = 0
+    while images_generated < gen_images:
+        sample_image(images_generated + 1, iteration, model)
+        images_generated = images_generated + 1
+    print("Done!")
+
+if __name__ == '__main__':
+    predict()