image_output_11.png application_vnd.jupyter.stderr_output_27.png application_vnd.jupyter.stderr_output_35.png application_vnd.jupyter.stderr_output_24.png image_output_24.png application_vnd.jupyter.stderr_output_9.png application_vnd.jupyter.stderr_output_52.png application_vnd.jupyter.stderr_output_53.png text_plain_output_43.png image_output_25.png application_vnd.jupyter.stderr_output_48.png text_plain_output_30.png application_vnd.jupyter.stderr_output_2.png application_vnd.jupyter.stderr_output_7.png image_output_17.png application_vnd.jupyter.stderr_output_11.png image_output_30.png application_vnd.jupyter.stderr_output_38.png image_output_14.png image_output_39.png image_output_28.png application_vnd.jupyter.stderr_output_4.png application_vnd.jupyter.stderr_output_26.png image_output_23.png application_vnd.jupyter.stderr_output_6.png image_output_34.png application_vnd.jupyter.stderr_output_31.png application_vnd.jupyter.stderr_output_33.png application_vnd.jupyter.stderr_output_25.png image_output_13.png text_plain_output_45.png image_output_40.png image_output_5.png application_vnd.jupyter.stderr_output_12.png text_plain_output_14.png image_output_18.png text_plain_output_32.png text_plain_output_29.png application_vnd.jupyter.stderr_output_8.png text_plain_output_58.png image_output_21.png image_output_7.png application_vnd.jupyter.stderr_output_10.png application_vnd.jupyter.stderr_output_23.png image_output_31.png text_plain_output_47.png application_vnd.jupyter.stderr_output_34.png image_output_20.png text_plain_output_18.png image_output_32.png text_plain_output_3.png application_vnd.jupyter.stderr_output_19.png image_output_4.png application_vnd.jupyter.stderr_output_44.png application_vnd.jupyter.stderr_output_13.png application_vnd.jupyter.stderr_output_42.png application_vnd.jupyter.stderr_output_5.png image_output_35.png image_output_36.png image_output_8.png image_output_37.png text_plain_output_16.png image_output_16.png application_vnd.jupyter.stderr_output_15.png text_plain_output_59.png application_vnd.jupyter.stderr_output_17.png image_output_27.png application_vnd.jupyter.stderr_output_28.png application_vnd.jupyter.stderr_output_46.png image_output_6.png application_vnd.jupyter.stderr_output_41.png application_vnd.jupyter.stderr_output_20.png application_vnd.jupyter.stderr_output_49.png application_vnd.jupyter.stderr_output_36.png application_vnd.jupyter.stderr_output_57.png application_vnd.jupyter.stderr_output_22.png image_output_12.png image_output_22.png application_vnd.jupyter.stderr_output_56.png application_vnd.jupyter.stderr_output_50.png text_plain_output_1.png image_output_3.png application_vnd.jupyter.stderr_output_51.png image_output_29.png image_output_2.png image_output_1.png image_output_10.png application_vnd.jupyter.stderr_output_39.png image_output_33.png application_vnd.jupyter.stderr_output_21.png application_vnd.jupyter.stderr_output_54.png image_output_15.png application_vnd.jupyter.stderr_output_55.png image_output_9.png image_output_19.png application_vnd.jupyter.stderr_output_40.png application_vnd.jupyter.stderr_output_37.png image_output_38.png image_output_26.png
from torchvision.utils import make_grid

from tqdm import tqdm

import copy

import matplotlib.pyplot as plt

import matplotlib.pyplot as plt

import torch

import torch.nn as nn

import torch.nn as nn

import torch.nn.functional as F

import torchvision.datasets as datasets

