Upload 20 files

LMAR_GAN_train.py

@@ -0,0 +1,331 @@
+import argparse
+import yaml
+import torchvision.transforms as transforms
+from utils import read_args, save_checkpoint, AverageMeter, CosineAnnealingWarmRestarts
+import time
+from tqdm import trange, tqdm
+from torchvision.utils import save_image
+# from tensorboardX import SummaryWriter
+import os
+import json
+import time
+import logging
+
+os.environ['CUDA_VISIBLE_DEVICES'] = '1'
+import torch
+from torch import optim
+import torch.nn as nn
+import torchvision.utils as vutils
+import torch.nn.functional as F
+
+from data import *
+from model import *
+from loss import *
+import pyiqa
+from torch.autograd import Variable
+import numpy as np
+
+global_step = 0
+psnr_calculator = pyiqa.create_metric('psnr').cuda()
+ssim_calculator = pyiqa.create_metric('ssimc', downsample=True).cuda()
+
+criterion_GAN = nn.MSELoss()
+Tensor = torch.cuda.FloatTensor
+
+mmdLoss = MMDLoss().cuda()
+
+# cos_loss = cos_loss
+# feature_extractor.eval()
+
+
+def train(model, data_loader, criterion, optimizer_G, optimizer_D, epoch, args, discriminator):
+    global global_step
+    iter_bar = tqdm(data_loader, desc='Iter (loss=X.XXX)')
+    nbatches = len(data_loader)
+
+    total_losses = AverageMeter()
+
+    pixel_losses = AverageMeter()
+    resize_losses = AverageMeter()
+    pseudo_losses = AverageMeter()
+    up_losses = AverageMeter()
+    dis_losses = AverageMeter()
+
+    psnrs = AverageMeter()
+    ssims = AverageMeter()
+
+    optimizer_G.zero_grad()
+    optimizer_D.zero_grad()
+
+    start_time = time.time()
+
+    if not os.path.exists(args.output_dir + '/image_train'):
+        os.mkdir(args.output_dir + '/image_train')
+
+    if not os.path.exists(args.output_dir + "/models"):
+        os.mkdir(args.output_dir + "/models")
+
+    for i, batch in enumerate(iter_bar):
+        optimizer_G.zero_grad()
+        optimizer_D.zero_grad()
+
+        inp_img, gt_img, down_h, down_w, inp_img_path = batch
+        batch_size = inp_img.size(0)
+        inp_img = inp_img.cuda()
+        gt_img = gt_img.cuda()
+
+        down_size = (down_h.item(), down_w.item())
+        up_size = eval(args.train_loader["img_size"])
+
+        down_x, hr_feature, new_lr_feature, ori_lr_feature, residual, res = model(inp_img, down_size, up_size)
+
+
+        dis_patch_lr = (1, down_size[0] // 2 ** 4, down_size[1] // 2 ** 4)
+        valid_lr = Variable(Tensor(np.ones((batch_size, *dis_patch_lr))), requires_grad=False)
+        fake_lr = Variable(Tensor(np.zeros((batch_size, *dis_patch_lr))), requires_grad=False)
+
+
+        pixel_loss = criterion_GAN(discriminator(down_x), valid_lr)
+        pixel_losses.update(pixel_loss.item(), batch_size)
+
+        resize_loss = criterion(hr_feature, new_lr_feature)
+        resize_losses.update(resize_loss.item(), batch_size)
+
+        pseudo_loss = similarity_loss(new_lr_feature, hr_feature) * 5000
+        pseudo_losses.update(pseudo_loss.item(), batch_size)
+
+        up_loss, gradient = feat_ssim(new_lr_feature, hr_feature, inp_img)
+        up_losses.update(up_loss.item(), batch_size)
+
+        total_loss = pixel_loss + resize_loss + pseudo_loss + up_loss
+        total_losses.update(total_loss.item(), batch_size)
+
+        total_loss.backward()
+        optimizer_G.step()
+
+
+
+        loss_real_lr = criterion_GAN(discriminator(resize(inp_img, out_shape=down_size, antialiasing=False)), valid_lr)
+        
+        loss_fake_lr = criterion_GAN(discriminator(down_x.detach()), fake_lr)
+
+        loss_D = (loss_fake_lr + loss_real_lr) * 0.5
+        dis_losses.update(loss_D.item(), batch_size)
+
+        loss_D.backward()
+        optimizer_D.step()
+
+        iter_bar.set_description('Iter (loss=%5.6f)' % (total_losses.avg + dis_losses.avg))
+
+        if i % 200 == 0:
+            error = torch.abs(resize(inp_img, out_shape=down_size, antialiasing=False) - down_x)
+            saved_image = torch.cat(
+                [resize(inp_img, out_shape=down_size, antialiasing=False)[0:2], down_x[0:2], error[0:2]],
+                dim=0)
+            save_image(saved_image, args.output_dir + '/image_train/epoch_{}_iter_down_{}.png'.format(epoch, i))
+
+            saved_image = torch.cat(
+                [torch.mean(hr_feature, dim=1, keepdim=True)[0:2], torch.mean(new_lr_feature, dim=1, keepdim=True)[0:2],
+                 torch.mean(ori_lr_feature, dim=1, keepdim=True)[0:2], torch.mean(torch.abs(new_lr_feature-ori_lr_feature), dim=1, keepdim=True)[0:2]],
+                dim=0)
+            save_image(saved_image, args.output_dir + '/image_train/epoch_{}_iter_feat_{}.png'.format(epoch, i))
+            residual = residual * 10
+            save_image(residual[0], args.output_dir + '/image_train/epoch_{}_iter_out_{}.png'.format(epoch, i))
+
+        if i % max(1, nbatches // 10) == 0:
+            psnr_val, ssim_val = 0.0, 0.0
+            psnrs.update(psnr_val, batch_size)
+            ssims.update(ssim_val, batch_size)
+
+            logging.info(
+                "Epoch {}, learning rates {:}, Iter {}, total_loss {:.4f}, pixel_loss {:.4f}, resize_loss {:.4f}, pseudo_loss {:.4f}, up_loss {:.4f}, dis_loss: {:.4f}, PSNR {:.4f}, SSIM {:.4f}, Elapse time {:.2f}\n".format(
+                    epoch, optimizer_G.param_groups[0]["lr"], i, total_losses.avg, pixel_losses.avg, resize_losses.avg,
+                    pseudo_losses.avg, up_losses.avg, dis_losses.avg,
+                    psnrs.avg, ssims.avg,
+                    time.time() - start_time))
+
+    if epoch % 1 == 0:
+        logging.info("** ** * Saving model and optimizer ** ** * ")
+
+        output_model_file = os.path.join(args.output_dir + "/models", "model.%d.bin" % (epoch))
+        state = {"epoch": epoch, "state_dict": model.state_dict(), "step": global_step}
+        save_checkpoint(state, output_model_file)
+
+        output_model_file = os.path.join(args.output_dir + "/models", "discriminator.%d.bin" % (epoch))
+        state = {"epoch": epoch, "state_dict": discriminator.state_dict(), "step": global_step}
+        save_checkpoint(state, output_model_file)
+        logging.info("Save model to %s", output_model_file)
+
+    logging.info(
+        "Finish training epoch %d, avg total_loss: %.4f, avg pixel_loss: %.4f, avg resize_loss: %.4f, avg pseudo_loss: %.4f, avg up_loss: %.4f, "
+        "avg dis_loss: %.4f, avg PSNR: %.2f, avg SSIM: %.2F, and takes %.2f seconds" % (
+            epoch, total_losses.avg, pixel_losses.avg, resize_losses.avg, pseudo_losses.avg, up_losses.avg, dis_losses.avg, psnrs.avg,
+            ssims.avg,
+            time.time() - start_time))
+
+    logging.info("***** CUDA.empty_cache() *****\n")
+    torch.cuda.empty_cache()
+
+
+def evaluate(model, load_path, data_loader, epoch):
+    checkpoint = torch.load(load_path)
+    model.load_state_dict(checkpoint["state_dict"])
+    model.cuda()
+    model.eval()
+
+    psnrs = AverageMeter()
+    ssims = AverageMeter()
+    random_index = torch.randint(low=0, high=5, size=(1,))
+    down_size = eval(args.test_loader["img_size"])
+    down_size = down_size[random_index]
+    logging.info("Inference at down size: {}".format(down_size))
+    up_size = eval(args.test_loader["gt_size"])
+
+    start_time = time.time()
+    with torch.no_grad():
+        for i, batch in enumerate(tqdm(data_loader)):
+            inp_img, gt_img, inp_img_path = batch
+            inp_img = inp_img.cuda()
+            batch_size = inp_img.size(0)
+            up_out, _ = model(inp_img, down_size, up_size, test_flag=True)
+
+            # metrics
+            clamped_out = torch.clamp(up_out, 0, 1)
+            psnr_val, ssim_val = psnr_calculator(clamped_out, gt_img), ssim_calculator(clamped_out, gt_img)
+            psnrs.update(torch.mean(psnr_val).item(), batch_size)
+            ssims.update(torch.mean(ssim_val).item(), batch_size)
+            torch.cuda.empty_cache()
+
+            if i % 100 == 0:
+                logging.info(
+                    "PSNR {:.4f}, SSIM {:.4f}, Elapse time {:.2f}\n".format(psnrs.avg, ssims.avg,
+                                                                            time.time() - start_time))
+
+        logging.info("avg PSNR: %.4f, avg SSIM: %.4F, and takes %.2f seconds" % (
+            psnrs.avg, ssims.avg, time.time() - start_time))
+
+
+def main(args):
+    global global_step
+
+    start_epoch = 1
+    global_step = 0
+
+    if not os.path.exists(args.output_dir):
+        os.mkdir(args.output_dir)
+
+    with open(os.path.join(args.output_dir, "args.json"), "w") as f:
+        json.dump(args.__dict__, f, sort_keys=True, indent=2)
+
+    log_format = "%(asctime)s %(levelname)-8s %(message)s"
+    log_file = os.path.join(args.output_dir, "train_log")
+    logging.basicConfig(filename=log_file, level=logging.INFO, format=log_format)
+    logging.getLogger().addHandler(logging.StreamHandler())
+
+    # device setting
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    args.device = device
+
+    logging.info(args.__dict__)
+
+    model = codebook_model(args)
+
+    discriminator = Discriminator(3).cuda()
+
+
+    optimizer_G = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=args.optimizer["lr"],
+                             betas=(0.9, 0.999))
+
+    optimizer_D = optim.Adam(list(discriminator.parameters()),
+                             lr=args.optimizer["lr"],
+                             betas=(0.9, 0.999))
+
+    logging.info("Building data loader")
+
+    if args.train_loader["loader"] == "resize":
+        train_transforms = transforms.Compose([transforms.Resize(eval(args.train_loader["img_size"])),
+                                               transforms.ToTensor()])
+        train_loader = get_loader(args.data["train_dir"],
+                                  eval(args.train_loader["img_size"]), train_transforms, False,
+                                  int(args.train_loader["batch_size"]), args.train_loader["num_workers"],
+                                  args.train_loader["shuffle"], random_flag=False)
+
+    elif args.train_loader["loader"] == "crop":
+        train_loader = get_loader(args.data["train_dir"],
+                                  eval(args.train_loader["img_size"]), False, True,
+                                  int(args.train_loader["batch_size"]), args.train_loader["num_workers"],
+                                  args.train_loader["shuffle"], random_flag=args.train_loader["random_flag"])
+
+    elif args.train_loader["loader"] == "default":
+        train_transforms = transforms.Compose([transforms.ToTensor()])
+        train_loader = get_loader(args.data["train_dir"],
+                                  eval(args.train_loader["img_size"]), train_transforms, False,
+                                  int(args.train_loader["batch_size"]), args.train_loader["num_workers"],
+                                  args.train_loader["shuffle"], random_flag=args.train_loader["random_flag"])
+    else:
+        raise NotImplementedError
+
+    if args.test_loader["loader"] == "default":
+
+        test_transforms = transforms.Compose([transforms.ToTensor()])
+        test_loader = get_loader(args.data["test_dir"],
+                                 None, test_transforms, False,
+                                 int(args.test_loader["batch_size"]), args.test_loader["num_workers"],
+                                 args.test_loader["shuffle"], random_flag=False)
+
+    elif args.test_loader["loader"] == "resize":
+
+        test_transforms = transforms.Compose([transforms.Resize(eval(args.test_loader["img_size"])),
+                                              transforms.ToTensor()])
+        test_loader = get_loader(args.data["test_dir"],
+                                 eval(args.test_loader["img_size"]), test_transforms, False,
+                                 int(args.test_loader["batch_size"]), args.test_loader["num_workers"],
+                                 args.test_loader["shuffle"], random_flag=False)
+    else:
+        raise NotImplementedError
+
+    # criterion = similarity_loss
+    criterion = nn.SmoothL1Loss()
+    # criterion = nn.L1Loss()
+
+    # vgg_loss = VGGLoss()
+
+    if args.optimizer["type"] == "cos":
+        lr_scheduler = CosineAnnealingWarmRestarts(optimizer, T_0=args.optimizer["T_0"],
+                                                   T_mult=args.optimizer["T_MULT"],
+                                                   eta_min=args.optimizer["ETA_MIN"],
+                                                   last_epoch=-1)
+    elif args.optimizer["type"] == "step":
+        lr_scheduler_G = torch.optim.lr_scheduler.StepLR(optimizer_G, step_size=args.optimizer["step"],
+                                                         gamma=args.optimizer["gamma"])
+        lr_scheduler_D = torch.optim.lr_scheduler.StepLR(optimizer_D, step_size=args.optimizer["step"],
+                                                         gamma=args.optimizer["gamma"])
+
+    t_total = int(len(train_loader) * args.optimizer["total_epoch"])
+    logging.info("***** CUDA.empty_cache() *****")
+    torch.cuda.empty_cache()
+
+    logging.info("***** Running training *****")
+    logging.info("  Batch size = %d", args.train_loader["batch_size"])
+    logging.info("  Num steps = %d", t_total)
+    logging.info("  Loader length = %d", len(train_loader))
+
+    model.train()
+    model.cuda()
+
+    logging.info("Begin training from epoch = %d\n", start_epoch)
+    for epoch in trange(start_epoch, args.optimizer["total_epoch"] + 1, desc="Epoch"):
+        train(model, train_loader, criterion, optimizer_G, optimizer_D, epoch, args, discriminator)
+        lr_scheduler_G.step()
+        lr_scheduler_D.step()
+        if epoch % args.evaluate_intervel == 0:
+            logging.info("***** Running testing *****")
+            load_path = os.path.join(args.output_dir + "/models", "model.%d.bin" % (epoch))
+            evaluate(model, load_path, test_loader, epoch)
+            logging.info("***** End testing *****")
+
+
+if __name__ == '__main__':
+    parser = read_args("/home/yuwei/code/cvpr/config/LMAR_config.yaml")
+    args = parser.parse_args()
+    main(args)

LMAR_VGG_train.py

@@ -0,0 +1,300 @@
+import os
+
+os.environ['CUDA_VISIBLE_DEVICES'] = '1'
+import argparse
+import yaml
+import torchvision.transforms as transforms
+from utils import read_args, save_checkpoint, AverageMeter, CosineAnnealingWarmRestarts
+import time
+from tqdm import trange, tqdm
+from torchvision.utils import save_image
+import json
+import time
+import logging
+import torch
+from torch import nn, optim
+import torchvision.utils as vutils
+import torch.nn.functional as F
+import torch.nn as nn
+
+from data import *
+from model import *
+from loss import *
+import pyiqa
+
+from torch.autograd import Variable
+
+global_step = 0
+psnr_calculator = pyiqa.create_metric('psnr').cuda()
+ssim_calculator = pyiqa.create_metric('ssimc', downsample=True).cuda()
+
+feature_extractor = VGGPerceptualLoss(resize=False).cuda()
+feature_extractor.eval()
+
+
+def weight_annealing(epoch):
+    initial_weight = 1
+    if epoch < 2:
+        return initial_weight  # 初始阶段保持权重不变
+    else:
+        return initial_weight * 0.001  # 后续阶段权重继续减小
+
+
+def train(model, data_loader, criterion, optimizer, epoch, args):
+    global global_step
+    iter_bar = tqdm(data_loader, desc='Iter (loss=X.XXX)')
+    nbatches = len(data_loader)
+
+    total_losses = AverageMeter()
+
+    pixel_losses = AverageMeter()
+    resize_losses = AverageMeter()
+    pseudo_losses = AverageMeter()
+
+    psnrs = AverageMeter()
+    ssims = AverageMeter()
+
+    optimizer.zero_grad()
+
+    start_time = time.time()
+
+    if not os.path.exists(args.output_dir + '/image_train'):
+        os.mkdir(args.output_dir + '/image_train')
+
+    if not os.path.exists(args.output_dir + "/models"):
+        os.mkdir(args.output_dir + "/models")
+
+    for i, batch in enumerate(iter_bar):
+        optimizer.zero_grad()
+        inp_img, gt_img, down_h, down_w, inp_img_path = batch
+        batch_size = inp_img.size(0)
+        inp_img = inp_img.cuda()
+        gt_img = gt_img.cuda()
+
+        down_size = (down_h.item(), down_w.item())
+        up_size = eval(args.train_loader["img_size"])
+
+
+        down_x, hr_feature, new_lr_feature, ori_lr_feature, residual, res = model(inp_img, down_size, up_size)
+
+        pixel_loss = criterion(new_lr_feature, hr_feature)
+
+        pixel_losses.update(pixel_loss.item(), batch_size)
+        
+        resize_loss = feature_extractor(down_x, resize(inp_img, out_shape=down_size, antialiasing=False),
+                                        feature_layers=[3])
+        resize_loss = resize_loss * weight_annealing(epoch)
+        resize_losses.update(resize_loss.item(), batch_size)
+
+
+        pseudo_loss, _ = feat_ssim(new_lr_feature, hr_feature, inp_img)
+        pseudo_losses.update(pseudo_loss.item(), batch_size)
+
+        total_loss = pixel_loss + resize_loss + pseudo_loss
+        total_losses.update(total_loss.item(), batch_size)
+
+        total_loss.backward()
+
+        optimizer.step()
+
+        iter_bar.set_description('Iter (loss=%5.6f)' % total_losses.avg)
+
+        if i % 200 == 0:
+            # print(residual.max())
+            error = torch.abs(resize(inp_img, out_shape=down_size, antialiasing=False) - down_x)
+            # error = (error - error.min()) / (error.max()-error.min())
+            saved_image = torch.cat(
+                [resize(inp_img, out_shape=down_size, antialiasing=False)[0:2], down_x[0:2], error[0:2]],
+                dim=0)
+            save_image(saved_image, args.output_dir + '/image_train/epoch_{}_iter_down_{}.png'.format(epoch, i))
+
+            saved_image = torch.cat(
+                [torch.mean(hr_feature, dim=1, keepdim=True)[0:2], torch.mean(new_lr_feature, dim=1, keepdim=True)[0:2],
+                 torch.mean(ori_lr_feature, dim=1, keepdim=True)[0:2],
+                 torch.mean(torch.abs(new_lr_feature - ori_lr_feature), dim=1, keepdim=True)[0:2]],
+                dim=0)
+            save_image(saved_image, args.output_dir + '/image_train/epoch_{}_iter_feat_{}.png'.format(epoch, i))
+            # residual = (residual - residual.min()) / (residual.max()-residual.min())
+            residual = residual * 10
+            save_image(residual[0], args.output_dir + '/image_train/epoch_{}_iter_out_{}.png'.format(epoch, i))
+
+        if i % max(1, nbatches // 10) == 0:
+            psnr_val, ssim_val = 0.0, 0.0
+            psnrs.update(psnr_val, batch_size)
+            ssims.update(ssim_val, batch_size)
+
+            logging.info(
+                "Epoch {}, learning rates {:}, Iter {}, total_loss {:.4f}, pixel_loss {:.4f}, resize_loss {:.4f}, pseudo_loss {:.4f}, PSNR {:.4f}, SSIM {:.4f}, Elapse time {:.2f}\n".format(
+                    epoch, optimizer.param_groups[0]["lr"], i, total_losses.avg, pixel_losses.avg, resize_losses.avg,
+                    pseudo_losses.avg,
+                    psnrs.avg, ssims.avg,
+                    time.time() - start_time))
+
+    if epoch % 1 == 0:
+        logging.info("** ** * Saving model and optimizer ** ** * ")
+
+        output_model_file = os.path.join(args.output_dir + "/models", "model.%d.bin" % (epoch))
+        state = {"epoch": epoch, "state_dict": model.state_dict(),
+                 "optimizer": optimizer.state_dict(), "step": global_step}
+
+        save_checkpoint(state, output_model_file)
+        logging.info("Save model to %s", output_model_file)
+
+    logging.info(
+        "Finish training epoch %d, avg total_loss: %.4f, avg pixel_loss: %.4f, avg resize_loss: %.4f, avg pseudo_loss: %.4f, avg PSNR: %.2f, avg SSIM: %.2F, and takes %.2f seconds" % (
+            epoch, total_losses.avg, pixel_losses.avg, resize_losses.avg, pseudo_losses.avg, psnrs.avg, ssims.avg,
+            time.time() - start_time))
+
+    logging.info("***** CUDA.empty_cache() *****\n")
+    torch.cuda.empty_cache()
+
+
+def evaluate(model, load_path, data_loader, epoch):
+    checkpoint = torch.load(load_path)
+    model.load_state_dict(checkpoint["state_dict"])
+    model.cuda()
+    model.eval()
+
+    psnrs = AverageMeter()
+    ssims = AverageMeter()
+    random_index = torch.randint(low=0, high=5, size=(1,))
+    down_size = eval(args.test_loader["img_size"])
+    down_size = down_size[random_index]
+    logging.info("Inference at down size: {}".format(down_size))
+    up_size = eval(args.test_loader["gt_size"])
+
+    start_time = time.time()
+    with torch.no_grad():
+        for i, batch in enumerate(tqdm(data_loader)):
+            inp_img, gt_img, inp_img_path = batch
+            inp_img = inp_img.cuda()
+            batch_size = inp_img.size(0)
+            up_out = model(inp_img, down_size, up_size, test_flag=True)
+
+            # metrics
+            clamped_out = torch.clamp(up_out, 0, 1)
+            psnr_val, ssim_val = psnr_calculator(clamped_out, gt_img), ssim_calculator(clamped_out, gt_img)
+            psnrs.update(torch.mean(psnr_val).item(), batch_size)
+            ssims.update(torch.mean(ssim_val).item(), batch_size)
+            torch.cuda.empty_cache()
+
+            if i % 100 == 0:
+                logging.info(
+                    "PSNR {:.4f}, SSIM {:.4f}, Elapse time {:.2f}\n".format(psnrs.avg, ssims.avg,
+                                                                            time.time() - start_time))
+
+        logging.info(f"Finish test at epoch {epoch}: avg PSNR: %.4f, avg SSIM: %.4F, and takes %.2f seconds" % (
+            psnrs.avg, ssims.avg, time.time() - start_time))
+
+
+def main(args):
+    global global_step
+
+    start_epoch = 1
+    global_step = 0
+
+    if not os.path.exists(args.output_dir):
+        os.mkdir(args.output_dir)
+
+    with open(os.path.join(args.output_dir, "args.json"), "w") as f:
+        json.dump(args.__dict__, f, sort_keys=True, indent=2)
+
+    log_format = "%(asctime)s %(levelname)-8s %(message)s"
+    log_file = os.path.join(args.output_dir, "train_log")
+    logging.basicConfig(filename=log_file, level=logging.INFO, format=log_format)
+    logging.getLogger().addHandler(logging.StreamHandler())
+
+    # device setting
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    args.device = device
+
+    logging.info(args.__dict__)
+
+    model = codebook_model(args)
+
+    optimizer = optim.Adam(model.parameters(), lr=args.optimizer["lr"],
+                           betas=(0.9, 0.999))
+
+    logging.info("Building data loader")
+
+    if args.train_loader["loader"] == "resize":
+        train_transforms = transforms.Compose([transforms.Resize(eval(args.train_loader["img_size"])),
+                                               transforms.ToTensor()])
+        train_loader = get_loader(args.data["train_dir"],
+                                  eval(args.train_loader["img_size"]), train_transforms, False,
+                                  int(args.train_loader["batch_size"]), args.train_loader["num_workers"],
+                                  args.train_loader["shuffle"], random_flag=False)
+
+    elif args.train_loader["loader"] == "crop":
+        train_loader = get_loader(args.data["train_dir"],
+                                  eval(args.train_loader["img_size"]), False, True,
+                                  int(args.train_loader["batch_size"]), args.train_loader["num_workers"],
+                                  args.train_loader["shuffle"], random_flag=args.train_loader["random_flag"])
+
+    elif args.train_loader["loader"] == "default":
+        train_transforms = transforms.Compose([transforms.ToTensor()])
+        train_loader = get_loader(args.data["train_dir"],
+                                  eval(args.train_loader["img_size"]), train_transforms, False,
+                                  int(args.train_loader["batch_size"]), args.train_loader["num_workers"],
+                                  args.train_loader["shuffle"], random_flag=args.train_loader["random_flag"])
+    else:
+        raise NotImplementedError
+
+    if args.test_loader["loader"] == "default":
+
+        test_transforms = transforms.Compose([transforms.ToTensor()])
+        test_loader = get_loader(args.data["test_dir"],
+                                 eval(args.test_loader["img_size"]), test_transforms, False,
+                                 int(args.test_loader["batch_size"]), args.test_loader["num_workers"],
+                                 args.test_loader["shuffle"], random_flag=False)
+
+    elif args.test_loader["loader"] == "resize":
+        test_transforms = transforms.Compose([transforms.Resize(eval(args.test_loader["img_size"])),
+                                              transforms.ToTensor()])
+        test_loader = get_loader(args.data["test_dir"],
+                                 eval(args.test_loader["img_size"]), test_transforms, False,
+                                 int(args.test_loader["batch_size"]), args.test_loader["num_workers"],
+                                 args.test_loader["shuffle"], random_flag=False)
+    else:
+        raise NotImplementedError
+
+    criterion = nn.SmoothL1Loss()
+    # vgg_loss = VGGLoss()
+
+    if args.optimizer["type"] == "cos":
+        lr_scheduler = CosineAnnealingWarmRestarts(optimizer, T_0=args.optimizer["T_0"],
+                                                   T_mult=args.optimizer["T_MULT"],
+                                                   eta_min=args.optimizer["ETA_MIN"],
+                                                   last_epoch=-1)
+    elif args.optimizer["type"] == "step":
+        lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=args.optimizer["step"],
+                                                       gamma=args.optimizer["gamma"])
+
+    t_total = int(len(train_loader) * args.optimizer["total_epoch"])
+    logging.info("***** CUDA.empty_cache() *****")
+    torch.cuda.empty_cache()
+
+    logging.info("***** Running training *****")
+    logging.info("  Batch size = %d", args.train_loader["batch_size"])
+    logging.info("  Num steps = %d", t_total)
+    logging.info("  Loader length = %d", len(train_loader))
+
+    model.train()
+    model.cuda()
+
+    logging.info("Begin training from epoch = %d\n", start_epoch)
+    # evaluate(model, "/home/yuwei/experiment/cvpr/prompt_final_vgg_gradient/models/model.1.bin", test_loader, 1)
+    for epoch in trange(start_epoch, args.optimizer["total_epoch"] + 1, desc="Epoch"):
+        train(model, train_loader, criterion, optimizer, epoch, args)
+        lr_scheduler.step()
+        if epoch % args.evaluate_intervel == 0:
+            logging.info("***** Running testing *****")
+            load_path = os.path.join(args.output_dir + "/models", "model.%d.bin" % (epoch))
+            evaluate(model, load_path, test_loader, epoch)
+            logging.info("***** End testing *****")
+
+
+if __name__ == '__main__':
+    parser = read_args("/home/yuwei/code/cvpr/config/LMAR_config.yaml")
+    args = parser.parse_args()
+    main(args)

LMAR_test.py

@@ -0,0 +1,98 @@
+import argparse
+import yaml
+import torchvision.transforms as transforms
+from utils import read_args, save_checkpoint, AverageMeter, CosineAnnealingWarmRestarts
+import time
+from tqdm import trange, tqdm
+from torchvision.utils import save_image
+import os
+os.environ['CUDA_VISIBLE_DEVICES'] = '0'
+import json
+import time
+import logging
+import torch
+from torch import nn, optim
+import numpy as np
+import torch.nn.functional as F
+
+import copy
+from model import *
+from data import *
+from PIL import Image
+from torch.optim import LBFGS
+import pyiqa
+from thop import profile
+from thop import clever_format
+
+from torchvision.models.feature_extraction import create_feature_extractor
+
+psnr_calculator = pyiqa.create_metric('psnr').cuda()
+ssim_calculator = pyiqa.create_metric('ssimc', downsample=True).cuda()
+
+
+def test(load_path, data_loader, args):
+    model = codebook_model(args)
+    checkpoint = torch.load(load_path)
+    model.load_state_dict(checkpoint["state_dict"])
+    model.cuda()
+    model.eval()
+
+    psnrs = AverageMeter()
+    ssims = AverageMeter()
+    lpipss = AverageMeter()
+    niqes = AverageMeter()
+    
+    down_size = (1440, 2560)
+    logging.info("Inference at down size: {}".format(down_size))
+    up_size = eval(args.test_loader["gt_size"])
+
+    start_time = time.time()
+    with torch.no_grad():
+        for i, batch in enumerate(tqdm(data_loader)):
+            inp_img, gt_img, inp_img_path = batch
+            inp_img = inp_img.cuda()
+            batch_size = inp_img.size(0)
+            gt_img = gt_img.cuda()
+            up_out = model(inp_img, down_size, up_size, test_flag=True)
+            name = inp_img_path[0].split("/")[-1]
+            # save_image(up_out[0], os.path.join(save_path, name))
+    
+            # metrics
+            clamped_out = torch.clamp(up_out, 0, 1)
+
+            psnr_val, ssim_val = psnr_calculator(clamped_out, gt_img), ssim_calculator(clamped_out, gt_img)
+            psnrs.update(psnr_val.item(), batch_size)
+            ssims.update(ssim_val.item(), batch_size)
+    
+            if i % 700 == 0:
+                logging.info(
+                    "PSNR {:.4f}, SSIM {:.4f}, LPIPS {:.4F}, NIQE {:.4F}, Elapse time {:.2f}\n".format(psnrs.avg, ssims.avg, lpipss.avg, niqes.avg,
+                                                                            time.time() - start_time))
+    
+        logging.info("Finish test: avg PSNR: %.4f, avg SSIM: %.4F, avg LPIPS: %.4F, avg NIQE: %.4F, and takes %.2f seconds" % (
+            psnrs.avg, ssims.avg, lpipss.avg, niqes.avg, time.time() - start_time))
+
+
+def main(args, load_path):
+    if not os.path.exists(args.output_dir):
+        os.mkdir(args.output_dir)
+    test_transforms = transforms.Compose([transforms.ToTensor()])
+
+    log_format = "%(asctime)s %(levelname)-8s %(message)s"
+    log_file = os.path.join(args.output_dir, "test_log")
+    logging.basicConfig(filename=log_file, level=logging.INFO, format=log_format)
+    logging.getLogger().addHandler(logging.StreamHandler())
+
+    logging.info("Building data loader")
+
+    test_loader = get_loader(args.data["test_dir"],
+                             eval(args.test_loader["img_size"]), test_transforms, False,
+                             int(args.test_loader["batch_size"]), args.test_loader["num_workers"],
+                             args.test_loader["shuffle"], random_flag=False)
+    test_time(load_path, test_loader, args)
+
+
+if __name__ == '__main__':
+    parser = read_args("/home/yuwei/code/cvpr/config/LMAR_config.yaml")
+    args = parser.parse_args()
+    main(args, "./pretrained_models\LMAR_model.bin")

base_test.py

@@ -0,0 +1,106 @@
+import argparse
+import yaml
+import torchvision.transforms as transforms
+from utils import read_args, save_checkpoint, AverageMeter, calculate_metrics, CosineAnnealingWarmRestarts
+# import torchvision.transforms.InterpolationMode
+import time
+from tqdm import trange, tqdm
+from torchvision.utils import save_image
+import os
+os.environ['CUDA_VISIBLE_DEVICES'] = '0'
+import json
+import time
+import logging
+import torch
+from torch import nn, optim
+import numpy as np
+import torch.nn.functional as F
+
+from model import *
+from data import *
+from PIL import Image
+from torchvision.transforms import Resize
+import pyiqa
+from thop import profile
+from thop import clever_format
+
+psnr_calculator = pyiqa.create_metric('psnr').cuda()
+ssim_calculator = pyiqa.create_metric('ssimc', downsample=True).cuda()
+lpips_calculator = pyiqa.create_metric('lpips').cuda()
+niqe_calculator = pyiqa.create_metric('niqe').cuda()
+
+
+def test(load_path, data_loader, args):
+    # if not os.path.exists(args.output_dir + '/out_my'):
+        # os.mkdir(args.output_dir + '/out_my')
+
+    # save_path = args.output_dir + "/out_my"
+    model = net(args)
+    checkpoint = torch.load(load_path)
+    model.load_state_dict(checkpoint["state_dict"])
+    model.cuda()
+    model.eval()
+
+    psnrs = AverageMeter()
+    ssims = AverageMeter()
+    lpipss = AverageMeter()
+    niqes = AverageMeter()
+    
+    start_time = time.time()
+    down_size = (1440, 2560)
+    logging.info("Inference at down size: {}".format(down_size))
+    with torch.no_grad():
+        for i, batch in enumerate(tqdm(data_loader)):
+            input_img, gt_img, inp_img_path = batch
+
+            name = inp_img_path[0].split("/")[-1]
+            input_img = input_img.cuda()
+            batch_size = input_img.size(0)
+            start_time = time.time()
+            input_img = resize(input_img, out_shape=down_size, antialiasing=False)
+            out_img = model(input_img)
+            out_img = resize(out_img, out_shape=eval(args.test_loader["gt_size"]), antialiasing=False)
+
+            # metrics
+            clamped_out = torch.clamp(out_img, 0, 1)
+            psnr_val, ssim_val = psnr_calculator(clamped_out, gt_img), ssim_calculator(clamped_out, gt_img)
+            psnrs.update(torch.mean(psnr_val).item(), batch_size)
+            ssims.update(torch.mean(ssim_val).item(), batch_size)
+
+            # lpips = lpips_calculator(clamped_out, gt_img)
+            # lpipss.update(torch.mean(lpips).item(), batch_size)
+            # niqe = niqe_calculator(clamped_out)
+            # niqes.update(torch.mean(niqe).item(), batch_size)
+            torch.cuda.empty_cache()
+
+            if i % 700 == 0:
+                logging.info(
+                    "PSNR {:.4f}, SSIM {:.4f}, LPIPS {:.4F}, NIQE {:.4F}, Elapse time {:.2f}\n".format(psnrs.avg, ssims.avg, lpipss.avg, niqes.avg,
+                                                                            time.time() - start_time))
+
+        logging.info("Finish test: avg PSNR: %.4f, avg SSIM: %.4F, avg LPIPS: %.4F, avg NIQE: %.4F, and takes %.2f seconds" % (
+            psnrs.avg, ssims.avg, lpipss.avg, niqes.avg, time.time() - start_time))
+
+def main(args, load_path):
+    if not os.path.exists(args.output_dir):
+        os.mkdir(args.output_dir)
+    test_transforms = transforms.Compose([transforms.ToTensor()])
+
+    log_format = "%(asctime)s %(levelname)-8s %(message)s"
+    log_file = os.path.join(args.output_dir, "baseline_log")
+    logging.basicConfig(filename=log_file, level=logging.INFO, format=log_format)
+    logging.getLogger().addHandler(logging.StreamHandler())
+
+    logging.info("Building data loader")
+
+    test_loader = get_loader(args.data["test_dir"],
+                             eval(args.test_loader["img_size"]), test_transforms, False,
+                             int(args.test_loader["batch_size"]), args.test_loader["num_workers"],
+                             args.test_loader["shuffle"], random_flag=False)
+    test(load_path, test_loader, args)
+
+
+if __name__ == '__main__':
+    parser = read_args("/home/yuwei/code/cvpr/config/base_config.yaml")
+    args = parser.parse_args()
+    main(args, "./pretrained_models/base_model.bin")

base_train.py

@@ -0,0 +1,261 @@
+import argparse
+import yaml
+import torchvision.transforms as transforms
+from utils import read_args, save_checkpoint, AverageMeter, calculate_metrics, CosineAnnealingWarmRestarts
+import time
+from tqdm import trange, tqdm
+from torchvision.utils import save_image
+# from tensorboardX import SummaryWriter
+import os
+import json
+import time
+import logging
+import torch
+from torch import nn, optim
+import torchvision.utils as vutils
+import torch.nn.functional as F
+
+from data import *
+from model import *
+from loss import *
+
+
+os.environ['CUDA_VISIBLE_DEVICES'] = '1'
+
+global_step = 0
+
+
+def train(model, data_loader, criterion, optimizer, epoch, args):
+    global global_step
+    iter_bar = tqdm(data_loader, desc='Iter (loss=X.XXX)')
+    nbatches = len(data_loader)
+
+    total_losses = AverageMeter()
+    pixel_losses = AverageMeter()
+    gradient_losses = AverageMeter()
+    psnrs = AverageMeter()
+    ssims = AverageMeter()
+
+    optimizer.zero_grad()
+
+    start_time = time.time()
+
+    if not os.path.exists(args.output_dir + '/image_train'):
+        os.mkdir(args.output_dir + '/image_train')
+
+    if not os.path.exists(args.output_dir + "/models"):
+        os.mkdir(args.output_dir + "/models")
+
+    for i, batch in enumerate(iter_bar):
+        optimizer.zero_grad()
+
+        input_img, gt_img, image_path = batch
+        input_img = input_img.cuda()
+        gt_img = gt_img.cuda()
+        batch_size = input_img.size(0)
+
+        out_img = model(input_img)
+
+        pixel_loss = criterion(out_img, gt_img)
+        pixel_losses.update(pixel_loss.item(), batch_size)
+
+        # gradient_loss = vggloss(out_img, gt_img).cuda()
+        # gradient_loss = args.hyper_params["x_lambda"] * gradient_loss
+        # gradient_losses.update(gradient_loss.item(), batch_size)
+
+        total_loss = pixel_loss
+        total_losses.update(total_loss.item(), batch_size)
+
+        total_loss.backward()
+        optimizer.step()
+
+        iter_bar.set_description('Iter (loss=%5.6f)' % total_losses.avg)
+
+        if i % 200 == 0:
+            saved_image = torch.cat([input_img[0:2], out_img[0:2], gt_img[0:2]], dim=0)
+            save_image(saved_image, args.output_dir + '/image_train/epoch_{}_iter_{}.jpg'.format(epoch, i))
+
+        # metrics
+        norm_out = torch.clamp(out_img, 0, 1)
+        #psnr_val, ssim_val = calculate_metrics(norm_out, gt_img)
+        #psnrs.update(psnr_val.item(), batch_size)
+        #ssims.update(ssim_val.item(), batch_size)
+
+        if i % max(1, nbatches // 10) == 0:
+            logging.info(
+                "Epoch {}, learning rates {:}, Iter {}, total_loss {:.4f}, pixel_loss {:.4f}, PSNR {:.4f}, SSIM {:.4f}, Elapse time {:.2f}\n".format(
+                    epoch, optimizer.param_groups[0]["lr"], i, total_losses.avg, pixel_losses.avg,
+                    psnrs.avg, ssims.avg,
+                    time.time() - start_time))
+
+    if epoch % 1 == 0:
+        logging.info("** ** * Saving model and optimizer ** ** * ")
+
+        output_model_file = os.path.join(args.output_dir + "/models", "model.%d.bin" % (epoch))
+        state = {"epoch": epoch, "state_dict": model.state_dict(),
+                 "optimizer": optimizer.state_dict(), "step": global_step}
+
+        save_checkpoint(state, output_model_file)
+        logging.info("Save model to %s", output_model_file)
+
+    logging.info(
+        "Finish training epoch %d, avg total_loss: %.4f, avg pixel_loss: %.4f, avg PSNR: %.2f, avg SSIM: %.2F, and takes %.2f seconds" % (
+            epoch, total_losses.avg, pixel_losses.avg, psnrs.avg, ssims.avg,
+            time.time() - start_time))
+
+    logging.info("***** CUDA.empty_cache() *****\n")
+    torch.cuda.empty_cache()
+
+
+def evaluate(model, load_path, data_loader, epoch):
+
+    checkpoint = torch.load(load_path)
+    model.load_state_dict(checkpoint["state_dict"])
+    model.cuda()
+    model.eval()
+
+    psnrs = AverageMeter()
+    ssims = AverageMeter()
+
+    start_time = time.time()
+    with torch.no_grad():
+        for i, batch in enumerate(tqdm(data_loader)):
+            input_img, gt_img, inp_img_path = batch
+            input_img = input_img.cuda()
+            batch_size = input_img.size(0)
+            out_img = model(input_img)
+
+            # metrics
+            norm_out = torch.clamp(out_img, 0, 1)
+            psnr_val, ssim_val = calculate_metrics(norm_out, gt_img)
+            psnrs.update(psnr_val.item(), batch_size)
+            ssims.update(ssim_val.item(), batch_size)
+            torch.cuda.empty_cache()
+
+            if i % 100 == 0:
+                logging.info(
+                    "PSNR {:.4f}, SSIM {:.4f}, Elapse time {:.2f}\n".format(psnrs.avg, ssims.avg,
+                                                                            time.time() - start_time))
+
+        logging.info(f"Finish test at epoch {epoch}: avg PSNR: %.4f, avg SSIM: %.4F, and takes %.2f seconds" % (
+            psnrs.avg, ssims.avg, time.time() - start_time))
+
+
+def main(args):
+    global global_step
+
+    start_epoch = 1
+    global_step = 0
+
+    if not os.path.exists(args.output_dir):
+        os.mkdir(args.output_dir)
+
+    with open(os.path.join(args.output_dir, "args.json"), "w") as f:
+        json.dump(args.__dict__, f, sort_keys=True, indent=2)
+
+    log_format = "%(asctime)s %(levelname)-8s %(message)s"
+    log_file = os.path.join(args.output_dir, "train_log")
+    logging.basicConfig(filename=log_file, level=logging.INFO, format=log_format)
+    logging.getLogger().addHandler(logging.StreamHandler())
+
+    # device setting
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    args.device = device
+
+    logging.info(args.__dict__)
+
+    if args.resume["flag"]:
+        model = net(args)
+        model.to(args.device)
+        check_point = torch.load(args.resume["checkpoint"])
+        model.load_state_dict(check_point["state_dict"])
+        optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=args.optimizer["lr"],
+                               betas=(0.9, 0.999))
+        optimizer.load_state_dict(check_point["optimizer"])
+        start_epoch = check_point["epoch"] + 1
+        # start_epoch = check_point["epoch"]
+
+    else:
+        model = net(args)
+        model.to(args.device)
+        optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=args.optimizer["lr"],
+                               betas=(0.9, 0.999))
+
+    logging.info("Building data loader")
+
+    if args.train_loader["loader"] == "resize":
+        train_transforms = transforms.Compose([transforms.Resize(eval(args.train_loader["img_size"])),
+                                               transforms.ToTensor()])
+        train_loader = get_loader(args.data["train_dir"],
+                                  eval(args.train_loader["img_size"]), train_transforms, False,
+                                  int(args.train_loader["batch_size"]), args.train_loader["num_workers"],
+                                  args.train_loader["shuffle"], inference_flag=False)
+
+    elif args.train_loader["loader"] == "crop":
+        train_loader = get_loader(args.data["train_dir"],
+                                  eval(args.train_loader["img_size"]), False, True,
+                                  int(args.train_loader["batch_size"]), args.train_loader["num_workers"],
+                                  args.train_loader["shuffle"], inference_flag=False)
+    else:
+        raise NotImplementedError
+
+    if args.test_loader["loader"] == "default":
+
+        test_transforms = transforms.Compose([transforms.ToTensor()])
+        test_loader = get_loader(args.data["test_dir"],
+                                  eval(args.test_loader["img_size"]), test_transforms, False,
+                                  int(args.test_loader["batch_size"]), args.test_loader["num_workers"],
+                                  args.test_loader["shuffle"], inference_flag=False)
+    
+    elif args.test_loader["loader"] == "resize":
+
+        test_transforms = transforms.Compose([transforms.Resize(eval(args.test_loader["img_size"])),
+                                               transforms.ToTensor()])
+        test_loader = get_loader(args.data["test_dir"],
+                                  eval(args.test_loader["img_size"]), test_transforms, False,
+                                  int(args.test_loader["batch_size"]), args.test_loader["num_workers"],
+                                  args.test_loader["shuffle"], inference_flag=False)
+
+    criterion = nn.L1Loss()
+    # vgg_loss = VGGLoss()
+
+    if args.optimizer["type"] == "cos":
+        lr_scheduler = CosineAnnealingWarmRestarts(optimizer, T_0=args.optimizer["T_0"],
+                                                   T_mult=args.optimizer["T_MULT"],
+                                                   eta_min=args.optimizer["ETA_MIN"],
+                                                   last_epoch=-1)
+    elif args.optimizer["type"] == "step":
+        lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=args.optimizer["step"],
+                                                       gamma=args.optimizer["gamma"])
+
+    if args.resume["flag"]:
+        for i in range(start_epoch):
+            lr_scheduler.step()
+
+    t_total = int(len(train_loader) * args.optimizer["total_epoch"])
+    logging.info("***** CUDA.empty_cache() *****")
+    torch.cuda.empty_cache()
+
+    logging.info("***** Running training *****")
+    logging.info("  Batch size = %d", args.train_loader["batch_size"])
+    logging.info("  Num steps = %d", t_total)
+    logging.info("  Loader length = %d", len(train_loader))
+
+    model.train()
+    model.cuda()
+
+    logging.info("Begin training from epoch = %d\n", start_epoch)
+    for epoch in trange(start_epoch, args.optimizer["total_epoch"] + 1, desc="Epoch"):
+        train(model, train_loader, criterion, optimizer, epoch, args)
+        lr_scheduler.step()
+        if epoch % args.evaluate_intervel == 0:
+            logging.info("***** Running testing *****")
+            load_path = os.path.join(args.output_dir + "/models", "model.%d.bin" % (epoch))
+            evaluate(model, load_path, test_loader, epoch)
+            logging.info("***** End testing *****")
+
+
+if __name__ == '__main__':
+    parser = read_args("/home/yuwei/code/cvpr/config/base_config.yaml")
+    args = parser.parse_args()
+    main(args)

