
import os, cv2, torch
import numpy as np
import matplotlib.pyplot as plt
import torchvision.transforms as tt
import torch.nn as nn
import torch.nn.functional as F
from tqdm import tqdm
from torchsummary import summary
from torch.utils.data import DataLoader
from torchvision.datasets import ImageFolder
from torchvision.utils import save_image, make_grid


# CONSTANTS
IMAGE_SIZE = 64
BATCH_SIZE = 128
MEAN, STD = (0.5, 0.5, 0.5), (0.5, 0.5, 0.5)
DATA_DIR = './data/'


train_ds = ImageFolder(DATA_DIR, transform = tt.Compose([tt.Resize(IMAGE_SIZE),
                                                         tt.CenterCrop(IMAGE_SIZE),
                                                         tt.ToTensor(),
                                                         tt.Normalize(mean=MEAN,
                                                                      std=STD)]))


train_dl = DataLoader(train_ds, BATCH_SIZE, shuffle = True, 
                      num_workers = 2, pin_memory = True)


def denorm(img_tensors):
    return img_tensors * MEAN[0] + STD[0]


def show_images(images, nmax=64):
    fig, ax = plt.subplots(figsize = (8, 8))
    ax.set_xticks([]); ax.set_yticks([])
    ax.imshow(make_grid(denorm(images.detach()[:nmax]), nrow=8).permute(1,2,0))


def show_batch(dl, nmax=64):
    for images, _ in dl:
        show_images(images, nmax)
        break


show_batch(train_dl)


def to_device(data, device):
    if isinstance(data, (list, tuple)):
        return [to_device(x, device) for x in data]
    return data.to(device, non_blocking = True)


class DeviceDataLoader():
    def __init__(self, dl, device):
        self.dl = dl
        self.device = device
    def __iter__(self):
        for b in self.dl:
            yield to_device(b, self.device)
    def __len__(self):
        return len(self.dl)


device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')


train_dl = DeviceDataLoader(train_dl, device)


# Define the discriminator architecture using nn.Sequential
discriminator = nn.Sequential(
    # First Convolutional Layer
    nn.Conv2d(3, 64, kernel_size=4, stride=2, padding=1, bias=False),
    nn.BatchNorm2d(64),
    nn.LeakyReLU(0.2, inplace=True),

    # Second Convolutional Layer
    nn.Conv2d(64, 128, kernel_size=4, stride=2, padding=1, bias=False),
    nn.BatchNorm2d(128),
    nn.LeakyReLU(0.2, inplace=True),

    # Third Convolutional Layer
    nn.Conv2d(128, 256, kernel_size=4, stride=2, padding=1, bias=False),
    nn.BatchNorm2d(256),
    nn.LeakyReLU(0.2, inplace=True),

    # Fourth Convolutional Layer
    nn.Conv2d(256, 512, kernel_size=4, stride=2, padding=1, bias=False),
    nn.BatchNorm2d(512),
    nn.LeakyReLU(0.2, inplace=True),

    # Fifth Convolutional Layer (Output Layer)
    nn.Conv2d(512, 1, kernel_size=4, stride=1, padding=0, bias=False),
    nn.Flatten(),
    nn.Sigmoid()
)


discriminator = to_device(discriminator, device)
summary(discriminator, (3, 64, 64))

----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Conv2d-1           [-1, 64, 32, 32]           3,072
       BatchNorm2d-2           [-1, 64, 32, 32]             128
         LeakyReLU-3           [-1, 64, 32, 32]               0
            Conv2d-4          [-1, 128, 16, 16]         131,072
       BatchNorm2d-5          [-1, 128, 16, 16]             256
         LeakyReLU-6          [-1, 128, 16, 16]               0
            Conv2d-7            [-1, 256, 8, 8]         524,288
       BatchNorm2d-8            [-1, 256, 8, 8]             512
         LeakyReLU-9            [-1, 256, 8, 8]               0
           Conv2d-10            [-1, 512, 4, 4]       2,097,152
      BatchNorm2d-11            [-1, 512, 4, 4]           1,024
        LeakyReLU-12            [-1, 512, 4, 4]               0
           Conv2d-13              [-1, 1, 1, 1]           8,192
          Flatten-14                    [-1, 1]               0
          Sigmoid-15                    [-1, 1]               0