import torchvision.transforms as transforms

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
runs = [{'latent': 20, 'nE': 100, 'bs': 64, 'cap': 64, 'lr': 0.001, 'wd': 1e-05, 'vb': 1}, {'latent': 20, 'nE': 200, 'bs': 64, 'cap': 64, 'lr': 0.001, 'wd': 1e-05, 'vb': 1}, {'latent': 20, 'nE': 100, 'bs': 64, 'cap': 64, 'lr': 0.001, 'wd': 1e-05, 'vb': 0.5}, {'latent': 20, 'nE': 200, 'bs': 64, 'cap': 64, 'lr': 0.001, 'wd': 1e-05, 'vb': 0.5}]
train_dir = '../input/covid19-chest-ct-image-augmentation-gan-dataset/COVID-19/COVID-19/train'
train_dataset = datasets.ImageFolder(train_dir, transforms.Compose([transforms.Resize((256, 256)), transforms.ToTensor()]))


def show_image(image_tensor, num_images=25, size=(1, 28, 28)):
    '''
    Function for visualizing images: Given a tensor of images, number of images, and
    size per image, plots and prints the images in an uniform grid.
    '''
    image_tensor = (image_tensor + 1) / 2
    image_unflat = image_tensor.detach().cpu()
    image_grid = make_grid(image_unflat[:num_images], nrow=5)
    plt.imshow(image_grid.permute(1, 2, 0).squeeze())
    plt.show()

def plot_loss(training_loss, settings):
    fig, ax = plt.subplots()
    ax.plot(range(len(training_loss)), training_loss)

    ax.set(xlabel='epochs', ylabel='BCE loss')
    ax.set_title(f'Loss for {settings}', y=1.1)
    ax.grid()

    fig.savefig("test.png")
    plt.show()

class Encoder(nn.Module):
    """encoder for VAE, goes from image conv net to linear latent layer"""
    def __init__(self, capacity, latent_dims):
        super(Encoder, self).__init__()
        c = capacity
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=c, kernel_size=4, stride=2, padding=1)
        self.conv2 = nn.Conv2d(in_channels=c, out_channels=c * 2, kernel_size=4, stride=2, padding=1)
        self.conv3 = nn.Conv2d(in_channels=c * 2, out_channels=c * 2 * 2, kernel_size=4, stride=2, padding=1)
        self.conv4 = nn.Conv2d(in_channels=c * 2 * 2, out_channels=c * 2 * 2 * 2, kernel_size=4, stride=2, padding=1)
        self.conv5 = nn.Conv2d(in_channels=c * 2 * 2 * 2, out_channels=c * 2 * 2 * 2 * 2, kernel_size=4, stride=2, padding=1)
        self.conv6 = nn.Conv2d(in_channels=c * 2 * 2 * 2 * 2, out_channels=c * 2 * 2 * 2 * 2 * 2, kernel_size=4, stride=2, padding=1)
        self.fc_mu = nn.Linear(in_features=32768, out_features=latent_dims)
        self.fc_logvar = nn.Linear(in_features=32768, out_features=latent_dims)
    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        x = F.relu(self.conv3(x))
        x = F.relu(self.conv4(x))
        x = F.relu(self.conv5(x))
        x = F.relu(self.conv6(x))
        x = x.view(x.size(0), -1)
        x_mu = self.fc_mu(x)
        x_logvar = self.fc_logvar(x)
        return (x_mu, x_logvar)
class Decoder(nn.Module):
    """decoder for VAE, goes from linear latent layer to deconv layers to reconstruct image"""
    def __init__(self, capacity, latent_dims):
        super(Decoder, self).__init__()
        c = capacity
        self.fc = nn.Linear(in_features=latent_dims, out_features=32768)
        self.conv1 = nn.ConvTranspose2d(in_channels=c, out_channels=3, kernel_size=4, stride=2, padding=1)
        self.conv2 = nn.ConvTranspose2d(out_channels=c, in_channels=c * 2, kernel_size=4, stride=2, padding=1)
        self.conv3 = nn.ConvTranspose2d(out_channels=c * 2, in_channels=c * 2 * 2, kernel_size=4, stride=2, padding=1)
        self.conv4 = nn.ConvTranspose2d(out_channels=c * 2 * 2, in_channels=c * 2 * 2 * 2, kernel_size=4, stride=2, padding=1)
        self.conv5 = nn.ConvTranspose2d(out_channels=c * 2 * 2 * 2, in_channels=c * 2 * 2 * 2 * 2, kernel_size=4, stride=2, padding=1)
        self.conv6 = nn.ConvTranspose2d(out_channels=c * 2 * 2 * 2 * 2, in_channels=c * 2 * 2 * 2 * 2 * 2, kernel_size=4, stride=2, padding=1)
    def forward(self, x):
        x = self.fc(x)
        x = x.view(x.size(0), 2048, 4, 4)
        x = F.relu(self.conv6(x))
        x = F.relu(self.conv5(x))
        x = F.relu(self.conv4(x))
        x = F.relu(self.conv3(x))
        x = F.relu(self.conv2(x))
        x = torch.sigmoid(self.conv1(x))
        return x