config/LMAR_config.yaml

@@ -0,0 +1,48 @@
+output_dir: '/home/yuwei/experiment/cvpr/LMAR_cubic'
+data:
+  train_dir: /home/data/yuwei/data/uhd4k_ll/train
+  test_dir: /home/data/yuwei/data/uhd4k_ll/test
+
+model:
+  in_channel: 3
+  model_channel: 8
+  sparsity_threshold: 0.01
+  num_blocks: 8
+  threslhold_frac: 0.6
+  hidden_channel: 48
+
+train_loader:
+  num_workers: 8
+  batch_size: 1
+  loader: crop
+  img_size: (1024, 1024)
+  shuffle: True
+  gt_size: (2160, 3840)
+  random_flag: True
+
+test_loader:
+  num_workers: 8
+  batch_size: 1
+  loader: default
+  img_size: ((1440, 2560), (1080, 1920), (1200, 1600), (720, 1280), (540, 960))
+  shuffle: False
+  gt_size: (2160, 3840)
+
+optimizer:
+  type: step
+  total_epoch: 12
+  lr: 0.0004
+  T_0: 0.00001
+  T_MULT: 1
+  ETA_MIN: 0.000001
+  step: 4
+  gamma: 0.75
+
+hyper_params:
+  lambda: 0.5
+
+resume:
+  flag: True
+  checkpoint: ./pretrained_models/base_model.bin
+
+evaluate_intervel: 1

