From bdaa67deb7ad095e9da4d217680d18690fd06007 Mon Sep 17 00:00:00 2001 From: James Betker Date: Wed, 12 Aug 2020 08:46:15 -0600 Subject: [PATCH] Misc --- codes/test.py | 41 +++++++++++++++++++++++++++++------ codes/train.py | 4 ++-- sandbox.py | 58 +++++++++++++++++++++++++++++++++++++++----------- 3 files changed, 82 insertions(+), 21 deletions(-) diff --git a/codes/test.py b/codes/test.py index 63b37eff..7f3fdd11 100644 --- a/codes/test.py +++ b/codes/test.py @@ -4,11 +4,13 @@ import time import argparse from collections import OrderedDict +import os import options.options as option import utils.util as util from data.util import bgr2ycbcr import models.archs.SwitchedResidualGenerator_arch as srg -from switched_conv_util import save_attention_to_image +from switched_conv_util import save_attention_to_image, save_attention_to_image_rgb +from switched_conv import compute_attention_specificity from data import create_dataset, create_dataloader from models import create_model from tqdm import tqdm @@ -22,14 +24,37 @@ import models.networks as networks def alter_srg(srg: srg.ConfigurableSwitchedResidualGenerator2): # First alteration, strip off switches one at a time. yield "naked" + + ''' for i in range(1, len(srg.switches)): srg.switches = srg.switches[:-i] yield "stripped-%i" % (i,) + ''' + + for sw in srg.switches: + sw.set_temperature(.001) + yield "specific" + + for sw in srg.switches: + sw.set_temperature(1000) + yield "normalized" + + for sw in srg.switches: + sw.set_temperature(1) + sw.switch.attention_norm = None + yield "no_anorm" return None def analyze_srg(srg: srg.ConfigurableSwitchedResidualGenerator2, path, alteration_suffix): - [save_attention_to_image(path, srg.attentions[i], srg.transformation_counts, i, "attention_" + alteration_suffix, - l_mult=5) for i in range(len(srg.attentions))] + mean_hists = [compute_attention_specificity(att, 2) for att in srg.attentions] + means = [i[0] for i in mean_hists] + hists = [torch.histc(i[1].clone().detach().cpu().flatten().float(), bins=srg.transformation_counts) for i in mean_hists] + hists = [h / torch.sum(h) for h in hists] + for i in range(len(means)): + print("%s - switch_%i_specificity" % (alteration_suffix, i), means[i]) + print("%s - switch_%i_histogram" % (alteration_suffix, i), hists[i]) + + [save_attention_to_image_rgb(path, srg.attentions[i], srg.transformation_counts, alteration_suffix, i) for i in range(len(srg.attentions))] def forward_pass(model, output_dir, alteration_suffix=''): @@ -60,7 +85,7 @@ if __name__ == "__main__": #### options torch.backends.cudnn.benchmark = True want_just_images = True - srg_analyze = True + srg_analyze = False parser = argparse.ArgumentParser() parser.add_argument('-opt', type=str, help='Path to options YMAL file.', default='../options/analyze_srg.yml') opt = option.parse(parser.parse_args().opt, is_train=False) @@ -106,14 +131,16 @@ if __name__ == "__main__": model_copy.load_state_dict(orig_model.state_dict()) model.netG = model_copy for alteration_suffix in alter_srg(model_copy): + alt_path = osp.join(dataset_dir, alteration_suffix) img_path = data['GT_path'][0] if need_GT else data['LQ_path'][0] - img_name = osp.splitext(osp.basename(img_path))[0] + img_name = osp.splitext(osp.basename(img_path))[0] + opt['name'] alteration_suffix += img_name + os.makedirs(alt_path, exist_ok=True) forward_pass(model, dataset_dir, alteration_suffix) - analyze_srg(model_copy, dataset_dir, alteration_suffix) + analyze_srg(model_copy, alt_path, alteration_suffix) # Reset model and do next alteration. model_copy = networks.define_G(opt).to(model.device) model_copy.load_state_dict(orig_model.state_dict()) model.netG = model_copy else: - forward_pass(model, dataset_dir) + forward_pass(model, dataset_dir, opt['name']) diff --git a/codes/train.py b/codes/train.py index 3c8a0fb8..c6f28bef 100644 --- a/codes/train.py +++ b/codes/train.py @@ -32,7 +32,7 @@ def init_dist(backend='nccl', **kwargs): def main(): #### options parser = argparse.ArgumentParser() - parser.add_argument('-opt', type=str, help='Path to option YAML file.', default='../options/train_imgset_spsr_switched.yml') + parser.add_argument('-opt', type=str, help='Path to option YAML file.', default='../options/train_imgset_spsr_switched_lr2.yml') parser.add_argument('--launcher', choices=['none', 'pytorch'], default='none', help='job launcher') parser.add_argument('--local_rank', type=int, default=0) @@ -161,7 +161,7 @@ def main(): current_step = resume_state['iter'] model.resume_training(resume_state) # handle optimizers and schedulers else: - current_step = -1 + current_step = 0 start_epoch = 0 #### training diff --git a/sandbox.py b/sandbox.py index 18825c15..97d8232d 100644 --- a/sandbox.py +++ b/sandbox.py @@ -1,22 +1,56 @@ import torch import torchvision from PIL import Image +from pytorch_wavelets import DWTForward, DWTInverse +import torch.nn.functional as F def load_img(path): - im = Image.open(path) + im = Image.open(path).convert(mode="RGB") return torchvision.transforms.ToTensor()(im) def save_img(t, path): torchvision.utils.save_image(t, path) -img = load_img("me.png") -# add zeros to the imaginary component -img = torch.stack([img, torch.zeros_like(img)], dim=-1) -fft = torch.fft(img, signal_ndim=2) -fft_d = torch.zeros_like(fft) -for i in range(-5, 5): - diag = torch.diagonal(fft, offset=i, dim1=1, dim2=2) - diag_em = torch.diag_embed(diag, offset=i, dim1=1, dim2=2) - fft_d += diag_em -resamp_img = torch.ifft(fft_d, signal_ndim=2)[:, :, :, 0] -save_img(resamp_img, "resampled.png") \ No newline at end of file +img = load_img("pu.jpg") +img = img.unsqueeze(0) + +# Reshape image to be multiple of 32 +w, h = img.shape[2:] +w = (w // 32) * 32 +h = (h // 32) * 32 +img = F.interpolate(img, size=(w, h)) +print("Input shape:", img.shape) + +J_spec = 5 + +Yl, Yh = DWTForward(J=J_spec, mode='periodization', wave='db3')(img) +print(Yl.shape, [h.shape for h in Yh]) + +imgLR = F.interpolate(img, scale_factor=.5) +LQYl, LQYh = DWTForward(J=J_spec-1, mode='periodization', wave='db3')(imgLR) +print(LQYl.shape, [h.shape for h in LQYh]) + +for i in range(J_spec): + smd = torch.sum(Yh[i], dim=2).cpu() + save_img(smd, "high_%i.png" % (i,)) +save_img(Yl, "lo.png") + +''' +Following code reconstructs the image with different high passes cancelled out. +''' +for i in range(J_spec): + corrupted_im = [y for y in Yh] + corrupted_im[i] = torch.zeros_like(corrupted_im[i]) + im = DWTInverse(mode='periodization', wave='db3')((Yl, corrupted_im)) + save_img(im, "corrupt_%i.png" % (i,)) +im = DWTInverse(mode='periodization', wave='db3')((torch.full_like(Yl, fill_value=torch.mean(Yl)), Yh)) +save_img(im, "corrupt_im.png") + + +''' +Following code reconstructs a hybrid image with the first high pass from the HR and the rest of the data from the LR. +highpass = [Yh[0]] + LQYh +im = DWTInverse(mode='periodization', wave='db3')((LQYl, highpass)) +save_img(im, "hybrid_lrhr.png") +save_img(F.interpolate(imgLR, scale_factor=2), "upscaled.png") +''' \ No newline at end of file