# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: animgan
#     language: python
#     name: animgan
# ---

# +
# %matplotlib inline
import os
import time
import shutil
import cv2
import imageio
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns

sns.set_style("dark")
plt.style.use("dark_background")

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.datasets
import torchvision.transforms as transforms

from datetime import datetime
from IPython import display

import utils
import model
import calculate_anim_mask
import batch_iterator
import train_model

training_timestamp = str(int(time.time()))
model_dir = f'trained_models/model_{training_timestamp}/'

if not os.path.exists(model_dir):
    os.makedirs(model_dir)

device = torch.device("cuda:0")


# -

def print_and_log(text):
    print(text)
    print(text, file=open(f'{model_dir}/log.txt', 'a'))


shutil.copy2('./animation_gan.ipynb', model_dir)

if os.path.isfile('preprocessed_data/processed_data.pt'):
    processed_data = torch.load('preprocessed_data/processed_data.pt')
else:
    processed_data = utils.preprocess_data(batch_size=64, data_size=40_000)
    torch.save(processed_data, 'preprocessed_data/processed_data.pt')

# +
batch_size = 8

dataset = torchvision.datasets.ImageFolder(root="data", transform=transforms.Compose([transforms.ToTensor(),
                                                                                      transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]))
dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=True)

batch_iter = batch_iterator.BatchIterator(processed_data, batch_size)
# -

real_images, _ = next(iter(dataloader))
cur_img = real_images[0].numpy().transpose(1, 2, 0)
plt.figure(figsize=(9*2, 4*2))
plt.imshow((cur_img+1.0)*0.5)

img_i = 0
frame_i = 0 # 0-35
cur_img_cube = utils.get_img_cube(real_images.numpy().transpose(0, 2, 3, 1), img_i)
plt.imshow((cur_img_cube[:,:,frame_i*3:(frame_i+1)*3]+1.0)*0.5)

batch_images = batch_iter.next_batch()
cur_img_plane = ((batch_images[0].type(torch.FloatTensor) / 255) - torch.Tensor([0.5])) / torch.Tensor([0.5])
plt.imshow((cur_img_plane[:3].permute(1, 2, 0)+1.0)*0.5)

anim_file = utils.animate_img_cube(cur_img_cube, f"{model_dir}/real.gif")
display.Image(filename=anim_file)

anim_file = utils.animate_img_batch(real_images.numpy().transpose(0, 2, 3, 1), f"{model_dir}/real_4.gif")
display.Image(filename=anim_file)

# +
latent_size = 100
# generator_net = model.Generator2D(latent_size=latent_size, input_channels=36*3, feature_map_size=64).to(device)
generator_net = model.Generator2D_10().to(device)
# generator_net = model.Generator3D().to(device)

# discriminator_net = model.Discriminator2D().to(device)
discriminator_net = model.Discriminator2D_10().to(device)
# discriminator_net = model.DiscriminatorTemporal().to(device)
# discriminator_net = model.Discriminator2D(input_channels=36*3, output_channels=36, feature_map_size=36*4, groups=36).to(device)
# discriminator_net = model.DiscriminatorSheet(input_channels=3, output_channels=1, feature_map_size=4).to(device)
# discriminator_net = model.Discriminator3D().to(device)

generator_net.apply(model.weights_init)
discriminator_net.apply(model.weights_init)

print_and_log(generator_net)
print_and_log(discriminator_net)
# -

learning_rate = 0.0002
beta1 = 0.5
discriminator_optimizer = optim.Adam(discriminator_net.parameters(), lr=learning_rate, betas=(beta1, 0.999))
generator_optimizer = optim.Adam(generator_net.parameters(), lr=learning_rate, betas=(beta1, 0.999))

# 2d
anim_mask = calculate_anim_mask.calculate_anim_mask(batch_size).to(device)
# 3d
# anim_mask = calculate_anim_mask.calculate_anim_mask(batch_size).reshape(-1, 3, 36, 64, 64).to(device)

print_and_log(f"{datetime.now()} Starting Training")
# generator_losses, discriminator_losses, generated_img_cubes = train_model.train_dataloader(discriminator_net, generator_net, discriminator_optimizer, generator_optimizer,
#                                                                                            model_dir, dataloader, anim_mask, latent_size, num_epochs=40, batch_size=batch_size)
generator_losses, discriminator_losses, generated_img_cubes = train_model.train_batchiterator(discriminator_net, generator_net, discriminator_optimizer, generator_optimizer,
                                                                                              model_dir, batch_iter, anim_mask, latent_size, epochs=40, batch_size=batch_size)

plt.figure(figsize=(10,5))
plt.title("Generator and Discriminator Loss During Training")
plt.plot(generator_losses,label="Generator")
plt.plot(discriminator_losses,label="Discriminator")
plt.xlabel("iterations")
plt.ylabel("Loss")
plt.legend()
plt.savefig(f"{model_dir}/loss.png")
plt.show()

torch.save(discriminator_net.state_dict(), f"{model_dir}/net_discriminator.pth")
torch.save(generator_net.state_dict(), f"{model_dir}/net_generator.pth")
np.save(f"{model_dir}/losses_generator.npy", np.array(generator_losses))
np.save(f"{model_dir}/losses_discriminator.npy", np.array(discriminator_losses))
np.save(f"{model_dir}/generated_img_cubes.npy", generated_img_cubes)

img_i = 0
frame_i = 0 # 0-35
eval_img_cube = generated_img_cubes[-1][img_i]
plt.imsave(f"{model_dir}/sample_output.png", cv2.resize((eval_img_cube[:,:,frame_i*3:(frame_i+1)*3]+1.0)*0.5, dsize=(0,0), fx=4.0, fy=4.0, interpolation=cv2.INTER_NEAREST))
plt.imshow((eval_img_cube[:,:,frame_i*3:(frame_i+1)*3]+1.0)*0.5)

anim_file = utils.animate_img_cube(eval_img_cube, f"{model_dir}/generated.gif")
display.Image(filename=anim_file)

anim_file = utils.animate_img_batch(generated_img_cubes[-1], f"{model_dir}/generated_4.gif", get_img_cubes=False)
display.Image(filename=anim_file)

anim_file = utils.animate_img_cube(generated_img_cubes, f"{model_dir}/training.gif", training_outputs=True)
display.Image(filename=anim_file)

anim_file = utils.animate_img_batch(generated_img_cubes, f"{model_dir}/training_4.gif", get_img_cubes=False, training_outputs=True)
display.Image(filename=anim_file)

plt.figure(figsize=(9*2, 5*2))
plt.imshow((utils.get_sample_frames(real_images[0])+1.0)*0.5)

plt.figure(figsize=(9*2, 5*2))
plt.imshow(utils.get_sample_frames(np.array([cv2.resize(frame, dsize=(0,0), fx=0.25, fy=0.25) for frame in imageio.get_reader(f"outputs/generated.gif")]), True))

plt.figure(figsize=(9*2, 5*2))
plt.imshow(utils.get_sample_frames(np.array([cv2.resize(frame, dsize=(0,0), fx=0.5, fy=0.5) for frame in imageio.get_reader(f"outputs/generated_4.gif")]), True, True))

plt.figure(figsize=(9*2, 5*2))
plt.imshow(utils.get_sample_frames(np.array([cv2.resize(frame, dsize=(0,0), fx=0.5, fy=0.5) for frame in imageio.get_reader(f"outputs/generated_4_64x64.gif")]), True, True, 90))

StarcoderdataJupyterScriptsDedupFiltered