================================================================
Total params: 2,765,696
Trainable params: 2,765,696
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.05
Forward/backward pass size (MB): 2.81
Params size (MB): 10.55
Estimated Total Size (MB): 13.41
----------------------------------------------------------------


LATENT_SIZE = 128


# Define the generator architecture using nn.Sequential
generator = nn.Sequential(
    # First Transposed Convolutional Layer
    nn.ConvTranspose2d(LATENT_SIZE, 512, kernel_size=4, stride=1, padding=0, bias=False),
    nn.BatchNorm2d(512),         # Batch Normalization for stabilization
    nn.ReLU(True),               # ReLU activation function for non-linearity

    # Second Transposed Convolutional Layer
    nn.ConvTranspose2d(512, 256, kernel_size=4, stride=2, padding=1, bias=False),
    nn.BatchNorm2d(256),
    nn.ReLU(True),

    # Third Transposed Convolutional Layer
    nn.ConvTranspose2d(256, 128, kernel_size=4, stride=2, padding=1, bias=False),
    nn.BatchNorm2d(128),
    nn.ReLU(True),

    # Fourth Transposed Convolutional Layer
    nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1, bias=False),
    nn.BatchNorm2d(64),
    nn.ReLU(True),

    # Fifth Transposed Convolutional Layer (Output Layer)
    nn.ConvTranspose2d(64, 3, kernel_size=4, stride=2, padding=1, bias=False),
    nn.Tanh()                    # Tanh activation for output normalization
)


generator = to_device(generator, device)
summary(generator, input_size=(LATENT_SIZE, 1, 1))

----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
   ConvTranspose2d-1            [-1, 512, 4, 4]       1,048,576
       BatchNorm2d-2            [-1, 512, 4, 4]           1,024
              ReLU-3            [-1, 512, 4, 4]               0
   ConvTranspose2d-4            [-1, 256, 8, 8]       2,097,152
       BatchNorm2d-5            [-1, 256, 8, 8]             512
              ReLU-6            [-1, 256, 8, 8]               0
   ConvTranspose2d-7          [-1, 128, 16, 16]         524,288
       BatchNorm2d-8          [-1, 128, 16, 16]             256
              ReLU-9          [-1, 128, 16, 16]               0
  ConvTranspose2d-10           [-1, 64, 32, 32]         131,072
      BatchNorm2d-11           [-1, 64, 32, 32]             128
             ReLU-12           [-1, 64, 32, 32]               0
  ConvTranspose2d-13            [-1, 3, 64, 64]           3,072
             Tanh-14            [-1, 3, 64, 64]               0
================================================================
Total params: 3,806,080
Trainable params: 3,806,080
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.00
Forward/backward pass size (MB): 3.00
Params size (MB): 14.52
Estimated Total Size (MB): 17.52
----------------------------------------------------------------


# Generate a batch of random latent vectors
xb = torch.randn(BATCH_SIZE, LATENT_SIZE, 1, 1)

# Generate fake images using the generator
fake_images = generator(xb)


# Print the shape of the generated fake images
print(fake_images.shape)

torch.Size([128, 3, 64, 64])


# Show the generated fake images
show_images(fake_images)


# Create a directory to save the generated images
sample_dir = 'generated'
os.makedirs(sample_dir, exist_ok=True)


# Generate a fixed set of random latent vectors for consistent visualization
fixed_latent = torch.randn(64, LATENT_SIZE, 1, 1, device=device)


# Define a function to save generated samples to the specified directory
def save_samples(index, latent_tensors, show=True):
    fake_images = generator(latent_tensors)
    fake_fname = f'generated-images-{index:0=4d}.png'
    
    # Save the generated images
    save_image(denorm(fake_images), os.path.join(sample_dir, fake_fname), nrow=8)
    print('Saving', fake_fname)
    
    # Optionally, display the saved images in a grid
    if show:
        fig, ax = plt.subplots(figsize=(8, 8))
        ax.set_xticks([]); ax.set_yticks([])
        ax.imshow(make_grid(fake_images.cpu().detach(), nrow=8).permute(1, 2, 0))


