Upload 2 files

core/logger.py

@@ -0,0 +1,118 @@
+import os
+import logging
+from collections import OrderedDict
+import json
+from datetime import datetime
+import argparse
+
+
+def mkdirs(paths):
+    if isinstance(paths, str):
+        os.makedirs(paths, exist_ok=True)
+    else:
+        for path in paths:
+            os.makedirs(path, exist_ok=True)
+
+def get_timestamp():
+    return datetime.now().strftime('%y%m%d_%H%M%S')
+
+
+def parse(args):
+    phase = args.phase
+    opt_path =args.config
+    gpu_ids = args.gpu_ids
+
+    json_str = ''
+    with open(opt_path, 'r') as f:
+        for line in f:
+            line = line.split('//')[0] + '\n'
+            json_str += line
+        #print(json_str)
+    opt =json.loads(json_str, object_pairs_hook=OrderedDict)
+    #print(opt)
+
+    #create experiments folder
+    experiments_root = os.path.join(
+        'experiments', '{}_{}'.format(opt['name'], get_timestamp()))
+    opt['path_cd']['experiments_root'] = experiments_root
+    for key, path in opt['path_cd'].items():
+        if 'resume' not in key and 'experiments' not in key:
+            opt['path_cd'][key] = os.path.join(experiments_root, path)
+            mkdirs(opt['path_cd'][key])
+
+    #chaneg dataset len
+    opt['phase'] = phase
+
+    # export CUDA_VISIBLE_DEVICES
+    if gpu_ids is not None:
+        opt['gpu_ids'] = [int(id) for id in gpu_ids.split(',')]
+        gpu_list = gpu_ids
+    else:
+        gpu_list = ','.join(str(x) for x in opt['gpu_ids'])
+    #print(gpu_list)
+    os.environ['CUDA_VISIBLE_DEVICES'] = gpu_list
+    print('expert CUDA_VISIBLE_DEVICES=' + gpu_list)
+    if len(gpu_list) > 1:
+        opt['distributed'] = True
+    else:
+        opt['distributed'] = False
+
+    return opt
+
+class NoneDict(dict):
+    def __missing__(self, key):
+        return None
+
+# convert to NoneDict, which return None for missing key.
+def dict_to_nonedict(opt):
+    if isinstance(opt, dict):
+        new_opt = dict()
+        for key, sub_opt in opt.items():
+            new_opt[key] = dict_to_nonedict(sub_opt)
+        return NoneDict(**new_opt)
+    elif isinstance(opt, list):
+        return [dict_to_nonedict(sub_opt) for sub_opt in opt]
+    else:
+        return opt
+
+def dict2str(opt, indent_l=1):
+    '''dict to string for logger'''
+    msg = ''
+    for k, v in opt.items():
+        if isinstance(v, dict):
+            msg += ' ' * (indent_l * 2) + k + ':[\n'
+            msg += dict2str(v, indent_l + 1)
+            msg += ' ' * (indent_l * 2) + ']\n'
+        else:
+            msg += ' ' * (indent_l * 2) + k + ': ' + str(v) + '\n'
+    return msg
+
+def setup_logger(logger_name, root, phase, level=logging.INFO, screen=False):
+    '''set up logger'''
+    l = logging.getLogger(logger_name)
+    formatter = logging.Formatter(
+        '%(asctime)s.%(msecs)03d - %(levelname)s: %(message)s', datefmt='%y-%m-%d %H:%M:%S')
+    print(formatter)
+    log_file = os.path.join(root, '{}.log'.format(phase))
+    print(log_file)
+    fh = logging.FileHandler(log_file, mode='w')
+    fh.setFormatter(formatter)
+    l.setLevel(level)
+    l.addHandler(fh)
+    if screen:
+        sh = logging.StreamHandler()
+        sh.setFormatter(formatter)
+        l.addHandler(sh)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument('-c', '--config', type=str, default='../config/levir.json')
+    parser.add_argument('-p', '--phase', type=str, choices=['train', 'test'], default='train')
+    parser.add_argument('-gpu', '--gpu_ids', type=str, default=None)
+
+    args = parser.parse_args()
+    opt = parse(args)
+    print(opt)
+    opt = dict_to_nonedict(opt)
+    print(opt)

core/metrics.py