config/base_config.yaml

@@ -0,0 +1,42 @@
+output_dir: '/home/yuwei/experiment/cvpr/uhd4k_ll_pretrain'
+data:
+  train_dir: /home/data/yuwei/data/uhd4k_ll/train
+  test_dir: /home/data/yuwei/data/uhd4k_ll/test
+
+model:
+  in_channel: 3
+  model_channel: 8
+
+train_loader:
+  num_workers: 8
+  batch_size: 2
+  loader: resize
+  img_size: (1024, 1024)
+  shuffle: True
+
+test_loader:
+  num_workers: 8
+  batch_size: 1
+  loader: default
+  img_size: (1200, 1600)
+  shuffle: False
+  gt_size: (2160, 3840)
+
+optimizer:
+  type: step
+  total_epoch: 100
+  lr: 0.001
+  T_0: 100
+  T_MULT: 1
+  ETA_MIN: 0.000001
+  step: 20
+  gamma: 0.75
+
+hyper_params:
+  x_lambda: 0.03
+
+resume:
+  flag: False
+  checkpoint: Null
+
+evaluate_intervel: 5

data/__init__.py

@@ -0,0 +1 @@
+from .loader import *

data/loader.py

@@ -0,0 +1,109 @@
+import os
+import numpy as np
+from tqdm import tqdm
+from torch.utils.data import DataLoader
+from torch.utils.data.dataset import Dataset
+from PIL import Image
+import torchvision.transforms.functional as TF
+import torchvision.transforms as tf
+from PIL import Image, ImageFile
+import random
+import math
+from model import *
+import torch
+# import cv2
+# cv2.setNumThreads(0)
+
+ImageFile.LOAD_TRUNCATED_IMAGES = True
+
+
+class base_dataset(Dataset):
+    def __init__(self, data_dir, img_size, transforms=False, crop=False):
+        imgs = sorted(os.listdir(data_dir + "/input"))
+        self.input_imgs = [os.path.join(data_dir + "/input", name) for name in imgs]
+        self.gt_imgs = [os.path.join(data_dir + "/gt", name) for name in imgs]
+        self.transforms = transforms
+        self.crop = crop
+        self.img_size = img_size
+
+    def __getitem__(self, index):
+        inp_img_path = self.input_imgs[index]
+        gt_img_path = self.gt_imgs[index]
+        inp_img = Image.open(inp_img_path).convert("RGB")
+        gt_img = Image.open(gt_img_path).convert("RGB")
+        if self.transforms:
+            inp_img = self.transforms(inp_img)
+            gt_img = self.transforms(gt_img)
+
+        if self.crop:
+            inp_img, gt_img = self.crop_image(inp_img, gt_img)
+
+        return inp_img, gt_img, inp_img_path
+
+    def __len__(self):
+        return len(self.gt_imgs)
+
+    def crop_image(self, inp_img, gt_img):
+        crop_h, crop_w = self.img_size
+        i, j, h, w = tf.RandomCrop.get_params(
+            inp_img, output_size=((crop_h, crop_w)))
+        inp_img = TF.crop(inp_img, i, j, h, w)
+        gt_img = TF.crop(gt_img, i, j, h, w)
+        inp_img = TF.to_tensor(inp_img)
+        gt_img = TF.to_tensor(gt_img)
+
+        return inp_img, gt_img
+
+
+class random_scale_dataset(Dataset):
+    def __init__(self, data_dir, img_size, transforms=False, crop=False):
+        imgs = sorted(os.listdir(data_dir + "/input"))
+        self.input_imgs = [os.path.join(data_dir + "/input", name) for name in imgs]
+        self.gt_imgs = [os.path.join(data_dir + "/gt", name) for name in imgs]
+        self.transforms = transforms
+        self.crop = crop
+        self.img_size = img_size
+
+    def __getitem__(self, index):
+        inp_img_path = self.input_imgs[index]
+        gt_img_path = self.gt_imgs[index]
+        inp_img = Image.open(inp_img_path).convert("RGB")
+        gt_img = Image.open(gt_img_path).convert("RGB")
+
+        random_scale_factor = random.randrange(self.img_size[0] * 0.25, self.img_size[0], 8)
+        down_h = down_w = random_scale_factor
+
+        if self.transforms:
+            inp_img = self.transforms(inp_img)
+            gt_img = self.transforms(gt_img)
+            return inp_img, gt_img, down_h, down_w, inp_img_path
+
+        if self.crop:
+            inp_img, gt_img = self.crop_image(inp_img, gt_img)
+            return inp_img, gt_img, down_h, down_w, inp_img_path
+
+    def __len__(self):
+        return len(self.gt_imgs)
+
+    def crop_image(self, inp_img, gt_img):
+        crop_h, crop_w = self.img_size
+        i, j, h, w = tf.RandomCrop.get_params(
+            inp_img, output_size=((crop_h, crop_w)))
+        inp_img = TF.crop(inp_img, i, j, h, w)
+        gt_img = TF.crop(gt_img, i, j, h, w)
+        inp_img = TF.to_tensor(inp_img)
+        gt_img = TF.to_tensor(gt_img)
+
+        return inp_img, gt_img
+
+
+def get_loader(data_dir, img_size, transforms, crop_flag, batch_size, num_workers, shuffle, random_flag=False, inference_flag=False):
+    if random_flag:
+        dataset = random_scale_dataset(data_dir, img_size, transforms, crop_flag)
+        dataloader = DataLoader(dataset, batch_size=batch_size,
+                                shuffle=shuffle, num_workers=num_workers, pin_memory=True)
+    else:
+        dataset = base_dataset(data_dir, img_size, transforms, crop_flag)
+        dataloader = DataLoader(dataset, batch_size=batch_size,
+                                shuffle=shuffle, num_workers=num_workers, pin_memory=True)
+    return dataloader