# Define a function to train a GAN model for a specified number of epochs
def fit(model, criterion, epochs, lr, start_idx=1):
    # Set discriminator and generator models to training mode
    model['discriminator'].train()
    model['generator'].train()

    # Clear GPU memory to avoid memory issues during training
    torch.cuda.empty_cache()

    # Initialize lists to store losses and scores
    losses_g = []
    losses_d = []
    real_scores = []
    fake_scores = []

    # Set up separate Adam optimizers for the discriminator and generator networks
    optimizer = {
        'discriminator': torch.optim.Adam(model['discriminator'].parameters(),
                                          lr=lr, betas=(0.5, 0.999)),
        'generator': torch.optim.Adam(model['generator'].parameters(),
                                      lr=lr, betas=(0.5, 0.999))
    }

    # Iterate over the specified number of epochs
    for epoch in range(epochs):
        # Initialize lists to store losses and scores for each batch in the epoch
        loss_d_per_epoch = []
        loss_g_per_epoch = []
        real_score_per_epoch = []
        fake_score_per_epoch = []

        # Iterate over batches in the training data loader
        for real_images, _ in tqdm(train_dl):
            # Zero out gradients for the discriminator optimizer
            optimizer['discriminator'].zero_grad()

            # Forward pass for real images through the discriminator
            real_preds = model['discriminator'](real_images)
            real_targets = torch.ones(real_images.size(0), 1, device=device)
            real_loss = criterion['discriminator'](real_preds, real_targets)
            cur_real_score = torch.mean(real_preds).item()

            # Generate fake images using the generator
            latent = torch.randn(BATCH_SIZE, LATENT_SIZE, 1, 1, device=device)
            fake_images = model['generator'](latent)

            # Forward pass for fake images through the discriminator
            fake_targets = torch.zeros(fake_images.size(0), 1, device=device)
            fake_preds = model['discriminator'](fake_images)
            fake_loss = criterion['discriminator'](fake_preds, fake_targets)
            cur_fake_score = torch.mean(fake_preds).item()

            # Append scores to lists
            real_score_per_epoch.append(cur_real_score)
            fake_score_per_epoch.append(cur_fake_score)

            # Calculate and backpropagate the total discriminator loss
            loss_d = real_loss + fake_loss
            loss_d.backward()
            optimizer['discriminator'].step()
            loss_d_per_epoch.append(loss_d.item())

            # Zero out gradients for the generator optimizer
            optimizer['generator'].zero_grad()

            # Generate new fake images and calculate generator loss
            latent = torch.randn(BATCH_SIZE, LATENT_SIZE, 1, 1, device=device)
            fake_images = model['generator'](latent)
            preds = model['discriminator'](fake_images)
            targets = torch.ones(BATCH_SIZE, 1, device=device)
            loss_g = criterion['generator'](preds, targets)

            # Backpropagate and optimize the generator's parameters
            loss_g.backward()
            optimizer['generator'].step()
            loss_g_per_epoch.append(loss_g.item())

        # Store the average losses and scores for the epoch
        losses_g.append(np.mean(loss_g_per_epoch))
        losses_d.append(np.mean(loss_d_per_epoch))
        real_scores.append(np.mean(real_score_per_epoch))
        fake_scores.append(np.mean(fake_score_per_epoch))

        # Log losses & scores for the last batch in each epoch
        print(f"Epoch [{epoch + 1}/{epochs}], "
              f"loss_g: {losses_g[-1]:.4f}, "
              f"loss_d: {losses_d[-1]:.4f}, "
              f"real_score: {real_scores[-1]:.4f}, "
              f"fake_score: {fake_scores[-1]:.4f}")

        # Save generated samples after the last epoch
        if epoch == epochs - 1:
            save_samples(epoch + start_idx, fixed_latent, show=False)

     # Save the final discriminator and generator models
    torch.save(model['discriminator'].state_dict(), 'discriminator.pth')
    torch.save(model['generator'].state_dict(), 'generator.pth')

    # Return the lists containing losses and scores for both the generator and discriminator
    return losses_g, losses_d, real_scores, fake_scores


model = {
    'discriminator': discriminator.to(device),
    'generator': generator.to(device)
}

