dcgan.py

import torch as t
from torch import nn
from torch.utils.data import Dataset, DataLoader
from torch.optim import Adam

import numpy as np
import cv2
import os

'''
Deep convolutional generative adversarial network to produce artificial bubble images to imitate
32x32 pixel preprocessed images of bubbles in PICO-60.  The training images can be produces with
PICO60_image_analysis.py.
'''

batch_size = 16
img_size = 32
latent_dim = 20  # Dimension of the latent space
channels = 1
lr = 0.00008 #0.0001
b1, b2 = 0.7, 0.999  # Beta parameters for Adam optimizer
n_epochs = 60
n_ch = 128  # Number of channels in initial convolution layers
sample_interval = 100  # Save a generated image every sample_interval number of batches

save_progress_samples = False  # Save images every sample_interval number of batches to check progress
write_dir = './dcgan_gen_images/'  # Directory to save images to if save_progress_samples is True
data_dir = './bub_single_24/'  # Directory containing training images
model_path = './dcgan_model.pt'  # File to save the model to once training is complete

if save_progress_samples:
    os.makedirs(write_dir, exist_ok=True)

class Data(Dataset):
    def __init__(self):
        self.data = [cv2.imread(data_dir + im, cv2.IMREAD_GRAYSCALE)//255 for im in os.listdir(data_dir) if '-0.png' in im]
        self.data = t.Tensor([im for im in self.data if im.shape == (img_size, img_size)])
        
        self.len = self.data.shape[0]

    def __len__(self):
        return self.len

    def __getitem__(self, index):
        return self.data[index]

class Generator(nn.Module):
    # Takes as input a random tensor of size (batch_size, latent_dim) generated by noise()
    def __init__(self):
        super(Generator, self).__init__()
        
        self.init_size = img_size//4
        
        self.l1 = nn.Sequential(nn.Linear(latent_dim, n_ch * self.init_size**2))
        
        self.conv_blocks = nn.Sequential(
            nn.BatchNorm2d(n_ch),
            nn.Upsample(scale_factor=2),
            nn.Conv2d(n_ch, n_ch, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(n_ch, 0.8),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Upsample(scale_factor=2),
            nn.Conv2d(n_ch, n_ch//2, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(n_ch//2, 0.8),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(n_ch//2, channels, kernel_size=3, stride=1, padding=1),
            nn.Sigmoid()
            #nn.Tanh()
            )
            
    def forward(self, z):
        out = self.l1(z)
        out = out.view(out.shape[0], n_ch, self.init_size, self.init_size)
        img = self.conv_blocks(out)
        return img
            

class Discriminator(nn.Module):
    # Takes as input a data image or image generated by Generator
    def __init__(self):
        super(Discriminator, self).__init__()

        def discriminator_block(in_filters, out_filters, bn=True):
            block = [nn.Conv2d(in_filters, out_filters, 3, 2, 1), nn.LeakyReLU(0.2, inplace=True), nn.Dropout2d(0.25)]
            if bn:
                block.append(nn.BatchNorm2d(out_filters, 0.8))
            return block

        self.model = nn.Sequential(
            *discriminator_block(channels, n_ch//8, bn=False),
            *discriminator_block(n_ch//8, n_ch//4),
            *discriminator_block(n_ch//4, n_ch//2),
            *discriminator_block(n_ch//2, n_ch),
        )

        # The height and width of downsampled image
        ds_size = img_size // 2 ** 4
        self.adv_layer = nn.Sequential(nn.Linear(n_ch * ds_size ** 2, 1), nn.Sigmoid())

    def forward(self, img):
        out = self.model(img)
        out = out.view(out.shape[0], -1)
        validity = self.adv_layer(out)
        return validity
        

def generate_noise(name):
    if name == 'Gaussian':
        return t.normal(mean=0, std=1, size=(batch_size, latent_dim))
    return t.rand((batch_size, latent_dim))
    
def save_model():
    print('Saving Saving generator model to ' + model_path)
    t.save(G, model_path)


def main():

    dataset = Data()

    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

    # Loss function
    adversarial_loss = nn.BCELoss()

    G = Generator()
    D = Discriminator()

    # Optimizers
    optimizer_G = Adam(G.parameters(), lr=lr, betas=(b1, b2))
    optimizer_D = Adam(D.parameters(), lr=lr, betas=(b1, b2))
    try:
        for epoch in range(n_epochs):
            for i, data_imgs in enumerate(dataloader):
                
                # Configure generator input
                data_shape = data_imgs.shape
                data_imgs = data_imgs.view(data_imgs.shape[0], channels, data_imgs.shape[1], data_imgs.shape[2])
                
                # Adversarial ground truths
                valid = t.ones((data_shape[0], 1), requires_grad=False)
                fake = t.zeros((data_shape[0], 1), requires_grad=False)
                
                # ----------------
                # Train Generator
                # ----------------
                
                optimizer_G.zero_grad()
                
                # Sample noise as generator input
                noise = generate_noise('Gaussian')
                
                # Generate a batch of images
                gen_imgs = G(noise)
                gen_imgs = gen_imgs[:data_shape[0]]
                
                # Loss measures generator's ability to fool the discriminator
                G_loss = adversarial_loss(D(gen_imgs), valid)

                G_loss.backward()
                optimizer_G.step()
                
                # --------------------
                # Train Discriminator
                # --------------------
                
                optimizer_D.zero_grad()
                
                # Measure discriminator's ability to classify real from generated samples
                real_loss = adversarial_loss(D(data_imgs), valid)
                fake_loss = adversarial_loss(D(gen_imgs.detach()), fake)
                D_loss = (real_loss + fake_loss) / 2
                
                D_loss.backward()
                optimizer_D.step()
                
                # ----------------
                # Log Progress
                # ----------------
                
                print(
                    "[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f]"
                    % (epoch, n_epochs, i, len(dataloader), D_loss.item(), G_loss.item())
                )
                if save_progress_samples:
                    batches_done = epoch*len(dataloader) + i
                    if batches_done % sample_interval == 0:
                        os.makedirs(write_dir + str(batches_done), exist_ok=True)
                        for j, gen_im in enumerate(gen_imgs):
                            cv2.imwrite('%s%d/%d.png' % (write_dir, batches_done, j), gen_im[0].detach().numpy()*255)
        
        save_model()
                
    except KeyboardInterrupt:
        save_model()


if __name__ == '__main__':
    main()