From aaea6dc0a3aba1fa1239fe5b7abd53e0cff208e2 Mon Sep 17 00:00:00 2001 From: wupeiwei Date: Tue, 3 Aug 2021 18:16:31 +0800 Subject: [PATCH 1/9] ref: update scikit-image to 1.16.X --- codes/official_metrics/metrics.py | 2 +- requirements.txt | 6 ++++++ 2 files changed, 7 insertions(+), 1 deletion(-) create mode 100644 requirements.txt diff --git a/codes/official_metrics/metrics.py b/codes/official_metrics/metrics.py index f6dd094..12e66dd 100644 --- a/codes/official_metrics/metrics.py +++ b/codes/official_metrics/metrics.py @@ -4,7 +4,7 @@ import pandas as pd from LPIPSmodels import util import LPIPSmodels.dist_model as dm -from skimage.measure import compare_ssim +from skimage.metrics import structural_similarity as compare_ssim from absl import flags flags.DEFINE_string('output', None, 'the path of output directory') diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..ce3f8b0 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,6 @@ +numpy +matplotlib +opencv-python +pyyaml +lmdb +scikit-image>=0.16.0 From e5e3876e91b7a8e693e2ea940607b67e41197577 Mon Sep 17 00:00:00 2001 From: wupeiwei Date: Tue, 3 Aug 2021 18:18:11 +0800 Subject: [PATCH 2/9] ref: save jpeg image to lmdb. --- .gitignore | 12 +++++++++++ codes/data/base_dataset.py | 5 +++-- scripts/create_lmdb.py | 44 +++++++++++++++++++++++++++----------- 3 files changed, 46 insertions(+), 15 deletions(-) create mode 100644 .gitignore diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..b348a96 --- /dev/null +++ b/.gitignore @@ -0,0 +1,12 @@ +# python +__pycache__ + +# IDE +.vscode +.history +experiments_* +data/ + +# project +*.pth +results \ No newline at end of file diff --git a/codes/data/base_dataset.py b/codes/data/base_dataset.py index 3131e73..b6ddb73 100644 --- a/codes/data/base_dataset.py +++ b/codes/data/base_dataset.py @@ -1,5 +1,5 @@ +import cv2 import lmdb - import numpy as np from torch.utils.data import Dataset @@ -62,7 +62,8 @@ def parse_lmdb_key(key): def read_lmdb_frame(env, key, size): with env.begin(write=False) as txn: buf = txn.get(key.encode('ascii')) - frm = np.frombuffer(buf, dtype=np.uint8).reshape(*size) + frm = np.frombuffer(buf, dtype=np.uint8) + frm = cv2.imdecode(frm, cv2.IMREAD_COLOR) return frm def crop_sequence(self, **kwargs): diff --git a/scripts/create_lmdb.py b/scripts/create_lmdb.py index 8ac6994..96120ee 100644 --- a/scripts/create_lmdb.py +++ b/scripts/create_lmdb.py @@ -1,16 +1,22 @@ -import os -import os.path as osp import argparse import glob -import lmdb +import os +import os.path as osp import pickle import random import cv2 +import lmdb import numpy as np -def create_lmdb(dataset, raw_dir, lmdb_dir, filter_file=''): +def get_FileSize(filePath): + + fsize = os.path.getsize(filePath) + return round(fsize, 3) + + +def create_lmdb(dataset, raw_dir, lmdb_dir, zip_img=False, filter_file=''): assert dataset in ('VimeoTecoGAN', 'VimeoTecoGAN-sub'), f'Unknown Dataset: {dataset}' print(f'Creating lmdb for dataset: {dataset}') @@ -27,7 +33,11 @@ def create_lmdb(dataset, raw_dir, lmdb_dir, filter_file=''): nbytes = 0 for seq_dir in seq_dir_lst: frm_path_lst = sorted(glob.glob(osp.join(raw_dir, seq_dir, '*.png'))) - nbytes_per_frm = cv2.imread(frm_path_lst[0], cv2.IMREAD_UNCHANGED).nbytes + if not zip_img: + nbytes_per_frm = cv2.imread(frm_path_lst[0], cv2.IMREAD_UNCHANGED).nbytes + else: + nbytes_per_frm = get_FileSize(frm_path_lst[0]) + nbytes += len(frm_path_lst)*nbytes_per_frm alloc_size = round(1.5*nbytes) print(f'{alloc_size / (1 << 30):.2f} GB') @@ -54,7 +64,10 @@ def create_lmdb(dataset, raw_dir, lmdb_dir, filter_file=''): h, w, c = frm.shape key = f'{seq_dir}_{n_frm}x{h}x{w}_{i:04d}' - txn.put(key.encode('ascii'), frm) + if not zip_img: + txn.put(key.encode('ascii'), frm) + else: + txn.put(key.encode('ascii'), cv2.imencode('.jpg', frm)[1]) keys.append(key) # commit @@ -73,13 +86,15 @@ def create_lmdb(dataset, raw_dir, lmdb_dir, filter_file=''): pickle.dump(meta_info, open(osp.join(lmdb_dir, 'meta_info.pkl'), 'wb')) -def check_lmdb(dataset, lmdb_dir, display=True): +def check_lmdb(dataset, lmdb_dir, zip_img=False, display=True): def visualize(win, img): if display: cv2.namedWindow(win, 0) cv2.resizeWindow(win, img.shape[-2], img.shape[-3]) - cv2.imshow(win, img[..., ::-1]) + size = img.shape + dst = cv2.resize(img, (size[1], size[0]), fx=0.5, fy=0.5) + cv2.imshow(win, dst[..., ::-1]) cv2.waitKey(0) cv2.destroyAllWindows() else: @@ -108,7 +123,11 @@ def visualize(win, img): with env.begin() as txn: buf = txn.get(key.encode('ascii')) - val = np.frombuffer(buf, dtype=np.uint8).reshape(*sz) # HWC + if not zip_img: + val = np.frombuffer(buf, dtype=np.uint8).reshape(*sz) # HWC + else: + val = np.frombuffer(buf, dtype=np.uint8) + val = cv2.imdecode(val, cv2.IMREAD_COLOR) visualize(key, val) @@ -131,10 +150,9 @@ def visualize(win, img): if osp.exists(lmdb_dir): print(f'Dataset [{args.dataset}] already exists') print('Checking the LMDB dataset ...') - check_lmdb(args.dataset, lmdb_dir) + check_lmdb(args.dataset, lmdb_dir, zip_img=True, display=False) else: - create_lmdb(args.dataset, raw_dir, lmdb_dir, filter_file) + create_lmdb(args.dataset, raw_dir, lmdb_dir, zip_img=True) print('Checking the LMDB dataset ...') - check_lmdb(args.dataset, lmdb_dir) - + check_lmdb(args.dataset, lmdb_dir, zip_img=True, display=False) From 49b2f61a41a28e75d34b253901250f7c04f68715 Mon Sep 17 00:00:00 2001 From: wupeiwei Date: Tue, 3 Aug 2021 18:24:32 +0800 Subject: [PATCH 3/9] ref: Avoid downloading data set files repeatedly. --- scripts/download/download_datasets.sh | 27 +++++++++++++++++++++------ 1 file changed, 21 insertions(+), 6 deletions(-) diff --git a/scripts/download/download_datasets.sh b/scripts/download/download_datasets.sh index bd242c2..d58f807 100644 --- a/scripts/download/download_datasets.sh +++ b/scripts/download/download_datasets.sh @@ -44,7 +44,9 @@ if [ ! -d "./data/Vid4/GT" ]; then echo ">>> Start to download [Vid4 GT] dataset" mkdir -p ${DATA_DIR} - download_large_file ${DATA_DIR} ${FID} GT + if [ ! -e "${DATA_DIR}/GT.zip" ]; then + download_large_file ${DATA_DIR} ${FID} GT + fi check_md5 ${DATA_DIR}/GT.zip ${MD5} unzip ${DATA_DIR}/GT.zip -d ${DATA_DIR} && rm ${DATA_DIR}/GT.zip fi @@ -60,7 +62,9 @@ if [ ! -d "./data/ToS3/GT" ]; then echo ">>> Start to download [ToS3 GT] dataset" mkdir -p ${DATA_DIR} - download_large_file ${DATA_DIR} ${FID} GT + if [ ! -e "${DATA_DIR}/GT.zip" ]; then + download_large_file ${DATA_DIR} ${FID} GT + fi check_md5 ${DATA_DIR}/GT.zip ${MD5} unzip ${DATA_DIR}/GT.zip -d ${DATA_DIR} && rm ${DATA_DIR}/GT.zip fi @@ -75,7 +79,10 @@ if [ $1 == BD ]; then echo ">>> Start to download [Vid4 LR] dataset (BD degradation)" - download_small_file ${DATA_DIR} ${FID} Gaussian4xLR + if [ ! -e "${DATA_DIR}/Gaussian4xLR.zip" ]; then + download_small_file ${DATA_DIR} ${FID} Gaussian4xLR + fi + check_md5 ${DATA_DIR}/Gaussian4xLR.zip ${MD5} unzip ${DATA_DIR}/Gaussian4xLR.zip -d ${DATA_DIR} && rm ${DATA_DIR}/Gaussian4xLR.zip fi @@ -90,7 +97,9 @@ if [ $1 == BD ]; then echo ">>> Start to download [ToS3 LR] dataset (BD degradation)" - download_large_file ${DATA_DIR} ${FID} Gaussian4xLR + if [ ! -e "${DATA_DIR}/Gaussian4xLR.zip" ]; then + download_large_file ${DATA_DIR} ${FID} Gaussian4xLR + fi check_md5 ${DATA_DIR}/Gaussian4xLR.zip ${MD5} unzip ${DATA_DIR}/Gaussian4xLR.zip -d ${DATA_DIR} && rm ${DATA_DIR}/Gaussian4xLR.zip fi @@ -104,7 +113,10 @@ elif [ $1 == BI ]; then echo ">>> Start to download [Vid4 LR] dataset (BI degradation)" - download_small_file ${DATA_DIR} ${FID} Bicubic4xLR + if [ ! -e "${DATA_DIR}/Bicubic4xLR.zip" ]; then + download_small_file ${DATA_DIR} ${FID} Bicubic4xLR + fi + check_md5 ${DATA_DIR}/Bicubic4xLR.zip ${MD5} unzip ${DATA_DIR}/Bicubic4xLR.zip -d ${DATA_DIR} && rm ${DATA_DIR}/Bicubic4xLR.zip fi @@ -119,7 +131,10 @@ elif [ $1 == BI ]; then echo ">>> Start to download [ToS3 LR] dataset (BI degradation)" - download_large_file ${DATA_DIR} ${FID} Bicubic4xLR + if [ ! -e "${DATA_DIR}/Bicubic4xLR.zip" ]; then + download_large_file ${DATA_DIR} ${FID} Bicubic4xLR + fi + check_md5 ${DATA_DIR}/Bicubic4xLR.zip ${MD5} unzip ${DATA_DIR}/Bicubic4xLR.zip -d ${DATA_DIR} && rm ${DATA_DIR}/Bicubic4xLR.zip fi From 25ddae07eea3625a7478f900c3901c9c9a70690d Mon Sep 17 00:00:00 2001 From: wupeiwei Date: Tue, 3 Aug 2021 20:58:44 +0800 Subject: [PATCH 4/9] fix: warning of lr.step() before opt.stop(). --- codes/main.py | 28 +++++++++++++++++----------- 1 file changed, 17 insertions(+), 11 deletions(-) diff --git a/codes/main.py b/codes/main.py index 699424a..fac8c2c 100644 --- a/codes/main.py +++ b/codes/main.py @@ -1,15 +1,15 @@ +import math import os import os.path as osp -import math import time import torch from data import create_dataloader +from metrics.metric_calculator import MetricCalculator from models import define_model from models.networks import define_generator -from metrics.metric_calculator import MetricCalculator -from utils import dist_utils, base_utils, data_utils +from utils import base_utils, data_utils, dist_utils def train(opt): @@ -44,10 +44,8 @@ def train(opt): # update iter iter += 1 curr_iter = start_iter + iter - if iter > total_iter: break - - # update learning rate - model.update_learning_rate() + if iter > total_iter: + break # prepare data model.prepare_data(data) @@ -55,19 +53,26 @@ def train(opt): # train a mini-batch model.train() + # update learning rate + model.update_learning_rate() + # update running log model.reduce_log() # for distributed training model.update_running_log() # print messages if log_freq > 0 and curr_iter % log_freq == 0: - msg = model.get_format_msg(epoch, curr_iter) + msg = model.get_format_msg(epoch, curr_iter, total_epoch, total_iter) base_utils.log_info(msg) # save model if ckpt_freq > 0 and curr_iter % ckpt_freq == 0: model.save(curr_iter) + filename = f'G_iter{curr_iter}.pth' + save_path = osp.join(model.ckpt_dir, filename) + base_utils.log_info(f"save model in {save_path}") + # evaluate model if test_freq > 0 and curr_iter % test_freq == 0: # set model index @@ -76,7 +81,8 @@ def train(opt): # for each testset for dataset_idx in sorted(opt['dataset'].keys()): # select test dataset - if 'test' not in dataset_idx: continue + if 'test' not in dataset_idx: + continue ds_name = opt['dataset'][dataset_idx]['name'] base_utils.log_info(f'Testing on {ds_name} dataset') @@ -135,7 +141,7 @@ def test(opt): for load_path in opt['model']['generator']['load_path_lst']: # set model index model_idx = osp.splitext(osp.split(load_path)[-1])[0] - + # log base_utils.log_info(f'\n{"="*40}') base_utils.log_info(f'Testing model: {model_idx}') @@ -175,7 +181,7 @@ def test(opt): data_utils.save_sequence( res_seq_dir, hr_seq, data['frm_idx'], to_bgr=True) - base_utils.log_info('-'*40) + base_utils.log_info('-' * 40) # logging base_utils.log_info(f'Finish testing\n{"="*40}') From 36b457dae474f716e2183881630715c4fe2d988d Mon Sep 17 00:00:00 2001 From: wupeiwei Date: Tue, 3 Aug 2021 20:58:52 +0800 Subject: [PATCH 5/9] ref: add logs. --- codes/models/base_model.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/codes/models/base_model.py b/codes/models/base_model.py index d530df4..42b8f49 100644 --- a/codes/models/base_model.py +++ b/codes/models/base_model.py @@ -154,9 +154,9 @@ def get_current_log(self): def get_running_log(self): return self.running_log_dict - def get_format_msg(self, epoch, iter): + def get_format_msg(self, epoch, iter, total_epoch, total_iter): # generic info - msg = f'[epoch: {epoch} | iter: {iter}' + msg = f'[epoch: {epoch:3d}/{total_epoch} | iter: {iter:6d}/{total_iter}' for lr_type, lr in self.get_current_learning_rate().items(): msg += f' | {lr_type}: {lr:.2e}' msg += '] ' From 415e5b22ce98c443eb345dde98f7b916ca826c2f Mon Sep 17 00:00:00 2001 From: wupeiwei Date: Tue, 3 Aug 2021 20:59:00 +0800 Subject: [PATCH 6/9] ref: init lmdb in __init__. --- codes/data/unpaired_lmdb_dataset.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/codes/data/unpaired_lmdb_dataset.py b/codes/data/unpaired_lmdb_dataset.py index aed356a..bb2f320 100644 --- a/codes/data/unpaired_lmdb_dataset.py +++ b/codes/data/unpaired_lmdb_dataset.py @@ -22,12 +22,12 @@ def __init__(self, data_opt, **kwargs): # use partial videos if hasattr(self, 'filter_file') and self.filter_file is not None: with open(self.filter_file, 'r') as f: - sel_seqs = { line.strip() for line in f } + sel_seqs = {line.strip() for line in f} self.keys = list(filter( lambda x: self.parse_lmdb_key(x)[0] in sel_seqs, self.keys)) # register parameters - self.env = None + self.env = self.init_lmdb(self.seq_dir) def __len__(self): return len(self.keys) @@ -64,7 +64,7 @@ def __getitem__(self, item): frms.append(frm[:, top: top + c_h, left: left + c_w].copy()) else: # read frames - for i in range(cur_frm, cur_frm + self.tempo_extent): + def get_frames(i): if i >= tot_frm: # reflect temporal paddding, e.g., (0,1,2) -> (0,1,2,1,0) key = '{}_{}x{}x{}_{:04d}'.format( @@ -75,6 +75,10 @@ def __getitem__(self, item): frm = self.read_lmdb_frame(self.env, key, size=(h, w, c)) frm = frm.transpose(2, 0, 1) # chw|rgb|uint8 + return frm + + for i in range(cur_frm, cur_frm + self.tempo_extent): + frm = get_frames(i) frms.append(frm) frms = np.stack(frms) # tchw|rgb|uint8 From a752726588c2b5966bd6cf53a679a588f0764cd4 Mon Sep 17 00:00:00 2001 From: wupeiwei Date: Tue, 3 Aug 2021 20:59:10 +0800 Subject: [PATCH 7/9] flake8 --- .flake8 | 28 ++ codes/data/paired_lmdb_dataset.py | 3 +- codes/main.py | 1 - codes/metrics/LPIPS/models/base_model.py | 18 +- codes/metrics/LPIPS/models/dist_model.py | 119 +++++---- codes/metrics/LPIPS/models/networks_basic.py | 137 +++++----- .../LPIPS/models/pretrained_networks.py | 21 +- codes/metrics/metric_calculator.py | 8 +- codes/metrics/model_summary.py | 9 +- codes/models/base_model.py | 19 +- codes/models/networks/tecogan_nets.py | 33 ++- codes/models/networks/vgg_nets.py | 2 +- codes/models/vsr_model.py | 3 +- codes/models/vsrgan_model.py | 8 +- .../LPIPSmodels/base_model.py | 20 +- .../LPIPSmodels/dist_model.py | 151 +++++------ .../LPIPSmodels/networks_basic.py | 215 ++++++++------- .../LPIPSmodels/pretrained_networks.py | 23 +- codes/official_metrics/LPIPSmodels/util.py | 177 +++++++------ codes/official_metrics/evaluate.py | 4 +- codes/official_metrics/metrics.py | 244 ++++++++++-------- codes/utils/data_utils.py | 7 +- codes/utils/net_utils.py | 15 +- scripts/create_lmdb.py | 4 +- scripts/monitor_training.py | 6 +- scripts/resize_BD.py | 5 +- 26 files changed, 683 insertions(+), 597 deletions(-) create mode 100644 .flake8 diff --git a/.flake8 b/.flake8 new file mode 100644 index 0000000..9ebb404 --- /dev/null +++ b/.flake8 @@ -0,0 +1,28 @@ +[flake8] +select = B,C,E,F,P,T4,W,B9 +max-line-length = 120 +# C408 ignored because we like the dict keyword argument syntax +# E501 is not flexible enough, we're using B950 instead +ignore = + E203,E305,E402,E501,E721,E741,F403,F405,F821,F999,W503,W504,C408,E302,W291,E303,W605,E722,E266,E265 + # shebang has extra meaning in fbcode lints, so I think it's not worth trying + # to line this up with executable bit + EXE001, + # these ignores are from flake8-bugbear; please fix! + B007,B008, + # these ignores are from flake8-comprehensions; please fix! + C400,C401,C402,C403,C404,C405,C407,C411,C413,C414,C415 +exclude = + .history, + .vscode, + tests/, + third_party, + build, + *.pyi, + .git, + .ipynb*, + __init__.py, + experiments_*, + data/ + + diff --git a/codes/data/paired_lmdb_dataset.py b/codes/data/paired_lmdb_dataset.py index 2bc4525..e1ac00f 100644 --- a/codes/data/paired_lmdb_dataset.py +++ b/codes/data/paired_lmdb_dataset.py @@ -2,7 +2,6 @@ import pickle import random -import cv2 import numpy as np import torch @@ -29,7 +28,7 @@ def __init__(self, data_opt, **kwargs): # use partial videos if hasattr(self, 'filter_file') and self.filter_file is not None: with open(self.filter_file, 'r') as f: - sel_seqs = { line.strip() for line in f } + sel_seqs = {line.strip() for line in f} self.gt_lr_keys = list(filter( lambda x: self.parse_lmdb_key(x[0])[0] in sel_seqs, self.gt_lr_keys)) diff --git a/codes/main.py b/codes/main.py index fac8c2c..ea108e0 100644 --- a/codes/main.py +++ b/codes/main.py @@ -1,5 +1,4 @@ import math -import os import os.path as osp import time diff --git a/codes/metrics/LPIPS/models/base_model.py b/codes/metrics/LPIPS/models/base_model.py index 273fd8c..d97d200 100644 --- a/codes/metrics/LPIPS/models/base_model.py +++ b/codes/metrics/LPIPS/models/base_model.py @@ -1,13 +1,13 @@ import os + +import numpy as np import torch -from torch.autograd import Variable -from pdb import set_trace as st class BaseModel(): def __init__(self): - pass; - + pass + def name(self): return 'BaseModel' @@ -18,9 +18,6 @@ def initialize(self, use_gpu=True, gpu_ids=[0]): def forward(self): pass - def get_image_paths(self): - pass - def optimize_parameters(self): pass @@ -43,7 +40,7 @@ def save_network(self, network, path, network_label, epoch_label): def load_network(self, network, network_label, epoch_label): save_filename = '%s_net_%s.pth' % (epoch_label, network_label) save_path = os.path.join(self.save_dir, save_filename) - print('Loading network from %s'%save_path) + print('Loading network from %s' % save_path) network.load_state_dict(torch.load(save_path)) def update_learning_rate(): @@ -53,6 +50,5 @@ def get_image_paths(self): return self.image_paths def save_done(self, flag=False): - np.save(os.path.join(self.save_dir, 'done_flag'),flag) - np.savetxt(os.path.join(self.save_dir, 'done_flag'),[flag,],fmt='%i') - + np.save(os.path.join(self.save_dir, 'done_flag'), flag) + np.savetxt(os.path.join(self.save_dir, 'done_flag'), [flag, ], fmt='%i') diff --git a/codes/metrics/LPIPS/models/dist_model.py b/codes/metrics/LPIPS/models/dist_model.py index 21a1d69..6daa2fd 100644 --- a/codes/metrics/LPIPS/models/dist_model.py +++ b/codes/metrics/LPIPS/models/dist_model.py @@ -1,33 +1,29 @@ from __future__ import absolute_import -import sys -import numpy as np -import torch -from torch import nn import os from collections import OrderedDict -from torch.autograd import Variable -import itertools -from .base_model import BaseModel + +import numpy as np +import torch from scipy.ndimage import zoom -import fractions -import functools -import skimage.transform +from torch.autograd import Variable from tqdm import tqdm -from IPython import embed +import models as util from . import networks_basic as networks -import models as util +from .base_model import BaseModel + class DistModel(BaseModel): def name(self): return self.model_name - def initialize(self, model='net-lin', net='alex', colorspace='Lab', pnet_rand=False, pnet_tune=False, model_path=None, - use_gpu=True, printNet=False, spatial=False, - is_train=False, lr=.0001, beta1=0.5, version='0.1', gpu_ids=[0]): + def initialize( + self, model='net-lin', net='alex', colorspace='Lab', pnet_rand=False, pnet_tune=False, model_path=None, + use_gpu=True, printNet=False, spatial=False, is_train=False, lr=.0001, beta1=0.5, version='0.1', + gpu_ids=[0]): ''' INPUTS model - ['net-lin'] for linearly calibrated network @@ -53,17 +49,18 @@ def initialize(self, model='net-lin', net='alex', colorspace='Lab', pnet_rand=Fa self.is_train = is_train self.spatial = spatial self.gpu_ids = gpu_ids - self.model_name = '%s [%s]'%(model,net) + self.model_name = '%s [%s]' % (model, net) - if(self.model == 'net-lin'): # pretrained net + linear layer + if(self.model == 'net-lin'): # pretrained net + linear layer self.net = networks.PNetLin(pnet_rand=pnet_rand, pnet_tune=pnet_tune, pnet_type=net, - use_dropout=True, spatial=spatial, version=version, lpips=True) + use_dropout=True, spatial=spatial, version=version, lpips=True) #kw = {} #if not use_gpu: # kw['map_location'] = 'cpu' if(model_path is None): import inspect - model_path = os.path.abspath(os.path.join(inspect.getfile(self.initialize), '..', 'weights/v%s/%s.pth'%(version,net))) + model_path = os.path.abspath(os.path.join(inspect.getfile(self.initialize), + '..', 'weights/v%s/%s.pth' % (version, net))) if(not is_train): # self.net.load_state_dict(torch.load(model_path, **kw), strict=False) @@ -71,34 +68,34 @@ def initialize(self, model='net-lin', net='alex', colorspace='Lab', pnet_rand=Fa model_path, map_location=lambda storage, loc: storage) self.net.load_state_dict(state_dict, strict=False) - elif(self.model=='net'): # pretrained network + elif(self.model == 'net'): # pretrained network self.net = networks.PNetLin(pnet_rand=pnet_rand, pnet_type=net, lpips=False) - elif(self.model in ['L2','l2']): - self.net = networks.L2(use_gpu=use_gpu,colorspace=colorspace) # not really a network, only for testing + elif(self.model in ['L2', 'l2']): + self.net = networks.L2(use_gpu=use_gpu, colorspace=colorspace) # not really a network, only for testing self.model_name = 'L2' - elif(self.model in ['DSSIM','dssim','SSIM','ssim']): - self.net = networks.DSSIM(use_gpu=use_gpu,colorspace=colorspace) + elif(self.model in ['DSSIM', 'dssim', 'SSIM', 'ssim']): + self.net = networks.DSSIM(use_gpu=use_gpu, colorspace=colorspace) self.model_name = 'SSIM' else: raise ValueError("Model [%s] not recognized." % self.model) self.parameters = list(self.net.parameters()) - if self.is_train: # training mode + if self.is_train: # training mode # extra network on top to go from distances (d0,d1) => predicted human judgment (h*) self.rankLoss = networks.BCERankingLoss() self.parameters += list(self.rankLoss.net.parameters()) self.lr = lr self.old_lr = lr self.optimizer_net = torch.optim.Adam(self.parameters, lr=lr, betas=(beta1, 0.999)) - else: # test mode + else: # test mode self.net.eval() if(use_gpu): self.net.to(gpu_ids[0]) #self.net = torch.nn.DataParallel(self.net, device_ids=gpu_ids) if(self.is_train): - self.rankLoss = self.rankLoss.to(device=gpu_ids[0]) # just put this on GPU0 + self.rankLoss = self.rankLoss.to(device=gpu_ids[0]) # just put this on GPU0 if(printNet): print('---------- Networks initialized -------------') @@ -125,8 +122,8 @@ def optimize_parameters(self): def clamp_weights(self): for module in self.net.modules(): - if(hasattr(module, 'weight') and module.kernel_size==(1,1)): - module.weight.data = torch.clamp(module.weight.data,min=0) + if(hasattr(module, 'weight') and module.kernel_size == (1, 1)): + module.weight.data = torch.clamp(module.weight.data, min=0) def set_input(self, data): self.input_ref = data['ref'] @@ -140,36 +137,36 @@ def set_input(self, data): self.input_p1 = self.input_p1.to(device=self.gpu_ids[0]) self.input_judge = self.input_judge.to(device=self.gpu_ids[0]) - self.var_ref = Variable(self.input_ref,requires_grad=True) - self.var_p0 = Variable(self.input_p0,requires_grad=True) - self.var_p1 = Variable(self.input_p1,requires_grad=True) + self.var_ref = Variable(self.input_ref, requires_grad=True) + self.var_p0 = Variable(self.input_p0, requires_grad=True) + self.var_p1 = Variable(self.input_p1, requires_grad=True) - def forward_train(self): # run forward pass + def forward_train(self): # run forward pass # print(self.net.module.scaling_layer.shift) # print(torch.norm(self.net.module.net.slice1[0].weight).item(), torch.norm(self.net.module.lin0.model[1].weight).item()) self.d0 = self.forward(self.var_ref, self.var_p0) self.d1 = self.forward(self.var_ref, self.var_p1) - self.acc_r = self.compute_accuracy(self.d0,self.d1,self.input_judge) + self.acc_r = self.compute_accuracy(self.d0, self.d1, self.input_judge) - self.var_judge = Variable(1.*self.input_judge).view(self.d0.size()) + self.var_judge = Variable(1. * self.input_judge).view(self.d0.size()) - self.loss_total = self.rankLoss.forward(self.d0, self.d1, self.var_judge*2.-1.) + self.loss_total = self.rankLoss.forward(self.d0, self.d1, self.var_judge * 2. - 1.) return self.loss_total def backward_train(self): torch.mean(self.loss_total).backward() - def compute_accuracy(self,d0,d1,judge): + def compute_accuracy(self, d0, d1, judge): ''' d0, d1 are Variables, judge is a Tensor ''' - d1_lt_d0 = (d1 %f' % (type,self.old_lr, lr)) + print('update lr [%s] decay: %f -> %f' % (type, self.old_lr, lr)) self.old_lr = lr + def score_2afc_dataset(data_loader, func, name=''): ''' Function computes Two Alternative Forced Choice (2AFC) score using distance function 'func' in dataset 'data_loader' @@ -232,16 +230,17 @@ def score_2afc_dataset(data_loader, func, name=''): gts = [] for data in tqdm(data_loader.load_data(), desc=name): - d0s+=func(data['ref'],data['p0']).data.cpu().numpy().flatten().tolist() - d1s+=func(data['ref'],data['p1']).data.cpu().numpy().flatten().tolist() - gts+=data['judge'].cpu().numpy().flatten().tolist() + d0s += func(data['ref'], data['p0']).data.cpu().numpy().flatten().tolist() + d1s += func(data['ref'], data['p1']).data.cpu().numpy().flatten().tolist() + gts += data['judge'].cpu().numpy().flatten().tolist() d0s = np.array(d0s) d1s = np.array(d1s) gts = np.array(gts) - scores = (d0s predicted human judgment (h*) self.rankLoss = networks.BCERankingLoss(use_gpu=use_gpu) - self.parameters+=self.rankLoss.parameters + self.parameters += self.rankLoss.parameters self.lr = lr self.old_lr = lr self.optimizer_net = torch.optim.Adam(self.parameters, lr=lr, betas=(beta1, 0.999)) - else: # test mode + else: # test mode self.net.eval() if(printNet): @@ -102,11 +105,11 @@ def initialize(self, model='net-lin', net='alex', pnet_rand=False, pnet_tune=Fal networks.print_network(self.net) print('-----------------------------------------------') - def forward_pair(self,in1,in2,retPerLayer=False): + def forward_pair(self, in1, in2, retPerLayer=False): if(retPerLayer): - return self.net.forward(in1,in2, retPerLayer=True) + return self.net.forward(in1, in2, retPerLayer=True) else: - return self.net.forward(in1,in2) + return self.net.forward(in1, in2) def forward(self, in0, in1, retNumpy=True): ''' Function computes the distance between image patches in0 and in1 @@ -124,8 +127,8 @@ def forward(self, in0, in1, retNumpy=True): self.input_ref = self.input_ref.cuda() self.input_p0 = self.input_p0.cuda() - self.var_ref = Variable(self.input_ref,requires_grad=True) - self.var_p0 = Variable(self.input_p0,requires_grad=True) + self.var_ref = Variable(self.input_ref, requires_grad=True) + self.var_p0 = Variable(self.input_p0, requires_grad=True) self.d0 = self.forward_pair(self.var_ref, self.var_p0) self.loss_total = self.d0 @@ -137,7 +140,8 @@ def convert_output(d0): ans = ans.flatten() else: assert(ans.shape[0] == 1 and len(ans.shape) == 4) - return ans[0,...].transpose([1, 2, 0]) # Reshape to usual numpy image format: (height, width, channels) + # Reshape to usual numpy image format: (height, width, channels) + return ans[0, ...].transpose([1, 2, 0]) return ans else: return d0 @@ -147,12 +151,13 @@ def convert_output(d0): spatial_shape = self.spatial_shape if spatial_shape is None: if(self.spatial_factor is None): - spatial_shape = (in0.size()[2],in0.size()[3]) + spatial_shape = (in0.size()[2], in0.size()[3]) else: - spatial_shape = (max([x.shape[0] for x in L])*self.spatial_factor, max([x.shape[1] for x in L])*self.spatial_factor) - + spatial_shape = (max([x.shape[0] for x in L]) * self.spatial_factor, + max([x.shape[1] for x in L]) * self.spatial_factor) + L = [skimage.transform.resize(x, spatial_shape, order=self.spatial_order, mode='edge') for x in L] - + L = np.mean(np.concatenate(L, 2) * len(L), 2) return L else: @@ -168,8 +173,8 @@ def optimize_parameters(self): def clamp_weights(self): for module in self.net.modules(): - if(hasattr(module, 'weight') and module.kernel_size==(1,1)): - module.weight.data = torch.clamp(module.weight.data,min=0) + if(hasattr(module, 'weight') and module.kernel_size == (1, 1)): + module.weight.data = torch.clamp(module.weight.data, min=0) def set_input(self, data): self.input_ref = data['ref'] @@ -183,33 +188,33 @@ def set_input(self, data): self.input_p1 = self.input_p1.cuda() self.input_judge = self.input_judge.cuda() - self.var_ref = Variable(self.input_ref,requires_grad=True) - self.var_p0 = Variable(self.input_p0,requires_grad=True) - self.var_p1 = Variable(self.input_p1,requires_grad=True) + self.var_ref = Variable(self.input_ref, requires_grad=True) + self.var_p0 = Variable(self.input_p0, requires_grad=True) + self.var_p1 = Variable(self.input_p1, requires_grad=True) - def forward_train(self): # run forward pass + def forward_train(self): # run forward pass self.d0 = self.forward_pair(self.var_ref, self.var_p0) self.d1 = self.forward_pair(self.var_ref, self.var_p1) - self.acc_r = self.compute_accuracy(self.d0,self.d1,self.input_judge) + self.acc_r = self.compute_accuracy(self.d0, self.d1, self.input_judge) # var_judge - self.var_judge = Variable(1.*self.input_judge).view(self.d0.size()) + self.var_judge = Variable(1. * self.input_judge).view(self.d0.size()) - self.loss_total = self.rankLoss.forward(self.d0, self.d1, self.var_judge*2.-1.) + self.loss_total = self.rankLoss.forward(self.d0, self.d1, self.var_judge * 2. - 1.) return self.loss_total def backward_train(self): torch.mean(self.loss_total).backward() - def compute_accuracy(self,d0,d1,judge): + def compute_accuracy(self, d0, d1, judge): ''' d0, d1 are Variables, judge is a Tensor ''' - d1_lt_d0 = (d1 %f' % (type,self.old_lr, lr)) + print('update lr [%s] decay: %f -> %f' % (type, self.old_lr, lr)) self.old_lr = lr - -def score_2afc_dataset(data_loader,func): +def score_2afc_dataset(data_loader, func): ''' Function computes Two Alternative Forced Choice (2AFC) score using distance function 'func' in dataset 'data_loader' INPUTS @@ -271,20 +275,21 @@ def score_2afc_dataset(data_loader,func): gts = [] # bar = pb.ProgressBar(max_value=data_loader.load_data().__len__()) - for (i,data) in enumerate(data_loader.load_data()): - d0s+=func(data['ref'],data['p0']).tolist() - d1s+=func(data['ref'],data['p1']).tolist() - gts+=data['judge'].cpu().numpy().flatten().tolist() + for (i, data) in enumerate(data_loader.load_data()): + d0s += func(data['ref'], data['p0']).tolist() + d1s += func(data['ref'], data['p1']).tolist() + gts += data['judge'].cpu().numpy().flatten().tolist() # bar.update(i) d0s = np.array(d0s) d1s = np.array(d1s) gts = np.array(gts) - scores = (d0s ori_h - 16): h = h - 32 while(w > ori_w - 16): w = w - 32 - + y = (ori_h - h) // 2 x = (ori_w - w) // 2 - crop_img = img[y:y+h, x:x+w] + crop_img = img[y:y + h, x:x + w] return crop_img, y, x + class Logger(object): def __init__(self): self.terminal = sys.stdout filename = "metricsfile.txt" - self.log = open(os.path.join(FLAGS.output, filename), "a") + self.log = open(os.path.join(FLAGS.output, filename), "a") + def write(self, message): self.terminal.write(message) - self.log.write(message) + self.log.write(message) + def flush(self): self.log.flush() - + + sys.stdout = Logger() print_configuration_op(FLAGS) @@ -112,131 +127,132 @@ def flush(self): model = dm.DistModel() -model.initialize(model='net-lin',net='alex',use_gpu=True) +model.initialize(model='net-lin', net='alex', use_gpu=True) cutfr = 2 # maxV = 0.4, for line 154-166 -keys = ["PSNR", "SSIM", "LPIPS", "tOF", "tLP100"] # keys = ["LPIPS"] -sum_dict = dict.fromkeys(["FrameAvg_"+_ for _ in keys], 0) -len_dict = dict.fromkeys(keys, 0) -avg_dict = dict.fromkeys(["Avg_"+_ for _ in keys], 0) -folder_dict = dict.fromkeys(["FolderAvg_"+_ for _ in keys], 0) +keys = ["PSNR", "SSIM", "LPIPS", "tOF", "tLP100"] # keys = ["LPIPS"] +sum_dict = dict.fromkeys(["FrameAvg_" + _ for _ in keys], 0) +len_dict = dict.fromkeys(keys, 0) +avg_dict = dict.fromkeys(["Avg_" + _ for _ in keys], 0) +folder_dict = dict.fromkeys(["FolderAvg_" + _ for _ in keys], 0) for folder_i in range(folder_n): result = listPNGinDir(result_list[folder_i]) target = listPNGinDir(target_list[folder_i]) image_no = len(target) - + list_dict = {} for key_i in keys: list_dict[key_i] = [] - - for i in range(cutfr, image_no-cutfr): - output_img = cv2.imread(result[i])[:,:,::-1] - target_img = cv2.imread(target[i])[:,:,::-1] - msg = "frame %d, tar %s, out %s, "%(i, str(target_img.shape), str(output_img.shape)) - if( target_img.shape[0] < output_img.shape[0]) or ( target_img.shape[1] < output_img.shape[1]): # target is not dividable by 4 - output_img = output_img[:target_img.shape[0],:target_img.shape[1]] - if( target_img.shape[0] > output_img.shape[0]) or ( target_img.shape[1] > output_img.shape[1]): # target is not dividable by 4 - target_img = target_img[:output_img.shape[0],:output_img.shape[1]] + + for i in range(cutfr, image_no - cutfr): + output_img = cv2.imread(result[i])[:, :, ::-1] + target_img = cv2.imread(target[i])[:, :, ::-1] + msg = "frame %d, tar %s, out %s, " % (i, str(target_img.shape), str(output_img.shape)) + if(target_img.shape[0] < output_img.shape[0]) or (target_img.shape[1] < output_img.shape[1]): # target is not dividable by 4 + output_img = output_img[:target_img.shape[0], :target_img.shape[1]] + if(target_img.shape[0] > output_img.shape[0]) or (target_img.shape[1] > output_img.shape[1]): # target is not dividable by 4 + target_img = target_img[:output_img.shape[0], :output_img.shape[1]] #print(result[i]) - - if "tOF" in keys:# tOF + + if "tOF" in keys: # tOF output_grey = cv2.cvtColor(output_img, cv2.COLOR_RGB2GRAY) target_grey = cv2.cvtColor(target_img, cv2.COLOR_RGB2GRAY) - if (i > cutfr): # temporal metrics - target_OF=cv2.calcOpticalFlowFarneback(pre_tar_grey, target_grey, None, 0.5, 3, 15, 3, 5, 1.2, 0) - output_OF=cv2.calcOpticalFlowFarneback(pre_out_grey, output_grey, None, 0.5, 3, 15, 3, 5, 1.2, 0) + if (i > cutfr): # temporal metrics + target_OF = cv2.calcOpticalFlowFarneback(pre_tar_grey, target_grey, None, 0.5, 3, 15, 3, 5, 1.2, 0) + output_OF = cv2.calcOpticalFlowFarneback(pre_out_grey, output_grey, None, 0.5, 3, 15, 3, 5, 1.2, 0) target_OF, ofy, ofx = crop_8x8(target_OF) output_OF, ofy, ofx = crop_8x8(output_OF) OF_diff = np.absolute(target_OF - output_OF) - if False: # for motion visualization - tOFpath = os.path.join(FLAGS.output,"%03d_tOF"%folder_i) - if(not os.path.exists(tOFpath)): os.mkdir(tOFpath) + if False: # for motion visualization + tOFpath = os.path.join(FLAGS.output, "%03d_tOF" % folder_i) + if(not os.path.exists(tOFpath)): + os.mkdir(tOFpath) hsv = np.zeros_like(output_img) - hsv[...,1] = 255 - out_path = os.path.join(tOFpath, "flow_%04d.jpg" %i) - mag, ang = cv2.cartToPolar(OF_diff[...,0], OF_diff[...,1]) + hsv[..., 1] = 255 + out_path = os.path.join(tOFpath, "flow_%04d.jpg" % i) + mag, ang = cv2.cartToPolar(OF_diff[..., 0], OF_diff[..., 1]) # print("tar max %02.6f, min %02.6f, avg %02.6f" % (mag.max(), mag.min(), mag.mean())) - mag = np.clip(mag, 0.0, maxV)/maxV - hsv[...,0] = ang*180/np.pi/2 - hsv[...,2] = mag * 255.0 # - bgr = cv2.cvtColor(hsv,cv2.COLOR_HSV2BGR) + mag = np.clip(mag, 0.0, maxV) / maxV + hsv[..., 0] = ang * 180 / np.pi / 2 + hsv[..., 2] = mag * 255.0 + bgr = cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR) cv2.imwrite(out_path, bgr) - - OF_diff = np.sqrt(np.sum(OF_diff * OF_diff, axis = -1)) # l1 vector norm + + OF_diff = np.sqrt(np.sum(OF_diff * OF_diff, axis=-1)) # l1 vector norm # OF_diff, ofy, ofx = crop_8x8(OF_diff) - list_dict["tOF"].append( OF_diff.mean() ) - msg += "tOF %02.2f, " %(list_dict["tOF"][-1]) - + list_dict["tOF"].append(OF_diff.mean()) + msg += "tOF %02.2f, " % (list_dict["tOF"][-1]) + pre_out_grey = output_grey pre_tar_grey = target_grey target_img, ofy, ofx = crop_8x8(target_img) output_img, ofy, ofx = crop_8x8(output_img) - - if "PSNR" in keys:# psnr - list_dict["PSNR"].append( psnr(target_img, output_img) ) - msg +="psnr %02.2f" %(list_dict["PSNR"][-1]) - - if "SSIM" in keys:# ssim - list_dict["SSIM"].append( ssim(target_img, output_img) ) - msg +=", ssim %02.2f" %(list_dict["SSIM"][-1]) - + + if "PSNR" in keys: # psnr + list_dict["PSNR"].append(psnr(target_img, output_img)) + msg += "psnr %02.2f" % (list_dict["PSNR"][-1]) + + if "SSIM" in keys: # ssim + list_dict["SSIM"].append(ssim(target_img, output_img)) + msg += ", ssim %02.2f" % (list_dict["SSIM"][-1]) + if "LPIPS" in keys or "tLP100" in keys: - img0 = util.im2tensor(target_img) # RGB image from [-1,1] + img0 = util.im2tensor(target_img) # RGB image from [-1,1] img1 = util.im2tensor(output_img) - - if "LPIPS" in keys: # LPIPS - dist01 = model.forward(img0,img1) - list_dict["LPIPS"].append( dist01[0] ) - msg +=", lpips %02.2f" %(dist01[0]) - - if "tLP100" in keys and (i > cutfr):# tLP, temporal metrics + + if "LPIPS" in keys: # LPIPS + dist01 = model.forward(img0, img1) + list_dict["LPIPS"].append(dist01[0]) + msg += ", lpips %02.2f" % (dist01[0]) + + if "tLP100" in keys and (i > cutfr): # tLP, temporal metrics dist0t = model.forward(pre_img0, img0) dist1t = model.forward(pre_img1, img1) # print ("tardis %f, outdis %f" %(dist0t, dist1t)) - dist01t = np.absolute(dist0t - dist1t) * 100.0 ##########!!!!! - list_dict["tLP100"].append( dist01t[0] ) - msg += ", tLPx100 %02.2f" %(dist01t[0]) + dist01t = np.absolute(dist0t - dist1t) * 100.0 # !!!!! + list_dict["tLP100"].append(dist01t[0]) + msg += ", tLPx100 %02.2f" % (dist01t[0]) pre_img0 = img0 pre_img1 = img1 - - msg +=", crop (%d, %d)" %(ofy, ofx) + + msg += ", crop (%d, %d)" % (ofy, ofx) #print(msg) - mode = 'w' if folder_i==0 else 'a' - + mode = 'w' if folder_i == 0 else 'a' + pd_dict = {} for cur_num_data in keys: - num_data = cur_num_data+"_%02d" % folder_i + num_data = cur_num_data + "_%02d" % folder_i cur_list = np.float32(list_dict[cur_num_data]) pd_dict[num_data] = pd.Series(cur_list) - + num_data_sum = cur_list.sum() num_data_len = cur_list.shape[0] num_data_mean = num_data_sum / num_data_len - #print("%s, max %02.4f, min %02.4f, avg %02.4f" % + #print("%s, max %02.4f, min %02.4f, avg %02.4f" % # (num_data, cur_list.max(), cur_list.min(), num_data_mean)) - + if folder_i == 0: - avg_dict["Avg_"+cur_num_data] = [num_data_mean] + avg_dict["Avg_" + cur_num_data] = [num_data_mean] else: - avg_dict["Avg_"+cur_num_data] += [num_data_mean] - - sum_dict["FrameAvg_"+cur_num_data] += num_data_sum + avg_dict["Avg_" + cur_num_data] += [num_data_mean] + + sum_dict["FrameAvg_" + cur_num_data] += num_data_sum len_dict[cur_num_data] += num_data_len - folder_dict["FolderAvg_"+cur_num_data] += num_data_mean - - pd.DataFrame(pd_dict).to_csv(os.path.join(FLAGS.output,"metrics.csv"), mode=mode) - + folder_dict["FolderAvg_" + cur_num_data] += num_data_mean + + pd.DataFrame(pd_dict).to_csv(os.path.join(FLAGS.output, "metrics.csv"), mode=mode) + for num_data in keys: - sum_dict["FrameAvg_"+num_data] = pd.Series([sum_dict["FrameAvg_"+num_data] / len_dict[num_data]]) - folder_dict["FolderAvg_"+num_data] = pd.Series([folder_dict["FolderAvg_"+num_data] / folder_n]) - avg_dict["Avg_"+num_data] = pd.Series(np.float32(avg_dict["Avg_"+num_data])) - print("%s, total frame %d, total avg %02.4f, folder avg %02.4f" % - (num_data, len_dict[num_data], sum_dict["FrameAvg_"+num_data][0], folder_dict["FolderAvg_"+num_data][0])) -pd.DataFrame(avg_dict).to_csv(os.path.join(FLAGS.output,"metrics.csv"), mode='a') -pd.DataFrame(folder_dict).to_csv(os.path.join(FLAGS.output,"metrics.csv"), mode='a') -pd.DataFrame(sum_dict).to_csv(os.path.join(FLAGS.output,"metrics.csv"), mode='a') + sum_dict["FrameAvg_" + num_data] = pd.Series([sum_dict["FrameAvg_" + num_data] / len_dict[num_data]]) + folder_dict["FolderAvg_" + num_data] = pd.Series([folder_dict["FolderAvg_" + num_data] / folder_n]) + avg_dict["Avg_" + num_data] = pd.Series(np.float32(avg_dict["Avg_" + num_data])) + print("%s, total frame %d, total avg %02.4f, folder avg %02.4f" % + (num_data, len_dict[num_data], sum_dict["FrameAvg_" + num_data][0], folder_dict["FolderAvg_" + num_data][0])) +pd.DataFrame(avg_dict).to_csv(os.path.join(FLAGS.output, "metrics.csv"), mode='a') +pd.DataFrame(folder_dict).to_csv(os.path.join(FLAGS.output, "metrics.csv"), mode='a') +pd.DataFrame(sum_dict).to_csv(os.path.join(FLAGS.output, "metrics.csv"), mode='a') print("Finished.") diff --git a/codes/utils/data_utils.py b/codes/utils/data_utils.py index aad7197..35300f0 100644 --- a/codes/utils/data_utils.py +++ b/codes/utils/data_utils.py @@ -1,11 +1,10 @@ import os import os.path as osp -from scipy import signal import cv2 import numpy as np import torch -import torch.nn.functional as F +from scipy import signal def create_kernel(sigma, ksize=None): @@ -37,9 +36,9 @@ def rgb_to_ycbcr(img): """ T = np.array([ - [0.256788235294118, -0.148223529411765, 0.439215686274510], + [0.256788235294118, -0.148223529411765, 0.439215686274510], [0.504129411764706, -0.290992156862745, -0.367788235294118], - [0.097905882352941, 0.439215686274510, -0.071427450980392], + [0.097905882352941, 0.439215686274510, -0.071427450980392], ], dtype=np.float64) O = np.array([16, 128, 128], dtype=np.float64) diff --git a/codes/utils/net_utils.py b/codes/utils/net_utils.py index b43aa6e..3938d15 100644 --- a/codes/utils/net_utils.py +++ b/codes/utils/net_utils.py @@ -116,15 +116,15 @@ def __init__(self, scale_factor, a=-0.75): # calculate weights (according to Eq.(6) in the reference paper) cubic = torch.FloatTensor([ - [0, a, -2*a, a], + [0, a, -2 * a, a], [1, 0, -(a + 3), a + 2], - [0, -a, (2*a + 3), -(a + 2)], + [0, -a, (2 * a + 3), -(a + 2)], [0, 0, a, -a] ]) kernels = [ torch.matmul(cubic, torch.FloatTensor([1, s, s**2, s**3])) - for s in [1.0*d/scale_factor for d in range(scale_factor)] + for s in [1.0 * d / scale_factor for d in range(scale_factor)] ] # s = x - floor(x) # register parameters @@ -136,7 +136,7 @@ def forward(self, input): f = self.scale_factor # merge n&c - input = input.reshape(n*c, 1, h, w) + input = input.reshape(n * c, 1, h, w) # pad input (left, right, top, bottom) input = F.pad(input, (1, 2, 1, 2), mode='replicate') @@ -144,15 +144,14 @@ def forward(self, input): # calculate output (vertical expansion) kernel_h = self.kernels.view(f, 1, 4, 1) output = F.conv2d(input, kernel_h, stride=1, padding=0) - output = output.permute(0, 2, 1, 3).reshape(n*c, 1, f*h, w + 3) + output = output.permute(0, 2, 1, 3).reshape(n * c, 1, f * h, w + 3) # calculate output (horizontal expansion) kernel_w = self.kernels.view(f, 1, 1, 4) output = F.conv2d(output, kernel_w, stride=1, padding=0) - output = output.permute(0, 2, 3, 1).reshape(n*c, 1, f*h, f*w) + output = output.permute(0, 2, 3, 1).reshape(n * c, 1, f * h, f * w) # split n&c - output = output.reshape(n, c, f*h, f*w) + output = output.reshape(n, c, f * h, f * w) return output - diff --git a/scripts/create_lmdb.py b/scripts/create_lmdb.py index 96120ee..f833ac1 100644 --- a/scripts/create_lmdb.py +++ b/scripts/create_lmdb.py @@ -38,8 +38,8 @@ def create_lmdb(dataset, raw_dir, lmdb_dir, zip_img=False, filter_file=''): else: nbytes_per_frm = get_FileSize(frm_path_lst[0]) - nbytes += len(frm_path_lst)*nbytes_per_frm - alloc_size = round(1.5*nbytes) + nbytes += len(frm_path_lst) * nbytes_per_frm + alloc_size = round(1.5 * nbytes) print(f'{alloc_size / (1 << 30):.2f} GB') # create lmdb environment diff --git a/scripts/monitor_training.py b/scripts/monitor_training.py index 497853c..1731ffa 100644 --- a/scripts/monitor_training.py +++ b/scripts/monitor_training.py @@ -1,9 +1,9 @@ -import os.path as osp import argparse +import bisect import json import math +import os.path as osp import re -import bisect import matplotlib.pyplot as plt @@ -201,4 +201,4 @@ def monitor(root_dir, testset, exp_id_lst, loss_lst, metric_lst): else: raise ValueError(f'Unrecoginzed model: {args.model}') - monitor(root_dir, args.dataset, exp_id_lst, loss_lst, metric_lst) \ No newline at end of file + monitor(root_dir, args.dataset, exp_id_lst, loss_lst, metric_lst) diff --git a/scripts/resize_BD.py b/scripts/resize_BD.py index b39e46a..b714e8d 100644 --- a/scripts/resize_BD.py +++ b/scripts/resize_BD.py @@ -1,12 +1,10 @@ +import glob import os import os.path as osp -import glob from multiprocessing import Pool -import numpy as np import cv2 - # setup params # default settings scale = 4 # downsampling scale @@ -52,4 +50,3 @@ def downsample_worker(filepath): pool.apply_async(downsample_worker, args=(filepath,)) pool.close() pool.join() - From 3a075c7b61db345cce4c5035aa519c7324879fae Mon Sep 17 00:00:00 2001 From: wupeiwei Date: Tue, 3 Aug 2021 21:00:33 +0800 Subject: [PATCH 8/9] feat: add eta time. --- codes/main.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/codes/main.py b/codes/main.py index ea108e0..5fc341f 100644 --- a/codes/main.py +++ b/codes/main.py @@ -1,6 +1,7 @@ import math import os.path as osp import time +from datetime import datetime import torch @@ -35,6 +36,7 @@ def train(opt): base_utils.log_info(f'Total epochs needed: {total_epoch} for {total_iter} iterations') # train + start_time = time.time() for epoch in range(total_epoch): if opt['dist']: train_loader.sampler.set_epoch(epoch) @@ -61,8 +63,12 @@ def train(opt): # print messages if log_freq > 0 and curr_iter % log_freq == 0: + ckpt_time = time.time() - start_time + eta = int(ckpt_time * (total_iter - curr_iter) / curr_iter) + eta = datetime.fromtimestamp(eta) - datetime.fromtimestamp(0) + msg = model.get_format_msg(epoch, curr_iter, total_epoch, total_iter) - base_utils.log_info(msg) + base_utils.log_info(f"{msg} | eta: {str(eta)}") # save model if ckpt_freq > 0 and curr_iter % ckpt_freq == 0: From 2960d79f2f77bde86c32ee2268a34348373bbab8 Mon Sep 17 00:00:00 2001 From: wupeiwei Date: Wed, 4 Aug 2021 09:25:21 +0800 Subject: [PATCH 9/9] fix: monitor with new log. --- scripts/monitor_training.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/scripts/monitor_training.py b/scripts/monitor_training.py index 1731ffa..cfc33a2 100644 --- a/scripts/monitor_training.py +++ b/scripts/monitor_training.py @@ -22,7 +22,7 @@ def split(pattern, string): def parse_log(log_file): # define loss patterns - loss_pattern = r'.*\[epoch:.*\| iter: (\d+).*\] (.*)' + loss_pattern = r'.*\[epoch:\s+.*\| iter:\s+(\d+).*\] (.*\|)' # load log file with open(log_file, 'r') as f: @@ -36,7 +36,10 @@ def parse_log(log_file): iter = int(loss_match.group(1)) append(loss_dict, 'iter', iter) for s in split(',', loss_match.group(2)): - if s: + if ':' in s: + if '|' in s: + s = str.split(s, '|')[0] + k, v = split(':', s) append(loss_dict, k, float(v))