criterion = {
    'discriminator': nn.BCELoss(),
    'generator': nn.BCELoss()
}

lr = 0.0002
epochs = 50


history = fit(model, criterion, epochs, lr)

100%|████████████████████████████████████████████████████████████████████████████████| 169/169 [11:48<00:00,  4.19s/it]

Epoch [1/50], loss_g: 6.0652, loss_d: 0.8181, real_score: 0.7618, fake_score: 0.2467

100%|████████████████████████████████████████████████████████████████████████████████| 169/169 [12:10<00:00,  4.32s/it]

Epoch [2/50], loss_g: 4.5312, loss_d: 0.7919, real_score: 0.7432, fake_score: 0.2549

100%|████████████████████████████████████████████████████████████████████████████████| 169/169 [11:59<00:00,  4.25s/it]

Epoch [3/50], loss_g: 4.5523, loss_d: 0.6902, real_score: 0.7637, fake_score: 0.2340

100%|████████████████████████████████████████████████████████████████████████████████| 169/169 [11:51<00:00,  4.21s/it]

Epoch [4/50], loss_g: 4.7291, loss_d: 0.7549, real_score: 0.7561, fake_score: 0.2433

100%|████████████████████████████████████████████████████████████████████████████████| 169/169 [11:55<00:00,  4.23s/it]

Epoch [5/50], loss_g: 5.0454, loss_d: 0.6721, real_score: 0.7698, fake_score: 0.2240


losses_g, losses_d, real_scores, fake_scores = history
plt.figure(figsize=(15,6))
plt.plot(losses_d, '-')
plt.plot(losses_g, '-')
plt.xlabel('epoch')
plt.ylabel('loss')
plt.legend(['Discriminator', 'Generator'])
plt.title('The Discriminator Loss and the Generator Loss');


generated_img = cv2.imread(f'./generated/generated-images-00{epochs}.png')
generated_img = generated_img[:, :, [2,1,0]]


fig, ax = plt.subplots(figsize=(8,8))
ax.set_xticks([]); ax.set_yticks([])
ax.imshow(generated_img);


generator_model = generator

# Load the generator model in PyTorch
loaded_generator = torch.load('generator_model.pth')


# Serialize the PyTorch generator model using pickle
with open('generator_model.pkl', 'wb') as f:
    pickle.dump(generator_model, f)

# Deserialize the PyTorch generator model using pickle
with open('generator_model.pkl', 'rb') as f:
    loaded_generator = pickle.load(f)

DCGAN: Anime Face Generator¶

Importing Necessary Libraries¶

Dataset Description¶

Data Preprocessing and Loading¶

¶

Constants:¶

Data Transformation and Loading:¶

Additional Functions for Visualization:¶

Transferring Data to the GPU Device¶

`to_device` Function:¶

`DeviceDataLoader` Class:¶

Setting the Device and Creating the Device DataLoader:¶

Discriminator Neural Network¶

Discriminator Architecture:¶

Device Placement:¶

Latent Size:¶

Generator Neural Network¶

Generator Architecture:¶

Device Placement:¶

Overall:¶

Initiating to save the generated images¶

Training the DCGAN¶

Visualizing the Discriminator Loss and the Generator Loss¶

Showing the Generated Images¶

Saving the Generator Model as a .pkl file¶

Streamlit User-Friendly Version¶

DCGAN: Anime Face Generator¶

Importing Necessary Libraries¶

Dataset Description¶

Data Preprocessing and Loading¶

¶

Constants:¶

Data Transformation and Loading:¶

Additional Functions for Visualization:¶

Transferring Data to the GPU Device¶

to_device Function:¶

DeviceDataLoader Class:¶

Setting the Device and Creating the Device DataLoader:¶

Discriminator Neural Network¶

Discriminator Architecture:¶

Device Placement:¶

Latent Size:¶

Generator Neural Network¶

Generator Architecture:¶

Device Placement:¶

Overall:¶

Initiating to save the generated images¶

Training the DCGAN¶

Visualizing the Discriminator Loss and the Generator Loss¶

Showing the Generated Images¶

Saving the Generator Model as a .pkl file¶

Streamlit User-Friendly Version¶

`to_device` Function:¶

`DeviceDataLoader` Class:¶