@@ -0,0 +1,107 @@
+import os
+import math
+import numpy as np
+import cv2
+from torchvision.utils import make_grid
+from data.colormap import second_colormap
+
+
+def tensor2img(tensor, out_type=np.uint8, min_max=(-1, 1)):
+    '''
+    Converts a torch Tensor into an image Numpy array
+    Input: 4D(B,(3/1),H,W), 3D(C,H,W), or 2D(H,W), any range, RGB channel order
+    Output: 3D(H,W,C) or 2D(H,W), [0,255], np.uint8 (default)
+    '''
+    tensor = tensor.squeeze().float().cpu().clamp_(*min_max)  # clamp
+    tensor = (tensor - min_max[0]) / \
+        (min_max[1] - min_max[0])  # to range [0,1]
+    n_dim = tensor.dim()
+    if n_dim == 4:
+        n_img = len(tensor)
+        img_np = make_grid(tensor, nrow=int(
+            math.sqrt(n_img)), normalize=False).numpy()
+        img_np = np.transpose(img_np, (1, 2, 0))  # HWC, RGB
+    elif n_dim == 3:
+        img_np = tensor.numpy()
+        img_np = np.transpose(img_np, (1, 2, 0))  # HWC, RGB
+    elif n_dim == 2:
+        img_np = tensor.numpy()
+    else:
+        raise TypeError(
+            'Only support 4D, 3D and 2D tensor. But received with dimension: {:d}'.format(n_dim))
+    if out_type == np.uint8:
+        img_np = (img_np * 255.0).round()
+        # Important. Unlike matlab, numpy.unit8() WILL NOT round by default.
+    return img_np.astype(out_type)
+
+
+def Index2Color(pred, cmap=second_colormap):
+    colormap = np.asarray(cmap, dtype='uint8')
+    x = np.asarray(pred, dtype='int32')
+    return colormap[x, :]
+
+def save_img(img, img_path, mode='RGB'):
+    cv2.imwrite(img_path, cv2.cvtColor(img, cv2.COLOR_RGB2BGR))
+    # cv2.imwrite(img_path, img)
+
+def save_scdimg(img, img_path, mode='RGB'):
+    cv2.imwrite(img_path, cv2.cvtColor(np.squeeze(img, axis=0), cv2.COLOR_RGB2BGR))
+
+def save_feat(img, img_path, mode='RGB'):
+    cv2.imwrite(img_path, cv2.applyColorMap(cv2.resize(img, (256,256), interpolation=cv2.INTER_CUBIC), cv2.COLORMAP_JET))
+    # cv2.imwrite(img_path, cv2.resize(img, (256,256), interpolation=cv2.INTER_))
+    # cv2.imwrite(img_path, img)
+
+
+def calculate_psnr(img1, img2):
+    # img1 and img2 have range [0, 255]
+    img1 = img1.astype(np.float64)
+    img2 = img2.astype(np.float64)
+    mse = np.mean((img1 - img2)**2)
+    if mse == 0:
+        return float('inf')
+    return 20 * math.log10(255.0 / math.sqrt(mse))
+
+
+def ssim(img1, img2):
+    C1 = (0.01 * 255)**2
+    C2 = (0.03 * 255)**2
+
+    img1 = img1.astype(np.float64)
+    img2 = img2.astype(np.float64)
+    kernel = cv2.getGaussianKernel(11, 1.5)
+    window = np.outer(kernel, kernel.transpose())
+
+    mu1 = cv2.filter2D(img1, -1, window)[5:-5, 5:-5]  # valid
+    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(img1**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(img1 * 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(img1, img2):
+    '''calculate SSIM
+    the same outputs as MATLAB's
+    img1, img2: [0, 255]
+    '''
+    if not img1.shape == img2.shape:
+        raise ValueError('Input images must have the same dimensions.')
+    if img1.ndim == 2:
+        return ssim(img1, img2)
+    elif img1.ndim == 3:
+        if img1.shape[2] == 3:
+            ssims = []
+            for i in range(3):
+                ssims.append(ssim(img1, img2))
+            return np.array(ssims).mean()
+        elif img1.shape[2] == 1:
+            return ssim(np.squeeze(img1), np.squeeze(img2))
+    else:
+        raise ValueError('Wrong input image dimensions.')