loss.py

@@ -0,0 +1,221 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision
+
+
+class VGG19(torch.nn.Module):
+    def __init__(self, requires_grad=False):
+        super().__init__()
+        vgg_pretrained_features = torchvision.models.vgg19(pretrained=True).features
+        self.slice1 = torch.nn.Sequential()
+        self.slice2 = torch.nn.Sequential()
+        self.slice3 = torch.nn.Sequential()
+        self.slice4 = torch.nn.Sequential()
+        self.slice5 = torch.nn.Sequential()
+        for x in range(2):
+            self.slice1.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(2, 7):
+            self.slice2.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(7, 12):
+            self.slice3.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(12, 21):
+            self.slice4.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(21, 30):
+            self.slice5.add_module(str(x), vgg_pretrained_features[x])
+        if not requires_grad:
+            for param in self.parameters():
+                param.requires_grad = False
+
+    def forward(self, X):
+        h_relu1 = self.slice1(X)
+        h_relu2 = self.slice2(h_relu1)
+        h_relu3 = self.slice3(h_relu2)
+        h_relu4 = self.slice4(h_relu3)
+        h_relu5 = self.slice5(h_relu4)
+        out = [h_relu1, h_relu2, h_relu3, h_relu4, h_relu5]
+        return out
+
+
+class VGGLoss(nn.Module):
+    def __init__(self):
+        super(VGGLoss, self).__init__()
+        self.vgg = VGG19().cuda()
+        self.criterion = nn.L1Loss()
+        self.weights = [1.0 / 32, 1.0 / 16, 1.0 / 8, 1.0 / 4, 1.0]
+
+    def forward(self, x, y):
+        x_vgg, y_vgg = self.vgg(x), self.vgg(y)
+        loss = 0
+        for i in range(len(x_vgg)):
+            loss += self.weights[i] * self.criterion(x_vgg[i], y_vgg[i].detach())
+        return loss
+
+
+class VGGPerceptualLoss(torch.nn.Module):
+    def __init__(self, lam=1, lam_p=1):
+        super(VGGPerceptualLoss, self).__init__()
+        self.loss_fn = VGGPerceptualLoss()
+
+    def forward(self, out, gt):
+        loss = self.loss_fn(out, gt, feature_layers=[2])
+
+        return loss
+
+
+class VGGPerceptualLoss(torch.nn.Module):
+    def __init__(self, resize=True):
+        super(VGGPerceptualLoss, self).__init__()
+        blocks = []
+        blocks.append(torchvision.models.vgg16(pretrained=True).features[:4].eval())
+        blocks.append(torchvision.models.vgg16(pretrained=True).features[4:9].eval())
+        blocks.append(torchvision.models.vgg16(pretrained=True).features[9:16].eval())
+        blocks.append(torchvision.models.vgg16(pretrained=True).features[16:23].eval())
+        for bl in blocks:
+            for p in bl:
+                p.requires_grad = False
+        self.blocks = torch.nn.ModuleList(blocks).cuda()
+        self.transform = torch.nn.functional.interpolate
+        self.mean = torch.nn.Parameter(torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1)).cuda()
+        self.std = torch.nn.Parameter(torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1)).cuda()
+        self.resize = resize
+
+    def forward(self, input, target, feature_layers=[0, 1, 2, 3], style_layers=[]):
+        if input.shape[1] != 3:
+            input = input.repeat(1, 3, 1, 1)
+            target = target.repeat(1, 3, 1, 1)
+        input = (input - self.mean) / self.std
+        target = (target - self.mean) / self.std
+        if self.resize:
+            input = self.transform(input, mode='bilinear', size=(224, 224), align_corners=False)
+            target = self.transform(target, mode='bilinear', size=(224, 224), align_corners=False)
+        loss = 0.0
+        x = input
+        y = target
+        for i, block in enumerate(self.blocks):
+            x = block(x)
+            y = block(y)
+            if i in feature_layers:
+                loss += torch.nn.functional.l1_loss(x, y)
+            if i in style_layers:
+                act_x = x.reshape(x.shape[0], x.shape[1], -1)
+                act_y = y.reshape(y.shape[0], y.shape[1], -1)
+                gram_x = act_x @ act_x.permute(0, 2, 1)
+                gram_y = act_y @ act_y.permute(0, 2, 1)
+                loss += torch.nn.functional.l1_loss(gram_x, gram_y)
+        return loss
+
+
+def scharr(x):  # 输入前对RGB通道求均值在灰度图上算
+    b, c, h, w = x.shape
+    pad = nn.ReplicationPad2d(padding=(1, 1, 1, 1))
+    x = pad(x)
+    kx = F.unfold(x, kernel_size=3, stride=1, padding=0)  # b,n*k*k,n_H*n_W
+    kx = kx.permute([0, 2, 1])  # b,n_H*n_W,n*k*k
+    # kx=kx.view(1, b*h*w, 9) #1,b*n_H*n_W,n*k*k
+
+    w1 = torch.tensor([-3, 0, 3, -10, 0, 10, -3, 0, 3]).float().cuda()
+    w2 = torch.tensor([-3, -10, -3, 0, 0, 0, 3, 10, 3]).float().cuda()
+
+    y1 = torch.matmul((kx * 255.0), w1)  # 1,b*n_H*n_W,1
+    y2 = torch.matmul((kx * 255.0), w2)  # 1,b*n_H*n_W,1
+    # y1=y1.view(b,h*w,1) #b,n_H*n_W,1
+    y1 = y1.unsqueeze(-1).permute([0, 2, 1])  # b,1,n_H*n_W
+    # y2=y2.view(b,h*w,1) #b,n_H*n_W,1
+    y2 = y2.unsqueeze(-1).permute([0, 2, 1])  # b,1,n_H*n_W
+
+    y1 = F.fold(y1, output_size=(h, w), kernel_size=1)  # b,m,n_H,n_W
+    y2 = F.fold(y2, output_size=(h, w), kernel_size=1)  # b,m,n_H,n_W
+    y1 = y1.clamp(-255, 255)
+    y2 = y2.clamp(-255, 255)
+    return (0.5 * torch.abs(y1) + 0.5 * torch.abs(y2)) / 255.0
+
+
+def gram_matrix(input):
+    a, b, c, d = input.size()  # a=batch size(=1)
+    # b=number of feature maps
+    # (c,d)=dimensions of a f. map (N=c*d)
+
+    features = input.reshape(a * b, c * d)  # resize F_XL into \hat F_XL
+
+    G = torch.mm(features, features.t())  # compute the gram product
+
+    # we 'normalize' the values of the gram matrix
+    # by dividing by the number of element in each feature maps.
+    return G.div(a * b * c * d)
+
+
+class StyleLoss(nn.Module):
+    def __init__(self):
+        super(StyleLoss, self).__init__()
+
+    def forward(self, input_fea, target_fea):
+        target = gram_matrix(target_fea).detach()
+        G = gram_matrix(input_fea)
+        loss = F.mse_loss(G, target)
+        return loss
+
+
+def cos_loss(feat1, feat2):
+    # maximize average cosine similarity
+    return -F.cosine_similarity(feat1, feat2).mean()
+
+
+def feat_scharr(x):
+    x = torch.mean(x, dim=1, keepdim=True)
+    x = (x - x.min()) / (x.max() - x.min())
+    x = x * 255
+    return scharr(x)
+
+
+def feat_ssim(feat1, feat2, gt):
+    mask = scharr(torch.mean(gt, dim=1, keepdim=True))
+    # mask = torch.nn.MaxPool2d(5, 1, 2)(mask)
+    mask = F.interpolate(mask, size=(feat1.shape[2], feat1.shape[3]), mode="bicubic")
+    loss = torch.abs(feat1 - feat2) * mask
+    return torch.mean(loss), mask
+
+
+def similarity_loss(f_s, f_t):
+    def at(f):
+        return F.normalize(f.pow(2).mean(1).view(f.size(0), -1))
+
+    return (at(f_s) - at(f_t)).pow(2).mean()
+
+
+class RBF(nn.Module):
+
+    def __init__(self, n_kernels=5, mul_factor=2.0, bandwidth=None):
+        super().__init__()
+        self.bandwidth_multipliers = mul_factor ** (torch.arange(n_kernels) - n_kernels // 2)
+        self.bandwidth = bandwidth
+
+    def get_bandwidth(self, L2_distances):
+        if self.bandwidth is None:
+            n_samples = L2_distances.shape[0]
+            return L2_distances.data.sum() / (n_samples ** 2 - n_samples)
+
+        return self.bandwidth
+
+    def forward(self, X):
+        L2_distances = torch.cdist(X, X) ** 2
+
+        return torch.exp(
+            -L2_distances[None, ...].cuda() / (self.get_bandwidth(L2_distances).cuda() * self.bandwidth_multipliers.cuda())[:, None,
+                                       None]).sum(dim=0)
+
+
+class MMDLoss(nn.Module):
+
+    def __init__(self, kernel=RBF()):
+        super().__init__()
+        self.kernel = kernel.cuda()
+
+    def forward(self, X, Y):
+        K = self.kernel(torch.vstack([X, Y]))
+
+        X_size = X.shape[0]
+        XX = K[:X_size, :X_size].mean()
+        XY = K[:X_size, X_size:].mean()
+        YY = K[X_size:, X_size:].mean()
+        return XX - 2 * XY + YY

metrics.py

@@ -0,0 +1,133 @@
+import cv2
+import numpy as np
+from PIL import Image
+import torchvision.transforms as transforms
+
+
+def calculate_psnr(img, img2, crop_border, input_order='HWC', test_y_channel=False, **kwargs):
+    """Calculate PSNR (Peak Signal-to-Noise Ratio).
+
+    Ref: https://en.wikipedia.org/wiki/Peak_signal-to-noise_ratio
+
+    Args:
+        img (ndarray): Images with range [0, 255].
+        img2 (ndarray): Images with range [0, 255].
+        crop_border (int): Cropped pixels in each edge of an image. These
+            pixels are not involved in the PSNR calculation.
+        input_order (str): Whether the input order is 'HWC' or 'CHW'.
+            Default: 'HWC'.
+        test_y_channel (bool): Test on Y channel of YCbCr. Default: False.
+
+    Returns:
+        float: psnr result.
+    """
+
+    assert img.shape == img2.shape, (f'Image shapes are different: {img.shape}, {img2.shape}.')
+    if input_order not in ['HWC', 'CHW']:
+        raise ValueError(f'Wrong input_order {input_order}. Supported input_orders are ' '"HWC" and "CHW"')
+    img = img.astype(np.float64)
+    img2 = img2.astype(np.float64)
+
+    if crop_border != 0:
+        img = img[crop_border:-crop_border, crop_border:-crop_border, ...]
+        img2 = img2[crop_border:-crop_border, crop_border:-crop_border, ...]
+
+    if test_y_channel:
+        img = to_y_channel(img)
+        img2 = to_y_channel(img2)
+
+    mse = np.mean((img - img2)**2)
+    if mse == 0:
+        return float('inf')
+    return 20. * np.log10(255. / np.sqrt(mse))
+
+
+def _ssim(img, img2):
+    """Calculate SSIM (structural similarity) for one channel images.
+
+    It is called by func:`calculate_ssim`.
+
+    Args:
+        img (ndarray): Images with range [0, 255] with order 'HWC'.
+        img2 (ndarray): Images with range [0, 255] with order 'HWC'.
+
+    Returns:
+        float: ssim result.
+    """
+
+    c1 = (0.01 * 255)**2
+    c2 = (0.03 * 255)**2
+
+    img = img.astype(np.float64)
+    img2 = img2.astype(np.float64)
+    kernel = cv2.getGaussianKernel(11, 1.5)
+    window = np.outer(kernel, kernel.transpose())
+
+    mu1 = cv2.filter2D(img, -1, window)[5:-5, 5:-5]
+    mu2 = cv2.filter2D(img2, -1, window)[5:-5, 5:-5]
+    mu1_sq = mu1**2
+    mu2_sq = mu2**2
+    mu1_mu2 = mu1 * mu2
+    sigma1_sq = cv2.filter2D(img**2, -1, window)[5:-5, 5:-5] - mu1_sq
+    sigma2_sq = cv2.filter2D(img2**2, -1, window)[5:-5, 5:-5] - mu2_sq
+    sigma12 = cv2.filter2D(img * img2, -1, window)[5:-5, 5:-5] - mu1_mu2
+
+    ssim_map = ((2 * mu1_mu2 + c1) * (2 * sigma12 + c2)) / ((mu1_sq + mu2_sq + c1) * (sigma1_sq + sigma2_sq + c2))
+    return ssim_map.mean()
+
+def calculate_ssim(img, img2, crop_border, input_order='HWC', test_y_channel=False, **kwargs):
+    """Calculate SSIM (structural similarity).
+
+    Ref:
+    Image quality assessment: From error visibility to structural similarity
+
+    The results are the same as that of the official released MATLAB code in
+    https://ece.uwaterloo.ca/~z70wang/research/ssim/.
+
+    For three-channel images, SSIM is calculated for each channel and then
+    averaged.
+
+    Args:
+        img (ndarray): Images with range [0, 255].
+        img2 (ndarray): Images with range [0, 255].
+        crop_border (int): Cropped pixels in each edge of an image. These
+            pixels are not involved in the SSIM calculation.
+        input_order (str): Whether the input order is 'HWC' or 'CHW'.
+            Default: 'HWC'.
+        test_y_channel (bool): Test on Y channel of YCbCr. Default: False.
+
+    Returns:
+        float: ssim result.
+    """
+
+    assert img.shape == img2.shape, (f'Image shapes are different: {img.shape}, {img2.shape}.')
+    if input_order not in ['HWC', 'CHW']:
+        raise ValueError(f'Wrong input_order {input_order}. Supported input_orders are ' '"HWC" and "CHW"')
+    # img = reorder_image(img, input_order=input_order)
+    # img2 = reorder_image(img2, input_order=input_order)
+    img = img.astype(np.float64)
+    img2 = img2.astype(np.float64)
+
+    if crop_border != 0:
+        img = img[crop_border:-crop_border, crop_border:-crop_border, ...]
+        img2 = img2[crop_border:-crop_border, crop_border:-crop_border, ...]
+
+    if test_y_channel:
+        img = to_y_channel(img)
+        img2 = to_y_channel(img2)
+
+    ssims = []
+    for i in range(img.shape[2]):
+        ssims.append(_ssim(img[..., i], img2[..., i]))
+    return np.array(ssims).mean()
+
+if __name__ == '__main__':
+
+    # test_transforms = transforms.Compose([transforms.Resize((512, 512)),transforms.ToTensor()])
+    # inp_img = Image.open("/mnt/disk1/yuwei/data/4Kdehaze/train/clear/0_000002.jpg").convert("RGB")
+    # img = test_transforms(inp_img)
+    img = cv2.imread("/mnt/disk1/yuwei/data/4Kdehaze/train/clear/0_000002.jpg")
+    psnr = calculate_psnr(img, img, 0)
+    ssim = calculate_ssim(img, img, 0)
+    print(psnr)
+    print(ssim)

model/LMAR_model.py

@@ -0,0 +1,277 @@
+from model import net
+import torch.nn as nn
+import torch
+import torch.nn.functional as F
+from torchvision.transforms import Resize
+
+try:
+    from resize_right import resize
+except:
+    from .resize_right import resize
+
+try:
+    from .interp_methods import *
+except:
+    from interp_methods import *
+
+from torchvision.models import vgg19
+from torchvision.models.feature_extraction import create_feature_extractor
+
+import tinycudann as tcnn
+from torchvision.utils import save_image
+import torchvision.transforms as transforms
+from torchviz import make_dot
+
+
+def make_coord(shape, ranges=None, flatten=True):
+    """ Make coordinates at grid centers.
+    """
+    coord_seqs = []
+    for i, n in enumerate(shape):
+        if ranges is None:
+            v0, v1 = -1, 1
+        else:
+            v0, v1 = ranges[i]
+        r = (v1 - v0) / (2 * n)
+        seq = v0 + r + (2 * r) * torch.arange(n).float()
+        coord_seqs.append(seq)
+    ret = torch.stack(torch.meshgrid(*coord_seqs), dim=-1)
+    if flatten:
+        ret = ret.view(-1, ret.shape[-1])
+    return ret
+
+def get_local_grid(img):
+    local_grid = make_coord(img.shape[-2:], flatten=False).cuda()
+    local_grid = local_grid.permute(2, 0, 1).unsqueeze(0)
+    local_grid = local_grid.expand(img.shape[0], 2, *img.shape[-2:])
+
+    return local_grid
+
+def creat_coord(x):
+    b = x.shape[0]
+    coord = make_coord(x.shape[-2:], flatten=False)
+    coord = coord.permute(2, 0, 1).contiguous().unsqueeze(0)
+    coord = coord.expand(b, 2, *coord.shape[-2:])
+
+    coord_ = coord.clone()
+    coord_ = coord_.clamp_(-1 + 1e-6, 1 - 1e-6)
+    coord_ = coord_.permute(0, 2, 3, 1).contiguous()
+    coord_ = coord_.view(b, -1, coord.size(1))
+    return coord.cuda(), coord_.cuda()
+
+
+def get_cell(img, local_grid):
+    cell = torch.ones_like(local_grid)
+    cell[:, 0] *= 2 / img.size(2)
+    cell[:, 1] *= 2 / img.size(3)
+
+    return cell
+
+
+class TcnnFCBlock(tcnn.Network):
+    def __init__(
+            self, in_features, out_features,
+            num_hidden_layers, hidden_features,
+            activation: str = 'LeakyRelu', last_activation: str = 'None',
+            seed=42):
+        assert hidden_features in [16, 32, 64, 128], "hidden_features can only be 16, 32, 64, or 128."
+        super().__init__(in_features, out_features, network_config={
+            "otype": "FullyFusedMLP",  # Component type.
+            "activation": activation,  # Activation of hidden layers.
+            "output_activation": last_activation,  # Activation of the output layer.
+            "n_neurons": hidden_features,  # Neurons in each hidden layer. # May only be 16, 32, 64, or 128.
+            "n_hidden_layers": num_hidden_layers,  # Number of hidden layers.
+        }, seed=seed)
+
+    def forward(self, x: torch.Tensor):
+        prefix = x.shape[:-1]
+        return super().forward(x.flatten(0, -2)).unflatten(0, prefix)
+
+
+class LMAR_model(nn.Module):
+    def __init__(self, args):
+        super().__init__()
+        self.resume_flag = args.resume["flag"]
+        self.load_path = args.resume["checkpoint"]
+
+        if self.resume_flag and self.load_path:
+            self.model = net(args)
+            checkpoint = torch.load(self.load_path)
+            self.model.load_state_dict(checkpoint["state_dict"])
+            for param in self.model.parameters():
+                param.requires_grad_(False)
+
+        self.in_channel = 3
+        self.out_channel = 3
+        self.kernel_size = 3
+        self.imnet = TcnnFCBlock(7, self.in_channel * self.out_channel * self.kernel_size * self.kernel_size, 5,
+                                 128).cuda()
+        self.mid_nodes = {"hr_backbone.skip2": "bottom"}
+        self.extractor_mid = create_feature_extractor(self.model, self.mid_nodes)
+        self.modulation = nn.Conv2d(6, 3, 1, 1, 0)
+        # self.projection = nn.Conv2d()
+
+    def forward(self, x, down_size, up_size, test_flag=False):
+        if test_flag:
+            up_out, _ = self.inference(x, down_size, up_size)
+            return up_out, _ 
+        else:
+            down_x, hr_feature, new_lr_feature, ori_lr_feature, residual, res = self.train_model(x, down_size, up_size)
+            return down_x, hr_feature, new_lr_feature, ori_lr_feature, residual, res
+
+    def train_model(self, x, down_size, up_size):
+        # down_sizer = transforms.Resize(size=down_size,
+        #                                interpolation=transforms.InterpolationMode.BILINEAR)
+        # up_sizer = transforms.Resize(size=up_size,
+        #                              interpolation=transforms.InterpolationMode.BILINEAR)
+
+        b = x.shape[0]
+        # down_x = down_sizer(x)
+        down_x = resize(x, out_shape=down_size, antialiasing=False)
+        # down_x = resize(x, out_shape=down_size, antialiasing=True)
+
+        hr_feature = self.extractor_mid(x)["bottom"]
+        # feature_sizer = transforms.Resize(size=(hr_feature.shape[2], hr_feature.shape[3]),
+        #                                   interpolation=transforms.InterpolationMode.BILINEAR)
+
+        hr_coord, hr_coord_ = self.creat_coord(x)
+        lr_coord, _ = self.creat_coord(down_x)
+        q_coord = F.grid_sample(lr_coord, hr_coord_.flip(-1).unsqueeze(1), mode='nearest', align_corners=False)
+        q_coord = q_coord.view(b, -1, hr_coord.size(2) * hr_coord.size(3)).permute(0, 2, 1).contiguous()
+
+        # test_coord = F.grid_sample(lr_coord, hr_coord.permute(0, 2, 3, 1), mode='bilinear', align_corners=False)
+        # test_rel_coord = hr_coord - test_coord
+        # test_rel_coord = test_rel_coord.view(b, -1, 2)
+
+        # test_rel_coord[:, :, 0] *= down_x.shape[-2]
+        # test_rel_coord[:, :, 1] *= down_x.shape[-1]
+
+        rel_coord = hr_coord_ - q_coord
+        rel_coord[:, :, 0] *= down_x.shape[-2]
+        rel_coord[:, :, 1] *= down_x.shape[-1]
+
+        laplacian = x - resize(down_x, out_shape=up_size, antialiasing=False)
+        # laplacian = x - resize(down_x, out_shape=up_size, antialiasing=True)
+
+        laplacian = laplacian.reshape(b, laplacian.size(1), -1).permute(0, 2, 1).contiguous()
+
+        # cell
+        hr_grid = self.get_local_grid(x)
+        hr_cell = self.get_cell(x, hr_grid)
+        hr_cell_ = hr_cell.clone()
+        hr_cell_ = hr_cell_.permute(0, 2, 3, 1).contiguous()
+        rel_cell = hr_cell_.view(b, -1, hr_cell.size(1))
+        rel_cell[:, :, 0] *= down_x.shape[-2]
+        rel_cell[:, :, 1] *= down_x.shape[-1]
+
+        inp = torch.cat([rel_coord.cuda(), rel_cell.cuda(), laplacian], dim=-1)
+        local_weight = self.imnet(inp)
+        local_weight = local_weight.type(torch.float32)
+        local_weight = local_weight.view(b, -1, x.shape[1] * 9, 3).contiguous()
+
+        unfolded_x = F.unfold(x, 3, padding=1).view(b, -1, x.shape[2] * x.shape[3]).permute(0, 2, 1).contiguous()
+        cols = unfolded_x.unsqueeze(2)
+        out = torch.matmul(cols, local_weight).squeeze(2).permute(0, 2, 1).contiguous().view(b, -1, x.size(2),
+                                                                                             x.size(3))
+        out = resize(out, out_shape=down_size, antialiasing=False)
+        # out = resize(out, out_shape=down_size, antialiasing=True)
+
+        # out = down_sizer(out)
+
+        # ori
+        ori_lr_feature = self.extractor_mid(down_x)["bottom"]
+        ori_lr_feature = resize(ori_lr_feature, out_shape=(hr_feature.shape[2], hr_feature.shape[3]),
+                                antialiasing=False)
+        # ori_lr_feature = resize(ori_lr_feature, out_shape=(hr_feature.shape[2], hr_feature.shape[3]), antialiasing=True)
+        # ori_lr_feature = feature_sizer(ori_lr_feature)
+
+        # new
+        down_x = self.modulation(torch.cat([down_x, out], dim=1))
+        new_lr_feature = self.extractor_mid(down_x)["bottom"]
+
+        new_lr_feature = resize(new_lr_feature, out_shape=(hr_feature.shape[2], hr_feature.shape[3]),
+                                antialiasing=False)
+        # new_lr_feature = resize(new_lr_feature, out_shape=(hr_feature.shape[2], hr_feature.shape[3]), antialiasing=True)
+
+        # new_lr_feature = feature_sizer(new_lr_feature)
+
+        # res = resize(self.model(self.modulation(torch.cat([down_x, out], dim=1))), out_shape=up_size,
+        #              antialiasing=False)
+
+        # res = up_sizer(self.model(self.modulation(torch.cat([down_x, out], dim=1))))
+        res = 0
+
+        return down_x, hr_feature, \
+               new_lr_feature, ori_lr_feature, out, res
+
+    def inference(self, x, down_size, up_size):
+        b = x.shape[0]
+        down_x = resize(x, out_shape=down_size, antialiasing=False)
+        hr_coord, hr_coord_ = self.creat_coord(x)
+        lr_coord, _ = self.creat_coord(down_x)
+        q_coord = F.grid_sample(lr_coord, hr_coord_.flip(-1).unsqueeze(1), mode='nearest', align_corners=False)
+        q_coord = q_coord.view(b, -1, hr_coord.size(2) * hr_coord.size(3)).permute(0, 2, 1).contiguous()
+
+        rel_coord = hr_coord_ - q_coord
+        rel_coord[:, :, 0] *= down_x.shape[-2]
+        rel_coord[:, :, 1] *= down_x.shape[-1]
+
+        hr_grid = self.get_local_grid(x)
+        hr_cell = self.get_cell(x, hr_grid)
+
+        hr_cell_ = hr_cell.clone()
+        hr_cell_ = hr_cell_.permute(0, 2, 3, 1).contiguous()
+
+        rel_cell = hr_cell_.view(b, -1, hr_cell.size(1))
+        rel_cell[:, :, 0] *= down_x.shape[-2]
+        rel_cell[:, :, 1] *= down_x.shape[-1]
+
+        laplacian = x - resize(down_x, out_shape=up_size, antialiasing=False)
+        # laplacian = x - resize(down_x, out_shape=up_size, antialiasing=True)
+
+        laplacian = laplacian.reshape(b, laplacian.size(1), -1).permute(0, 2, 1).contiguous()
+        # laplacian = F.unfold(laplacian, 3, padding=1).view(b, -1, laplacian.shape[2] * laplacian.shape[3]).permute(0, 2, 1).contiguous()
+
+        inp = torch.cat([rel_coord.cuda(), rel_cell.cuda(), laplacian], dim=-1)
+        local_weight = self.imnet(inp)
+        local_weight = local_weight.type(torch.float32)
+        local_weight = local_weight.view(b, -1, x.shape[1] * 9, 3)
+
+        unfolded_x = F.unfold(x, 3, padding=1).view(b, -1, x.shape[2] * x.shape[3]).permute(0, 2, 1).contiguous()
+
+        cols = unfolded_x.unsqueeze(2)
+
+        out = torch.matmul(cols, local_weight).squeeze(2).permute(0, 2, 1).contiguous().view(b, -1, x.size(2),
+                                                                                             x.size(3))
+        out = resize(out, out_shape=down_size, antialiasing=False)
+        down_x = self.modulation(torch.cat([down_x, out], dim=1))
+
+        res = resize(self.model(down_x), out_shape=up_size, antialiasing=False)
+        return res, down_x
+
+    def creat_coord(self, x):
+        b = x.shape[0]
+        coord = make_coord(x.shape[-2:], flatten=False)
+        coord = coord.permute(2, 0, 1).contiguous().unsqueeze(0)
+        coord = coord.expand(b, 2, *coord.shape[-2:])
+
+        coord_ = coord.clone()
+        coord_ = coord_.clamp_(-1 + 1e-6, 1 - 1e-6)
+        coord_ = coord_.permute(0, 2, 3, 1).contiguous()
+        coord_ = coord_.view(b, -1, coord.size(1))
+        return coord.cuda(), coord_.cuda()
+
+    def get_local_grid(self, img):
+        local_grid = make_coord(img.shape[-2:], flatten=False).cuda()
+        local_grid = local_grid.permute(2, 0, 1).unsqueeze(0)
+        local_grid = local_grid.expand(img.shape[0], 2, *img.shape[-2:])
+
+        return local_grid
+
+    def get_cell(self, img, local_grid):
+        cell = torch.ones_like(local_grid)
+        cell[:, 0] *= 2 / img.size(2)
+        cell[:, 1] *= 2 / img.size(3)
+
+        return cell
+

model/__init__.py

@@ -0,0 +1,5 @@
+from .model import net
+from .resize_right import resize
+from .interp_methods import *
+from .module import Discriminator, Discriminator_new
+from .LMAR_model import *

model/interp_methods.py

@@ -0,0 +1,69 @@
+from math import pi
+
+try:
+    import torch
+except ImportError:
+    torch = None
+
+try:
+    import numpy
+except ImportError:
+    numpy = None
+
+if numpy is None and torch is None:
+    raise ImportError("Must have either Numpy or PyTorch but both not found")
+
+
+def set_framework_dependencies(x):
+    if type(x) is numpy.ndarray:
+        to_dtype = lambda a: a
+        fw = numpy
+    else:
+        to_dtype = lambda a: a.to(x.dtype)
+        fw = torch
+    eps = fw.finfo(fw.float32).eps
+    return fw, to_dtype, eps
+
+
+def support_sz(sz):
+    def wrapper(f):
+        f.support_sz = sz
+        return f
+    return wrapper
+
+
+@support_sz(4)
+def cubic(x):
+    fw, to_dtype, eps = set_framework_dependencies(x)
+    absx = fw.abs(x)
+    absx2 = absx ** 2
+    absx3 = absx ** 3
+    return ((1.5 * absx3 - 2.5 * absx2 + 1.) * to_dtype(absx <= 1.) +
+            (-0.5 * absx3 + 2.5 * absx2 - 4. * absx + 2.) *
+            to_dtype((1. < absx) & (absx <= 2.)))
+
+
+@support_sz(4)
+def lanczos2(x):
+    fw, to_dtype, eps = set_framework_dependencies(x)
+    return (((fw.sin(pi * x) * fw.sin(pi * x / 2) + eps) /
+            ((pi**2 * x**2 / 2) + eps)) * to_dtype(abs(x) < 2))
+
+
+@support_sz(6)
+def lanczos3(x):
+    fw, to_dtype, eps = set_framework_dependencies(x)
+    return (((fw.sin(pi * x) * fw.sin(pi * x / 3) + eps) /
+            ((pi**2 * x**2 / 3) + eps)) * to_dtype(abs(x) < 3))
+
+
+@support_sz(2)
+def linear(x):
+    fw, to_dtype, eps = set_framework_dependencies(x)
+    return ((x + 1) * to_dtype((-1 <= x) & (x < 0)) + (1 - x) *
+            to_dtype((0 <= x) & (x <= 1)))
+
+@support_sz(1)
+def box(x):
+    fw, to_dtype, eps = set_framework_dependencies(x)
+    return to_dtype((-1 <= x) & (x < 0)) + to_dtype((0 <= x) & (x <= 1))

model/model.py

@@ -0,0 +1,194 @@
+try:
+    from .module import *
+except:
+    from module import *
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.nn.init as init
+
+class SuperUnet_MS(nn.Module):
+    def __init__(self, channels, block="INV"):
+        super(SuperUnet_MS, self).__init__()
+        # ---------ENCODE
+        self.layer_dowm1 = basic_block(channels, channels, block)
+        self.dowm1 = nn.Sequential(nn.Conv2d(channels, channels * 2, 4, 2, 1, bias=True),
+                                   nn.InstanceNorm2d(channels * 2, affine=True), nn.LeakyReLU(0.2, inplace=True))
+        self.layer_dowm2 = basic_block(channels * 2, channels * 2, block)
+        self.dowm2 = nn.Sequential(nn.Conv2d(channels * 2, channels * 4, 4, 2, 1, bias=True),
+                                   nn.InstanceNorm2d(channels * 4, affine=True), nn.LeakyReLU(0.2, inplace=True))
+        # ---------DECODE
+        self.layer_bottom = basic_block(channels * 4, channels * 4, block)
+        self.up2 = nn.Sequential(nn.ConvTranspose2d(channels * 4, channels * 2, 4, 2, 1, bias=True),
+                                 nn.InstanceNorm2d(channels * 2, affine=True), nn.LeakyReLU(0.2, inplace=True))
+        self.layer_up2 = basic_block(channels * 2, channels * 2, block)
+        self.up1 = nn.Sequential(nn.ConvTranspose2d(channels * 2, channels, 4, 2, 1, bias=True),
+                                 nn.InstanceNorm2d(channels, affine=True), nn.LeakyReLU(0.2, inplace=True))
+        self.layer_up1 = basic_block(channels, channels, block)
+        # ---------SKIP
+        self.fus2 = skip(channels * 4, channels * 2, "HIN")
+        self.fus1 = skip(channels * 2, channels, "HIN")
+        # ---------SKIP
+        self.skip_down1 = nn.Sequential(nn.Conv2d(channels, channels, 4, 2, 1, bias=True),
+                                        nn.InstanceNorm2d(channels, affine=True), nn.LeakyReLU(0.2, inplace=True))
+        self.skip1 = skip(channels * 3, channels * 2, "CONV")
+        self.skip_down2 = nn.Sequential(nn.Conv2d(channels * 2, channels, 4, 2, 1, bias=True),
+                                        nn.InstanceNorm2d(channels, affine=True), nn.LeakyReLU(0.2, inplace=True))
+        self.skip2 = skip(channels * 5, channels * 4, "CONV")
+        # self.skip3 = skip(channels*2, channels, "CONV")
+        self.skip_up4 = nn.Sequential(nn.ConvTranspose2d(channels * 4, channels, 4, 2, 1, bias=True),
+                                      nn.InstanceNorm2d(channels, affine=True), nn.LeakyReLU(0.2, inplace=True))
+        self.skip4 = skip(channels * 3, channels * 2, "CONV")
+        # self.skip5 = skip(channels*2, channels, "CONV")
+        self.skip_up6 = nn.Sequential(nn.ConvTranspose2d(channels * 2, channels, 4, 2, 1, bias=True),
+                                      nn.InstanceNorm2d(channels, affine=True), nn.LeakyReLU(0.2, inplace=True))
+        self.skip6 = skip(channels * 2, channels, "CONV")
+
+    def forward(self, x):
+        # ---------ENCODE
+        x_11 = self.layer_dowm1(x)
+        x_down1 = self.dowm1(x_11)
+        # x  =self.skip_down1(x)
+        # print(x.shape, x_down1.shape)
+
+        x_down1 = self.skip1(torch.cat([self.skip_down1(x), x_down1], 1), x_down1)
+
+        x_12 = self.layer_dowm2(x_down1)
+        x_down2 = self.dowm2(x_12)
+        x_down2 = self.skip2(torch.cat([self.skip_down2(x_down1), x_down2], 1), x_down2)
+
+        x_bottom = self.layer_bottom(x_down2)
+
+        # ---------DECODE
+        x_up2 = self.up2(x_bottom)
+        x_22 = self.layer_up2(x_up2)
+        x_22 = self.skip4(torch.cat([self.skip_up4(x_bottom), x_22], 1), x_22)
+        x_22 = self.fus2(torch.cat([x_12, x_22], 1), x_22)
+
+        x_up1 = self.up1(x_22)
+        x_21 = self.layer_up1(x_up1)
+        x_21 = self.skip6(torch.cat([self.skip_up6(x_22), x_21], 1), x_21)
+        x_21 = self.fus1(torch.cat([x_11, x_21], 1), x_21)
+        return x_21, x_down2
+
+
+class skip(nn.Module):
+    def __init__(self, channels_in, channels_out, block):
+        super(skip, self).__init__()
+        if block == "CONV":
+            self.body = nn.Sequential(nn.Conv2d(channels_in, channels_out, 1, 1, 0, bias=True),
+                                      nn.InstanceNorm2d(channels_out, affine=True), nn.ReLU(inplace=True), )
+        if block == "ID":
+            self.body = nn.Identity()
+        if block == "INV":
+            self.body = nn.Sequential(InvBlock(channels_in, channels_in // 2),
+                                      nn.Conv2d(channels_in, channels_out, 1, 1, 0, bias=True), )
+        if block == "HIN":
+            self.body = nn.Sequential(HinBlock(channels_in, channels_out))
+        # --------------------------------------
+        self.alpha1 = nn.Parameter(torch.FloatTensor(1), requires_grad=True)
+        self.alpha1.data.fill_(1.0)
+        self.alpha2 = nn.Parameter(torch.FloatTensor(1), requires_grad=True)
+        self.alpha2.data.fill_(0.5)
+
+    def forward(self, x, y):
+        out = self.alpha1 * self.body(x) + self.alpha2 * y
+        return out
+
+
+def subnet(net_structure, init='xavier'):
+    def constructor(channel_in, channel_out):
+        if net_structure == 'HIN':
+            return HinBlock(channel_in, channel_out)
+
+    return constructor
+
+
+class InvBlock(nn.Module):
+    def __init__(self, channel_num, channel_split_num, subnet_constructor=subnet('HIN'),
+                 clamp=0.8):  ################  split_channel一般设为channel_num的一半
+        super(InvBlock, self).__init__()
+        # channel_num: 3
+        # channel_split_num: 1
+
+        self.split_len1 = channel_split_num  # 1
+        self.split_len2 = channel_num - channel_split_num  # 2
+
+        self.clamp = clamp
+
+        self.F = subnet_constructor(self.split_len2, self.split_len1)
+        self.G = subnet_constructor(self.split_len1, self.split_len2)
+        self.H = subnet_constructor(self.split_len1, self.split_len2)
+
+    def forward(self, x):
+        x1, x2 = (x.narrow(1, 0, self.split_len1), x.narrow(1, self.split_len1, self.split_len2))
+
+        y1 = x1 + self.F(x2)  # 1 channel
+        self.s = self.clamp * (torch.sigmoid(self.H(y1)) * 2 - 1)
+        y2 = x2.mul(torch.exp(self.s)) + self.G(y1)  # 2 channel
+        out = torch.cat((y1, y2), 1)
+
+        return out + x
+
+
+class sample_block(nn.Module):
+    def __init__(self, channels_in, channels_out, size, dil):
+        super(sample_block, self).__init__()
+        # ------------------------------------------
+        if size == "DOWN":
+            self.conv = nn.Sequential(
+                nn.Conv2d(channels_in, channels_out, 3, 1, dil, dilation=dil),
+                nn.InstanceNorm2d(channels_out, affine=True),
+                nn.ReLU(inplace=True),
+            )
+        if size == "UP":
+            self.conv = nn.Sequential(
+                nn.ConvTranspose2d(channels_in, channels_out, 3, 1, dil, dilation=dil),
+                nn.InstanceNorm2d(channels_out, affine=True),
+                nn.ReLU(inplace=True),
+            )
+
+    def forward(self, x):
+        return self.conv(x)
+
+
+class HinBlock(nn.Module):
+    def __init__(self, in_size, out_size):
+        super(HinBlock, self).__init__()
+        self.identity = nn.Conv2d(in_size, out_size, 1, 1, 0)
+        self.norm = nn.InstanceNorm2d(out_size // 2, affine=True)
+
+        self.conv_1 = nn.Conv2d(in_size, out_size, kernel_size=3, stride=1, padding=1, bias=True)
+        self.relu_1 = nn.Sequential(nn.LeakyReLU(0.2, inplace=False), )
+        self.conv_2 = nn.Sequential(nn.Conv2d(out_size, out_size, kernel_size=3, stride=1, padding=1, bias=True),
+                                    nn.LeakyReLU(0.2, inplace=False), )
+
+    def forward(self, x):
+        out = self.conv_1(x)
+        out_1, out_2 = torch.chunk(out, 2, dim=1)
+        out = torch.cat([self.norm(out_1), out_2], dim=1)
+        out = self.relu_1(out)
+        out = self.conv_2(out)
+        out += self.identity(x)
+        return out
+
+
+class net(nn.Module):
+    def __init__(self, args):
+        super().__init__()
+        self.args = args.model
+        self.hr_inc = DoubleConv(self.args["in_channel"], self.args["model_channel"] * 2)
+        self.hr_backbone = SuperUnet_MS(self.args["model_channel"] * 2)
+        self.final_out = nn.Conv2d(self.args["model_channel"] * 2, 3, kernel_size=1, bias=False)
+
+    def forward(self, x):
+        x = self.hr_inc(x)
+        x, mid_feat = self.hr_backbone(x)
+        out = self.final_out(x)
+        return out
+

model/module.py

@@ -0,0 +1,248 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+from torchvision.transforms.functional import rgb_to_grayscale
+import numpy as np
+
+class DoubleConv(nn.Module):
+    """(convolution => [BN] => ReLU) * 2"""
+
+    def __init__(self, in_channels, out_channels, mid_channels=None):
+        super().__init__()
+        if not mid_channels:
+            mid_channels = out_channels
+        self.double_conv = nn.Sequential(
+            nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1, bias=False),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1, bias=False),
+            nn.ReLU(inplace=True)
+        )
+        self.apply(self._init_weights)
+
+    def forward(self, x):
+        return self.double_conv(x)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Conv2d):
+            n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+            m.weight.data.normal_(0, math.sqrt(2. / n))
+        elif isinstance(m, nn.BatchNorm2d):
+            m.weight.data.fill_(1)
+            m.bias.data.zero_()
+
+
+class Down(nn.Module):
+    """Downscaling with maxpool then double conv"""
+
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        self.maxpool_conv = nn.Sequential(
+            nn.MaxPool2d(2),
+            DoubleConv(in_channels, out_channels)
+        )
+
+    def forward(self, x):
+        return self.maxpool_conv(x)
+
+
+class Up(nn.Module):
+    """Upscaling then double conv"""
+
+    def __init__(self, in_channels, out_channels, bilinear=True):
+        super().__init__()
+
+        # if bilinear, use the normal convolutions to reduce the number of channels
+        if bilinear:
+            self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
+            self.conv = DoubleConv(in_channels, out_channels, in_channels // 2)
+        else:
+            self.up = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2)
+            self.conv = DoubleConv(in_channels, out_channels)
+
+    def forward(self, x1, x2):
+        x1 = self.up(x1)
+        # input is CHW
+        diffY = x2.size()[2] - x1.size()[2]
+        diffX = x2.size()[3] - x1.size()[3]
+
+        x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2,
+                        diffY // 2, diffY - diffY // 2])
+        # if you have padding issues, see
+        # https://github.com/HaiyongJiang/U-Net-Pytorch-Unstructured-Buggy/commit/0e854509c2cea854e247a9c615f175f76fbb2e3a
+        # https://github.com/xiaopeng-liao/Pytorch-UNet/commit/8ebac70e633bac59fc22bb5195e513d5832fb3bd
+        x = torch.cat([x2, x1], dim=1)
+        return self.conv(x)
+
+
+# spatial attention
+class SpatialGate(nn.Module):
+    def __init__(self, in_channels):
+        super(SpatialGate, self).__init__()
+        self.spatial = nn.Conv2d(in_channels, 1, kernel_size=1)
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, x):
+        x_out = self.spatial(x)
+        scale = self.sigmoid(x_out)
+        return scale * x
+
+
+# sobel
+class SobelOperator(nn.Module):
+    def __init__(self):
+        super(SobelOperator, self).__init__()
+        self.conv_x = nn.Conv2d(1, 1, kernel_size=3, padding=1, bias=False)
+        self.conv_y = nn.Conv2d(1, 1, kernel_size=3, padding=1, bias=False)
+        self.conv_x.weight[0].data[:, :, :] = torch.FloatTensor([[[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]]])
+        self.conv_y.weight[0].data[:, :, :] = torch.FloatTensor([[[-1, -2, -1], [0, 0, 0], [1, 2, 1]]])
+
+    def forward(self, x):
+        G_x = self.conv_x(x)
+        G_y = self.conv_y(x)
+        grad_mag = torch.sqrt(torch.pow(G_x, 2) + torch.pow(G_y, 2))
+        return grad_mag
+
+
+class offset_estimator(nn.Sequential):
+    def __init__(self, kernel_size, fwhm, in_channel, mid_channel, out_channel) -> None:
+        super().__init__()
+        model = []
+        assert len(kernel_size) == len(fwhm), "length error"
+        for i in range(len(kernel_size)):
+            if i == 0:
+                gaussian_weight = torch.FloatTensor(gaussian_2d(kernel_size[i], fwhm=fwhm[i]))
+                gauss_filter = nn.Conv2d(in_channel, mid_channel, kernel_size[i], padding=(kernel_size[i] - 1) // 2,
+                                         bias=False)
+                gauss_filter.weight[0].data[:, :, :] = gaussian_weight
+                model += [gauss_filter, nn.ReLU(inplace=True)]
+            elif i == len(kernel_size) - 1:
+                gaussian_weight = torch.FloatTensor(gaussian_2d(kernel_size[i], fwhm=fwhm[i]))
+                gauss_filter = nn.Conv2d(mid_channel, out_channel, kernel_size[i], padding=(kernel_size[i] - 1) // 2,
+                                         bias=False)
+                gauss_filter.weight[0].data[:, :, :] = gaussian_weight
+                model += [gauss_filter, nn.ReLU(inplace=True)]
+            else:
+                gaussian_weight = torch.FloatTensor(gaussian_2d(kernel_size[i], fwhm=fwhm[i]))
+                gauss_filter = nn.Conv2d(mid_channel, mid_channel, kernel_size[i], padding=(kernel_size[i] - 1) // 2,
+                                         bias=False)
+                gauss_filter.weight[0].data[:, :, :] = gaussian_weight
+                model += [gauss_filter, nn.ReLU(inplace=True)]
+        self.model = nn.Sequential(*model)
+
+    def forward(self, x):
+        return self.model(x)
+
+
+# Channel attention
+def logsumexp_2d(tensor):
+    tensor_flatten = tensor.view(tensor.size(0), tensor.size(1), -1)
+    s, _ = torch.max(tensor_flatten, dim=2, keepdim=True)
+    outputs = s + (tensor_flatten - s).exp().sum(dim=2, keepdim=True).log()
+    return outputs
+
+
+class Flatten(nn.Module):
+    def forward(self, x):
+        return x.view(x.size(0), -1)
+
+
+class ChannelGate(nn.Module):
+    def __init__(self, gate_channels, reduction_ratio=16, pool_types=['avg', 'max']):
+        super(ChannelGate, self).__init__()
+        self.gate_channels = gate_channels
+        self.mlp = nn.Sequential(
+            Flatten(),
+            nn.Linear(gate_channels, gate_channels // reduction_ratio),
+            nn.ReLU(),
+            nn.Linear(gate_channels // reduction_ratio, gate_channels)
+        )
+        self.pool_types = pool_types
+
+    def forward(self, x):
+        channel_att_sum = None
+        for pool_type in self.pool_types:
+            if pool_type == 'avg':
+                avg_pool = F.avg_pool2d(x, (x.size(2), x.size(3)), stride=(x.size(2), x.size(3)))
+                channel_att_raw = self.mlp(avg_pool)
+            elif pool_type == 'max':
+                max_pool = F.max_pool2d(x, (x.size(2), x.size(3)), stride=(x.size(2), x.size(3)))
+                channel_att_raw = self.mlp(max_pool)
+            elif pool_type == 'lp':
+                lp_pool = F.lp_pool2d(x, 2, (x.size(2), x.size(3)), stride=(x.size(2), x.size(3)))
+                channel_att_raw = self.mlp(lp_pool)
+            elif pool_type == 'lse':
+                # LSE pool only
+                lse_pool = logsumexp_2d(x)
+                channel_att_raw = self.mlp(lse_pool)
+
+            if channel_att_sum is None:
+                channel_att_sum = channel_att_raw
+            else:
+                channel_att_sum = channel_att_sum + channel_att_raw
+
+        scale = torch.sigmoid(channel_att_sum).unsqueeze(2).unsqueeze(3).expand_as(x)
+        return x * scale
+
+
+# LBP
+def LBP(image):  # b, 3, h, w tensor
+    radius = 2
+    n_points = 8 * radius
+    method = 'uniform'
+    gray_img = rgb_to_grayscale(image)  # b, 1, h, w
+    gray_img = gray_img.squeeze(1)
+    lbf_feature = np.zeros((gray_img.shape[0], gray_img.shape[1], gray_img.shape[2]))
+    for i in range(gray_img.shape[0]):
+        lbf_feature[i] = feature.local_binary_pattern(gray_img[i], n_points, radius, method)
+    return torch.FloatTensor(lbf_feature).unsqueeze(1)
+
+
+class Discriminator(nn.Module):
+    def __init__(self, in_channel):
+        super().__init__()
+        self.in_channel = in_channel
+
+        def discriminator_block(in_filters, out_filters):
+            layers = [nn.Conv2d(in_filters, out_filters, 4, stride=2, padding=1), nn.LeakyReLU(0.2, inplace=False)]
+            return layers
+
+        self.model = nn.Sequential(
+            *discriminator_block(self.in_channel, 4),
+            *discriminator_block(4, 4),
+            *discriminator_block(4, 4),
+            *discriminator_block(4, 4),
+            nn.ZeroPad2d((1, 0, 1, 0)),
+            nn.Conv2d(4, 1, 4, padding=1, bias=False)
+        )
+
+    def forward(self, x):
+        return self.model(x)
+
+
+class Discriminator_new(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+        def discriminator_block(in_filters, out_filters, first_block=False):
+            layers = []
+            layers.append(nn.Conv2d(in_filters, out_filters, kernel_size=3, stride=1, padding=1))
+            layers.append(nn.LeakyReLU(0.2, inplace=True))
+            layers.append(nn.Conv2d(out_filters, out_filters, kernel_size=3, stride=2, padding=1))
+            layers.append(nn.LeakyReLU(0.2, inplace=True))
+            return layers
+
+        layers = []
+        in_filters = 3
+        for i, out_filters in enumerate([4, 6, 8, 10]):
+            layers.extend(discriminator_block(in_filters, out_filters, first_block=(i == 0)))
+            in_filters = out_filters
+
+        layers.append(nn.ZeroPad2d((1, 0, 1, 0)))
+        layers.append(nn.Conv2d(out_filters, 1, kernel_size=3, stride=1, padding=1))
+
+        self.model = nn.Sequential(*layers)
+
+    def forward(self, img):
+        return self.model(img)
+

model/resize_right.py

@@ -0,0 +1,437 @@
+from typing import Tuple
+import warnings
+from math import ceil
+
+try:
+    from .interp_methods import *
+except:
+    from interp_methods import *
+from fractions import Fraction
+
+
+class NoneClass:
+    pass
+
+
+try:
+    import torch
+    from torch import nn
+
+    nnModuleWrapped = nn.Module
+except ImportError:
+    warnings.warn('No PyTorch found, will work only with Numpy')
+    torch = None
+    nnModuleWrapped = NoneClass
+
+try:
+    import numpy
+except ImportError:
+    warnings.warn('No Numpy found, will work only with PyTorch')
+    numpy = None
+
+if numpy is None and torch is None:
+    raise ImportError("Must have either Numpy or PyTorch but both not found")
+
+
+def resize(input, scale_factors=None, out_shape=None,
+           interp_method=lanczos3, support_sz=None,
+           antialiasing=True, by_convs=False, scale_tolerance=None,
+           max_numerator=10, pad_mode='constant', adv_weights=None):
+    # get properties of the input tensor
+    in_shape, n_dims = input.shape, input.ndim
+
+    # fw stands for framework that can be either numpy or torch,
+    # determined by the input type
+    fw = numpy if type(input) is numpy.ndarray else torch
+    eps = fw.finfo(fw.float32).eps
+    device = input.device if fw is torch else None
+    weights_container = []
+
+    # set missing scale factors or output shapem one according to another,
+    # scream if both missing. this is also where all the defults policies
+    # take place. also handling the by_convs attribute carefully.
+    scale_factors, out_shape, by_convs = set_scale_and_out_sz(in_shape,
+                                                              out_shape,
+                                                              scale_factors,
+                                                              by_convs,
+                                                              scale_tolerance,
+                                                              max_numerator,
+                                                              eps, fw)
+
+    # sort indices of dimensions according to scale of each dimension.
+    # since we are going dim by dim this is efficient
+    sorted_filtered_dims_and_scales = [(dim, scale_factors[dim], by_convs[dim],
+                                        in_shape[dim], out_shape[dim])
+                                       for dim in sorted(range(n_dims),
+                                                         key=lambda ind: scale_factors[ind])
+                                       if scale_factors[dim] != 1.]
+
+    # unless support size is specified by the user, it is an attribute
+    # of the interpolation method
+    if support_sz is None:
+        support_sz = interp_method.support_sz
+
+    # output begins identical to input and changes with each iteration
+    output = input
+
+    # iterate over dims
+    for i, (dim, scale_factor, dim_by_convs, in_sz, out_sz
+            ) in enumerate(sorted_filtered_dims_and_scales):
+        # STEP 1- PROJECTED GRID: The non-integer locations of the projection
+        # of output pixel locations to the input tensor
+        projected_grid = get_projected_grid(in_sz, out_sz,
+                                            scale_factor, fw, dim_by_convs,
+                                            device)
+
+        # STEP 1.5: ANTIALIASING- If antialiasing is taking place, we modify
+        # the window size and the interpolation method (see inside function)
+        cur_interp_method, cur_support_sz = apply_antialiasing_if_needed(
+            interp_method,
+            support_sz,
+            scale_factor,
+            antialiasing)
+
+        # STEP 2- FIELDS OF VIEW: for each output pixels, map the input pixels
+        # that influence it. Also calculate needed padding and update grid
+        # accoedingly
+        field_of_view = get_field_of_view(projected_grid, cur_support_sz, fw,
+                                          eps, device)
+
+        # STEP 2.5- CALCULATE PAD AND UPDATE: according to the field of view,
+        # the input should be padded to handle the boundaries, coordinates
+        # should be updated. actual padding only occurs when weights are
+        # aplied (step 4). if using by_convs for this dim, then we need to
+        # calc right and left boundaries for each filter instead.
+        pad_sz, projected_grid, field_of_view = calc_pad_sz(in_sz, out_sz,
+                                                            field_of_view,
+                                                            projected_grid,
+                                                            scale_factor,
+                                                            dim_by_convs, fw,
+                                                            device)
+
+        # STEP 3- CALCULATE WEIGHTS: Match a set of weights to the pixels in
+        # the field of view for each output pixel
+        if adv_weights != None:
+            weights = adv_weights[i]
+        else:
+            weights = get_weights(cur_interp_method, projected_grid, field_of_view)
+            weights_container.append(weights)
+
+        # STEP 4- APPLY WEIGHTS: Each output pixel is calculated by multiplying
+        # its set of weights with the pixel values in its field of view.
+        # We now multiply the fields of view with their matching weights.
+        # We do this by tensor multiplication and broadcasting.
+        # if by_convs is true for this dim, then we do this action by
+        # convolutions. this is equivalent but faster.
+        if not dim_by_convs:
+            output = apply_weights(output, field_of_view, weights, dim, n_dims,
+                                   pad_sz, pad_mode, fw)
+        else:
+            output = apply_convs(output, scale_factor, in_sz, out_sz, weights,
+                                 dim, pad_sz, pad_mode, fw)
+    return output
+
+
+def get_projected_grid(in_sz, out_sz, scale_factor, fw, by_convs, device=None):
+    # we start by having the ouput coordinates which are just integer locations
+    # in the special case when usin by_convs, we only need two cycles of grid
+    # points. the first and last.
+    grid_sz = out_sz if not by_convs else scale_factor.numerator
+    out_coordinates = fw_arange(grid_sz, fw, device)
+
+    # This is projecting the ouput pixel locations in 1d to the input tensor,
+    # as non-integer locations.
+    # the following fomrula is derived in the paper
+    # "From Discrete to Continuous Convolutions" by Shocher et al.
+    return (out_coordinates / float(scale_factor) +
+            (in_sz - 1) / 2 - (out_sz - 1) / (2 * float(scale_factor)))
+
+
+def get_field_of_view(projected_grid, cur_support_sz, fw, eps, device):
+    # for each output pixel, map which input pixels influence it, in 1d.
+    # we start by calculating the leftmost neighbor, using half of the window
+    # size (eps is for when boundary is exact int)
+    left_boundaries = fw_ceil(projected_grid - cur_support_sz / 2 - eps, fw)
+
+    # then we simply take all the pixel centers in the field by counting
+    # window size pixels from the left boundary
+    ordinal_numbers = fw_arange(ceil(cur_support_sz - eps), fw, device)
+    return left_boundaries[:, None] + ordinal_numbers
+
+
+def calc_pad_sz(in_sz, out_sz, field_of_view, projected_grid, scale_factor,
+                dim_by_convs, fw, device):
+    if not dim_by_convs:
+        # determine padding according to neighbor coords out of bound.
+        # this is a generalized notion of padding, when pad<0 it means crop
+        pad_sz = [-field_of_view[0, 0].item(),
+                  field_of_view[-1, -1].item() - in_sz + 1]
+
+        # since input image will be changed by padding, coordinates of both
+        # field_of_view and projected_grid need to be updated
+        field_of_view += pad_sz[0]
+        projected_grid += pad_sz[0]
+
+    else:
+        # only used for by_convs, to calc the boundaries of each filter the
+        # number of distinct convolutions is the numerator of the scale factor
+        num_convs, stride = scale_factor.numerator, scale_factor.denominator
+
+        # calculate left and right boundaries for each conv. left can also be
+        # negative right can be bigger than in_sz. such cases imply padding if
+        # needed. however if# both are in-bounds, it means we need to crop,
+        # practically apply the conv only on part of the image.
+        left_pads = -field_of_view[:, 0]
+
+        # next calc is tricky, explanation by rows:
+        # 1) counting output pixels between the first position of each filter
+        #    to the right boundary of the input
+        # 2) dividing it by number of filters to count how many 'jumps'
+        #    each filter does
+        # 3) multiplying by the stride gives us the distance over the input
+        #    coords done by all these jumps for each filter
+        # 4) to this distance we add the right boundary of the filter when
+        #    placed in its leftmost position. so now we get the right boundary
+        #    of that filter in input coord.
+        # 5) the padding size needed is obtained by subtracting the rightmost
+        #    input coordinate. if the result is positive padding is needed. if
+        #    negative then negative padding means shaving off pixel columns.
+        right_pads = (((out_sz - fw_arange(num_convs, fw, device) - 1)  # (1)
+                       // num_convs)  # (2)
+                      * stride  # (3)
+                      + field_of_view[:, -1]  # (4)
+                      - in_sz + 1)  # (5)
+
+        # in the by_convs case pad_sz is a list of left-right pairs. one per
+        # each filter
+
+        pad_sz = list(zip(left_pads, right_pads))
+
+    return pad_sz, projected_grid, field_of_view
+
+
+def get_weights(interp_method, projected_grid, field_of_view):
+    # the set of weights per each output pixels is the result of the chosen
+    # interpolation method applied to the distances between projected grid
+    # locations and the pixel-centers in the field of view (distances are
+    # directed, can be positive or negative)
+    weights = interp_method(projected_grid[:, None] - field_of_view)
+
+    # we now carefully normalize the weights to sum to 1 per each output pixel
+    sum_weights = weights.sum(1, keepdims=True)
+    sum_weights[sum_weights == 0] = 1
+    return weights / sum_weights
+
+
+def apply_weights(input, field_of_view, weights, dim, n_dims, pad_sz, pad_mode,
+                  fw):
+    # for this operation we assume the resized dim is the first one.
+    # so we transpose and will transpose back after multiplying
+    tmp_input = fw_swapaxes(input, dim, 0, fw)
+
+    # apply padding
+    tmp_input = fw_pad(tmp_input, fw, pad_sz, pad_mode)
+
+    # field_of_view is a tensor of order 2: for each output (1d location
+    # along cur dim)- a list of 1d neighbors locations.
+    # note that this whole operations is applied to each dim separately,
+    # this is why it is all in 1d.
+    # neighbors = tmp_input[field_of_view] is a tensor of order image_dims+1:
+    # for each output pixel (this time indicated in all dims), these are the
+    # values of the neighbors in the 1d field of view. note that we only
+    # consider neighbors along the current dim, but such set exists for every
+    # multi-dim location, hence the final tensor order is image_dims+1.
+    neighbors = tmp_input[field_of_view]
+
+    # weights is an order 2 tensor: for each output location along 1d- a list
+    # of weights matching the field of view. we augment it with ones, for
+    # broadcasting, so that when multiplies some tensor the weights affect
+    # only its first dim.
+    tmp_weights = fw.reshape(weights, (*weights.shape, *[1] * (n_dims - 1)))
+
+    # now we simply multiply the weights with the neighbors, and then sum
+    # along the field of view, to get a single value per out pixel
+    tmp_output = (neighbors * tmp_weights).sum(1)
+
+    # we transpose back the resized dim to its original position
+    return fw_swapaxes(tmp_output, 0, dim, fw)
+
+
+def apply_convs(input, scale_factor, in_sz, out_sz, weights, dim, pad_sz,
+                pad_mode, fw):
+    # for this operations we assume the resized dim is the last one.
+    # so we transpose and will transpose back after multiplying
+    input = fw_swapaxes(input, dim, -1, fw)
+
+    # the stride for all convs is the denominator of the scale factor
+    stride, num_convs = scale_factor.denominator, scale_factor.numerator
+
+    # prepare an empty tensor for the output
+    tmp_out_shape = list(input.shape)
+    tmp_out_shape[-1] = out_sz
+    tmp_output = fw_empty(tuple(tmp_out_shape), fw, input.device)
+
+    # iterate over the conv operations. we have as many as the numerator
+    # of the scale-factor. for each we need boundaries and a filter.
+    for conv_ind, (pad_sz, filt) in enumerate(zip(pad_sz, weights)):
+        # apply padding (we pad last dim, padding can be negative)
+        pad_dim = input.ndim - 1
+        tmp_input = fw_pad(input, fw, pad_sz, pad_mode, dim=pad_dim)
+
+        # apply convolution over last dim. store in the output tensor with
+        # positional strides so that when the loop is comlete conv results are
+        # interwind
+        tmp_output[..., conv_ind::num_convs] = fw_conv(tmp_input, filt, stride)
+
+    return fw_swapaxes(tmp_output, -1, dim, fw)
+
+
+def set_scale_and_out_sz(in_shape, out_shape, scale_factors, by_convs,
+                         scale_tolerance, max_numerator, eps, fw):
+    # eventually we must have both scale-factors and out-sizes for all in/out
+    # dims. however, we support many possible partial arguments
+    if scale_factors is None and out_shape is None:
+        raise ValueError("either scale_factors or out_shape should be "
+                         "provided")
+    if out_shape is not None:
+        # if out_shape has less dims than in_shape, we defaultly resize the
+        # first dims for numpy and last dims for torch
+        out_shape = (list(out_shape) + list(in_shape[len(out_shape):])
+                     if fw is numpy
+                     else list(in_shape[:-len(out_shape)]) + list(out_shape))
+        if scale_factors is None:
+            # if no scale given, we calculate it as the out to in ratio
+            # (not recomended)
+            scale_factors = [out_sz / in_sz for out_sz, in_sz
+                             in zip(out_shape, in_shape)]
+    if scale_factors is not None:
+        # by default, if a single number is given as scale, we assume resizing
+        # two dims (most common are images with 2 spatial dims)
+        scale_factors = (scale_factors
+                         if isinstance(scale_factors, (list, tuple))
+                         else [scale_factors, scale_factors])
+        # if less scale_factors than in_shape dims, we defaultly resize the
+        # first dims for numpy and last dims for torch
+        scale_factors = (list(scale_factors) + [1] *
+                         (len(in_shape) - len(scale_factors)) if fw is numpy
+                         else [1] * (len(in_shape) - len(scale_factors)) +
+                              list(scale_factors))
+        if out_shape is None:
+            # when no out_shape given, it is calculated by multiplying the
+            # scale by the in_shape (not recomended)
+            out_shape = [ceil(scale_factor * in_sz)
+                         for scale_factor, in_sz in
+                         zip(scale_factors, in_shape)]
+        # next part intentionally after out_shape determined for stability
+        # we fix by_convs to be a list of truth values in case it is not
+        if not isinstance(by_convs, (list, tuple)):
+            by_convs = [by_convs] * len(out_shape)
+
+        # next loop fixes the scale for each dim to be either frac or float.
+        # this is determined by by_convs and by tolerance for scale accuracy.
+        for ind, (sf, dim_by_convs) in enumerate(zip(scale_factors, by_convs)):
+            # first we fractionaize
+            if dim_by_convs:
+                frac = Fraction(1 / sf).limit_denominator(max_numerator)
+                frac = Fraction(numerator=frac.denominator, denominator=frac.numerator)
+
+            # if accuracy is within tolerance scale will be frac. if not, then
+            # it will be float and the by_convs attr will be set false for
+            # this dim
+            if scale_tolerance is None:
+                scale_tolerance = eps
+            if dim_by_convs and abs(frac - sf) < scale_tolerance:
+                scale_factors[ind] = frac
+            else:
+                scale_factors[ind] = float(sf)
+                by_convs[ind] = False
+
+        return scale_factors, out_shape, by_convs
+
+
+def apply_antialiasing_if_needed(interp_method, support_sz, scale_factor,
+                                 antialiasing):
+    # antialiasing is "stretching" the field of view according to the scale
+    # factor (only for downscaling). this is low-pass filtering. this
+    # requires modifying both the interpolation (stretching the 1d
+    # function and multiplying by the scale-factor) and the window size.
+    scale_factor = float(scale_factor)
+    if scale_factor >= 1.0 or not antialiasing:
+        return interp_method, support_sz
+    cur_interp_method = (lambda arg: scale_factor *
+                                     interp_method(scale_factor * arg))
+    cur_support_sz = support_sz / scale_factor
+    return cur_interp_method, cur_support_sz
+
+
+def fw_ceil(x, fw):
+    if fw is numpy:
+        return fw.int_(fw.ceil(x))
+    else:
+        return x.ceil().long()
+
+
+def fw_floor(x, fw):
+    if fw is numpy:
+        return fw.int_(fw.floor(x))
+    else:
+        return x.floor().long()
+
+
+def fw_cat(x, fw):
+    if fw is numpy:
+        return fw.concatenate(x)
+    else:
+        return fw.cat(x)
+
+
+def fw_swapaxes(x, ax_1, ax_2, fw):
+    if fw is numpy:
+        return fw.swapaxes(x, ax_1, ax_2)
+    else:
+        return x.transpose(ax_1, ax_2)
+
+
+def fw_pad(x, fw, pad_sz, pad_mode, dim=0):
+    if pad_sz == (0, 0):
+        return x
+    if fw is numpy:
+        pad_vec = [(0, 0)] * x.ndim
+        pad_vec[dim] = pad_sz
+        return fw.pad(x, pad_width=pad_vec, mode=pad_mode)
+    else:
+        if x.ndim < 3:
+            x = x[None, None, ...]
+
+        pad_vec = [0] * ((x.ndim - 2) * 2)
+        pad_vec[0:2] = pad_sz
+        return fw.nn.functional.pad(x.transpose(dim, -1), pad=pad_vec,
+                                    mode=pad_mode).transpose(dim, -1)
+
+
+def fw_conv(input, filter, stride):
+    # we want to apply 1d conv to any nd array. the way to do it is to reshape
+    # the input to a 4D tensor. first two dims are singeletons, 3rd dim stores
+    # all the spatial dims that we are not convolving along now. then we can
+    # apply conv2d with a 1xK filter. This convolves the same way all the other
+    # dims stored in the 3d dim. like depthwise conv over these.
+    # TODO: numpy support
+    reshaped_input = input.reshape(1, 1, -1, input.shape[-1])
+    reshaped_output = torch.nn.functional.conv2d(reshaped_input,
+                                                 filter.view(1, 1, 1, -1),
+                                                 stride=(1, stride))
+    return reshaped_output.reshape(*input.shape[:-1], -1)
+
+
+def fw_arange(upper_bound, fw, device):
+    if fw is numpy:
+        return fw.arange(upper_bound)
+    else:
+        return fw.arange(upper_bound, device=device)
+
+
+def fw_empty(shape, fw, device):
+    if fw is numpy:
+        return fw.empty(shape)
+    else:
+        return fw.empty(size=(*shape,), device=device)

pretrained_models/LMAR_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:27f1ada04c3297053af030ec2547a06f54d5de5e1ec20f3b430a9dd2f2f666ff
+size 1475245

pretrained_models/base_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:45d49de91c08e7a6080d60f7059482fcd443377982e1908045625759e5931772
+size 3417093

utils.py

@@ -0,0 +1,177 @@
+import torchvision.utils as vutils
+import argparse
+import yaml
+import torch
+import torchvision
+from metrics import calculate_psnr, calculate_ssim
+import torchvision.transforms as transforms
+import numpy as np
+from torch.optim.lr_scheduler import _LRScheduler
+import math
+
+
+class AverageMeter(object):
+    """Computes and stores the average and current value"""
+
+    def __init__(self):
+        self.reset()
+
+    def reset(self):
+        self.val = 0
+        self.avg = 0
+        self.sum = 0
+        self.count = 0
+
+    def update(self, val, n=1):
+        self.val = val
+        self.sum += val * n
+        self.count += n
+        self.avg = self.sum / self.count
+
+
+def calculate_metrics(imgs_1, imgs_2):
+    psnrs = []
+    ssims = []
+    assert imgs_1.shape[0] == imgs_2.shape[0]
+    batch_size = imgs_1.shape[0]
+    for i in range(batch_size):
+        img1 = imgs_1[i]
+        img2 = imgs_2[i]
+        img1 = np.asarray(transforms.ToPILImage()(img1))
+        img2 = np.asarray(transforms.ToPILImage()(img2))
+        psnr = calculate_psnr(img1, img2, 0)
+        ssim = calculate_ssim(img1, img2, 0)
+        psnrs.append(psnr)
+        ssims.append(ssim)
+    return np.asarray(psnrs).mean(), np.asarray(ssims).mean()
+
+
+def read_args(config_file):
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", default=config_file)
+    file = open(config_file)
+    config = yaml.safe_load(file)
+    for k, v in config.items():
+        parser.add_argument(f"--{k}", default=v)
+    return parser
+
+
+def save_checkpoint(state, filename):
+    torch.save(state, filename)
+
+
+class CosineAnnealingWarmRestarts(_LRScheduler):
+    r"""Set the learning rate of each parameter group using a cosine annealing
+    schedule, where :math:`\eta_{max}` is set to the initial lr, :math:`T_{cur}`
+    is the number of epochs since the last restart and :math:`T_{i}` is the number
+    of epochs between two warm restarts in SGDR:
+    .. math::
+        \eta_t = \eta_{min} + \frac{1}{2}(\eta_{max} - \eta_{min})\left(1 +
+        \cos\left(\frac{T_{cur}}{T_{i}}\pi\right)\right)
+    When :math:`T_{cur}=T_{i}`, set :math:`\eta_t = \eta_{min}`.
+    When :math:`T_{cur}=0` after restart, set :math:`\eta_t=\eta_{max}`.
+    It has been proposed in
+    `SGDR: Stochastic Gradient Descent with Warm Restarts`_.
+    Args:
+        optimizer (Optimizer): Wrapped optimizer.
+        T_0 (int): Number of iterations for the first restart.
+        T_mult (int, optional): A factor increases :math:`T_{i}` after a restart. Default: 1.
+        eta_min (float, optional): Minimum learning rate. Default: 0.
+        last_epoch (int, optional): The index of last epoch. Default: -1.
+        verbose (bool): If ``True``, prints a message to stdout for
+            each update. Default: ``False``.
+    .. _SGDR\: Stochastic Gradient Descent with Warm Restarts:
+        https://arxiv.org/abs/1608.03983
+    """
+
+    def __init__(self, optimizer, T_0, T_mult=1, eta_min=0, last_epoch=-1, verbose=False):
+        if T_0 <= 0 or not isinstance(T_0, int):
+            raise ValueError("Expected positive integer T_0, but got {}".format(T_0))
+        if T_mult < 1 or not isinstance(T_mult, int):
+            raise ValueError("Expected integer T_mult >= 1, but got {}".format(T_mult))
+        self.T_0 = T_0
+        self.T_i = T_0
+        self.T_mult = T_mult
+        self.eta_min = eta_min
+
+        self.T_cur = 0 if last_epoch < 0 else last_epoch
+        super(CosineAnnealingWarmRestarts, self).__init__(optimizer, last_epoch, verbose)
+
+    def get_lr(self):
+        if not self._get_lr_called_within_step:
+            warnings.warn("To get the last learning rate computed by the scheduler, "
+                          "please use `get_last_lr()`.", UserWarning)
+        return [self.eta_min + (base_lr - self.eta_min) * (1 + math.cos(math.pi * self.T_cur / self.T_i)) / 2
+                for base_lr in self.base_lrs]
+
+    def step(self, epoch=None):
+        """Step could be called after every batch update
+        Example:
+            >>> scheduler = CosineAnnealingWarmRestarts(optimizer, T_0, T_mult)
+            >>> iters = len(dataloader)
+            >>> for epoch in range(20):
+            >>>     for i, sample in enumerate(dataloader):
+            >>>         inputs, labels = sample['inputs'], sample['labels']
+            >>>         optimizer.zero_grad()
+            >>>         outputs = net(inputs)
+            >>>         loss = criterion(outputs, labels)
+            >>>         loss.backward()
+            >>>         optimizer.step()
+            >>>         scheduler.step(epoch + i / iters)
+        This function can be called in an interleaved way.
+        Example:
+            >>> scheduler = CosineAnnealingWarmRestarts(optimizer, T_0, T_mult)
+            >>> for epoch in range(20):
+            >>>     scheduler.step()
+            >>> scheduler.step(26)
+            >>> scheduler.step() # scheduler.step(27), instead of scheduler(20)
+        """
+        if epoch is None and self.last_epoch < 0:
+            epoch = 0
+        if epoch is None:
+            epoch = self.last_epoch + 1
+            self.T_cur = self.T_cur + 1
+            if self.T_cur >= self.T_i:
+                self.T_cur = self.T_cur - self.T_i
+                self.T_i = self.T_i * self.T_mult
+        else:
+            if epoch < 0:
+                raise ValueError("Expected non-negative epoch, but got {}".format(epoch))
+            if epoch >= self.T_0:
+                if self.T_mult == 1:
+                    self.T_cur = epoch % self.T_0
+                else:
+                    n = int(math.log((epoch / self.T_0 * (self.T_mult - 1) + 1), self.T_mult))
+                    self.T_cur = epoch - self.T_0 * (self.T_mult ** n - 1) / (self.T_mult - 1)
+                    self.T_i = self.T_0 * self.T_mult ** (n)
+            else:
+                self.T_i = self.T_0
+                self.T_cur = epoch
+        self.last_epoch = math.floor(epoch)
+
+        class _enable_get_lr_call:
+            def __init__(self, o):
+                self.o = o
+
+            def __enter__(self):
+                self.o._get_lr_called_within_step = True
+                return self
+
+            def __exit__(self, type, value, traceback):
+                self.o._get_lr_called_within_step = False
+                return self
+
+        with _enable_get_lr_call(self):
+            for i, data in enumerate(zip(self.optimizer.param_groups, self.get_lr())):
+                param_group, lr = data
+                param_group['lr'] = lr
+                self.print_lr(self.verbose, i, lr, epoch)
+        self._last_lr = [group['lr'] for group in self.optimizer.param_groups]
+
+
+def set_seed(seed):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)