class VAE(nn.Module):
    """VAE architecture for encoder -> sample from latent -> decode latent sample"""
    def __init__(self, capacity, latent_dims):
        super(VAE, self).__init__()
        self.encoder = Encoder(capacity, latent_dims)
        self.decoder = Decoder(capacity, latent_dims)
    def forward(self, x):
        latent_mu, latent_logvar = self.encoder(x)
        latent = self.latent_sample(latent_mu, latent_logvar)
        x_recon = self.decoder(latent)
        return (x_recon, latent_mu, latent_logvar)
    def latent_sample(self, mu, logvar):
        if self.training:
            std = logvar.mul(0.5).exp_()
            eps = torch.empty_like(std).normal_()
            return eps.mul(std).add_(mu)
        else:
            return mu
def vae_loss(recon_x, x, mu, logvar, variational_beta):
    recon_loss = F.binary_cross_entropy(recon_x.view(-1, 65536), x.view(-1, 65536), reduction='sum')
    kldivergence = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
    return recon_loss + variational_beta * kldivergence

def setup_model(capacity, latent_dims):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    vae = VAE(capacity=capacity, latent_dims=latent_dims).to(device)
    return vae

def train(vae, train_loader, n_epochs, learning_rate, weight_decay, variational_beta):
    optimizer = torch.optim.Adam(params=vae.parameters(), lr=learning_rate, weight_decay=weight_decay)
    vae.train()
    train_loss = []
    best_model_wts = None
    bmw_epoch = 0
    for epoch in tqdm(range(n_epochs)):
        num_batches = 0
        avg_loss = 0
        best_loss = 0
        image_batch_recon = None
        for image_batch, _ in train_loader:
            image_batch = image_batch.to(device)
            image_batch_recon, latent_mu, latent_logvar = vae(image_batch)
            loss = vae_loss(image_batch_recon, image_batch, latent_mu, latent_logvar, variational_beta)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            avg_loss += loss.item()
            num_batches += 1
        avg_loss /= num_batches
        train_loss.append(avg_loss)
        if epoch == 0:
            best_loss = avg_loss
            best_model_wts = copy.deepcopy(vae.state_dict())
        if avg_loss < best_loss:
            best_model_wts = copy.deepcopy(vae.state_dict())
            best_loss = avg_loss
            bmw_epoch = epoch + 1
    return (vae, best_model_wts, bmw_epoch, train_loss)

for k, settings in enumerate(runs):
    print(f'starting run ... {k}/{len(runs)}')
    print(settings)
    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=settings['bs'], shuffle=True, num_workers=0)
    vae = setup_model(settings['cap'], settings['latent'])
    vae, best_model_wts, bmw_epoch, train_loss = train(vae, train_loader, settings['nE'], settings['lr'], settings['wd'], settings['vb'])
    last_modeL_wts = vae.state_dict()
    plot_loss(train_loss, settings)
    print('-------------------------------------------------------------------')
    print('saving weights for model')
    torch.save(vae.state_dict(), f"n_{settings['nE']}.ld_{settings['latent']}.lr_{settings['lr']}.vb_{settings['vb']}.last_model.wts")
    vae.load_state_dict(best_model_wts)
    torch.save(vae.state_dict(), f"n_{bmw_epoch}.ld_{settings['latent']}.lr_{settings['lr']}.vb_{settings['vb']}.lowest_loss_model.wts")
    print('-------------------------------------------------------------------')
    print()