Code - yjyoon/Raindrop

import sys import os import torch import glob import numpy as np import time import pandas as pd import subprocess import atexit import torchvision.transforms import cv2 from visdom import Visdom from torchvision.utils import save_image from torchvision import transforms from torchvision.transforms import RandomCrop, RandomPerspective, Compose from torch.utils.data import DataLoader from time import gmtime, strftime from model.Autoencoder import AutoEncoder from model.Generator import Generator from model.Discriminator import DiscriminativeNet as Discriminator from model.AttentiveRNN import AttentiveRNN from tools.argparser import get_param from tools.logger import Logger from tools.dataloader import ImagePairDataset args = get_param() # I am doing this for easier debugging # when you have error on those variables without doing this, # you will be in trouble because error message will not say anything. epochs = args.epochs batch_size = args.batch_size save_interval = args.save_interval sample_interval = args.sample_interval num_worker = args.num_worker device = args.device load = args.load generator_learning_rate = args.generator_learning_rate generator_ARNN_learning_rate = args.generator_ARNN_learning_rate if args.discriminator_learning_rate is not None else args.generator_learning_rate generator_learning_miniepoch = args.generator_learning_miniepoch generator_attentivernn_blocks = args.generator_attentivernn_blocks generator_resnet_depth = args.generator_resnet_depth discriminator_learning_rate = args.discriminator_learning_rate if args.discriminator_learning_rate is not None else args.generator_learning_rate logger = Logger() cuda = True if torch.cuda.is_available() else False device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') generator = Generator(generator_attentivernn_blocks, generator_resnet_depth).to(device=device) discriminator = Discriminator().to(device=device) if load is not None: generator.attentiveRNN.load_state_dict(torch.load(load)) generator.autoencoder.load_state_dict(torch.load(load)) discriminator.load_state_dict(torch.load(load)) else: pass # 이건 땜빵이고 차후에 데이터 관리 모듈 만들꺼임 rainy_data_path = glob.glob("data/source/Oxford_raindrop_dataset/dirty/*.png") rainy_data_path = sorted(rainy_data_path) clean_data_path = glob.glob("data/source/Oxford_raindrop_dataset/clean/*.png") clean_data_path = sorted(clean_data_path) height = 480 width = 720 transform = Compose([ RandomPerspective(), RandomCrop((height, width)) ]) resize = torchvision.transforms.Resize((height, width), antialias=True) dataset = ImagePairDataset(clean_data_path, rainy_data_path, transform=transform) dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=4) # declare generator loss optimizer_G = torch.optim.Adam(generator.parameters(), lr=generator_learning_rate) optimizer_G_ARNN = torch.optim.Adam(generator.attentiveRNN.parameters(), lr=generator_ARNN_learning_rate, betas=(0.5, 0.999)) optimizer_G_AE = torch.optim.Adam(generator.autoencoder.parameters(), lr=generator_learning_rate, betas=(0.5, 0.999)) optimizer_D = torch.optim.Adam(discriminator.parameters(), lr=discriminator_learning_rate) # ------visdom visualizer ---------- server_process = subprocess.Popen("python -m visdom.server", shell=True) # to ensure the visdom server process must stop whenever the script is terminated. def cleanup(): server_process.terminate() atexit.register(cleanup) time.sleep(10) vis = Visdom(server="http://localhost", port=8097) ARNN_loss_window = vis.line(Y=np.array([0]), X=np.array([0]), opts=dict(title='Generator-AttentionRNN Loss')) AE_loss_window = vis.line(Y=np.array([0]), X=np.array([0]), opts=dict(title='Generator-AutoEncoder Loss')) Discriminator_loss_window = vis.line(Y=np.array([0]), X=np.array([0]), opts=dict(title='Discriminator Loss')) Attention_map_visualizer = vis.image(np.zeros((height, width)), opts=dict(title='Attention Map')) Difference_mask_map_visualizer = vis.image(np.zeros((height, width)), opts=dict(title='Mask Map')) Generator_output_visualizer = vis.image(np.zeros((height, width)), opts=dict(title='Generated Derain Output')) Input_image_visualizer = vis.image(np.zeros((height, width)), opts=dict(title='input clean image')) for epoch_num, epoch in enumerate(range(epochs)): for i, imgs in enumerate(dataloader): img_batch = imgs clean_img = img_batch["clean_image"] / 255 rainy_img = img_batch["rainy_image"] / 255 # clean_img = clean_img.to() # rainy_img = rainy_img.to() clean_img = clean_img.to(device=device) rainy_img = rainy_img.to(device=device) optimizer_G_ARNN.zero_grad() optimizer_G_AE.zero_grad() attentiveRNNresults = generator.attentiveRNN(rainy_img) generator_attention_map = attentiveRNNresults['attention_map_list'] binary_difference_mask = generator.binary_diff_mask(clean_img, rainy_img, thresold=0.24) generator_loss_ARNN = generator.attentiveRNN.loss(rainy_img, binary_difference_mask) generator_loss_ARNN.backward() optimizer_G_ARNN.step() generator_outputs = generator.autoencoder(rainy_img * attentiveRNNresults['attention_map_list'][-1]) generator_result = generator_outputs['skip_3'] generator_output = generator_outputs['output'] generator_loss_AE = generator.autoencoder.loss(clean_img, rainy_img) generator_loss_AE.backward() optimizer_G_AE.step() optimizer_D.zero_grad() real_clean_prediction = discriminator(clean_img) fake_clean_prediction = discriminator(generator_result)["fc_out"] discriminator_loss = discriminator.loss(clean_img, generator_result, generator_attention_map) discriminator_loss.backward() optimizer_D.step() # Total loss optimizer_G.zero_grad() generator_loss_whole = generator_loss_AE + generator_loss_ARNN + torch.mean( torch.log(torch.subtract(1, fake_clean_prediction)) ) optimizer_G.step() losses = { "generator_loss_ARNN": generator_loss_ARNN, "generator_loss_AE" : generator_loss_AE, "discriminator loss" : discriminator_loss } logger.print_training_log(epoch_num, epochs, i, len(dataloader), losses) # visdom logger vis.line(Y=np.array([generator_loss_ARNN.item()]), X=np.array([epoch_num * epochs + i]), win=ARNN_loss_window, update='append') vis.line(Y=np.array([generator_loss_AE.item()]), X=np.array([epoch_num * epochs + i]), win=AE_loss_window, update='append') vis.line(Y=np.array([discriminator_loss.item()]), X=np.array([epoch_num * epochs + i]), win=Discriminator_loss_window, update='append') vis.image(generator_attention_map[-1][0, 0, :, :], win=Attention_map_visualizer, opts=dict(title="Attention Map")) vis.image(binary_difference_mask[-1], win=Difference_mask_map_visualizer, opts=dict(title="Binary Mask Map")) vis.image(generator_result[-1], win=Generator_output_visualizer, opts=dict(title="Generator Output")) vis.image(clean_img[-1], win=Input_image_visualizer, opts=dict(title="input clean image")) day = strftime("%Y-%m-%d %H:%M:%S", gmtime()) if epoch % save_interval == 0 and epoch != 0: torch.save(generator.attentiveRNN.state_dict(), f"weight/Attention_RNN_{epoch}_{day}.pt") torch.save(generator.state_dict(), f"weight/Generator_{epoch}_{day}.pt") torch.save(discriminator.state_dict(), f"weight/Discriminator_{epoch}_{day}.pt") server_process.terminate() ## RNN 따로 돌리고 CPU로 메모리 옳기고 ## Autoencoder 따로 돌리고 메모리 옳기고 ## 안되는가 ## 대충 열심히 GAN 구성하는 코드 ## 대충 그래서 weight export해서 inference용과 training용으로 나누는 코드 ## 대충 그래서 inference용은 attention map까지 하는 녀석과 deraining까지 하는 녀석 두개가 나오는 코드 ## 학습용은 그래서 풀 weight 나옴 ## GAN은 학습 시키면 Nash equilibrium ... 나오게 할 수 있으려나? ## 대충 학습은 어떻게 돌려야 되지 하는 코드 ## generator에서 튀어 나온 애들을 따로 저장해야 하는건가