import os import cv2 import torch import logging import datetime import numpy as np from pprint import pprint from utils import util from utils.config import CONFIG from tensorboardX import SummaryWriter LEVELS = { "DEBUG": logging.DEBUG, "INFO": logging.INFO, "WARNING": logging.WARNING, "ERROR": logging.ERROR, "CRITICAL": logging.CRITICAL, } def make_color_wheel(): # from https://github.com/JiahuiYu/generative_inpainting/blob/master/inpaint_ops.py RY, YG, GC, CB, BM, MR = (15, 6, 4, 11, 13, 6) ncols = RY + YG + GC + CB + BM + MR colorwheel = np.zeros([ncols, 3]) col = 0 # RY colorwheel[0:RY, 0] = 255 colorwheel[0:RY, 1] = np.transpose(np.floor(255*np.arange(0, RY) / RY)) col += RY # YG colorwheel[col:col+YG, 0] = 255 - np.transpose(np.floor(255*np.arange(0, YG) / YG)) colorwheel[col:col+YG, 1] = 255 col += YG # GC colorwheel[col:col+GC, 1] = 255 colorwheel[col:col+GC, 2] = np.transpose(np.floor(255*np.arange(0, GC) / GC)) col += GC # CB colorwheel[col:col+CB, 1] = 255 - np.transpose(np.floor(255*np.arange(0, CB) / CB)) colorwheel[col:col+CB, 2] = 255 col += CB # BM colorwheel[col:col+BM, 2] = 255 colorwheel[col:col+BM, 0] = np.transpose(np.floor(255*np.arange(0, BM) / BM)) col += + BM # MR colorwheel[col:col+MR, 2] = 255 - np.transpose(np.floor(255 * np.arange(0, MR) / MR)) colorwheel[col:col+MR, 0] = 255 return colorwheel COLORWHEEL = make_color_wheel() def compute_color(u,v): # from https://github.com/JiahuiYu/generative_inpainting/blob/master/inpaint_ops.py h, w = u.shape img = np.zeros([h, w, 3]) nanIdx = np.isnan(u) | np.isnan(v) u[nanIdx] = 0 v[nanIdx] = 0 colorwheel = COLORWHEEL # colorwheel = make_color_wheel() ncols = np.size(colorwheel, 0) rad = np.sqrt(u**2+v**2) a = np.arctan2(-v, -u) / np.pi fk = (a+1) / 2 * (ncols - 1) + 1 k0 = np.floor(fk).astype(int) k1 = k0 + 1 k1[k1 == ncols+1] = 1 f = fk - k0 for i in range(np.size(colorwheel,1)): tmp = colorwheel[:, i] col0 = tmp[k0-1] / 255 col1 = tmp[k1-1] / 255 col = (1-f) * col0 + f * col1 idx = rad <= 1 col[idx] = 1-rad[idx]*(1-col[idx]) notidx = np.logical_not(idx) col[notidx] *= 0.75 img[:, :, i] = np.uint8(np.floor(255 * col*(1-nanIdx))) return img def flow_to_image(flow): # part from https://github.com/JiahuiYu/generative_inpainting/blob/master/inpaint_ops.py maxrad = -1 u = flow[0, :, :] v = flow[1, :, :] rad = np.sqrt(u ** 2 + v ** 2) maxrad = max(maxrad, np.max(rad)) u = u/(maxrad + np.finfo(float).eps) v = v/(maxrad + np.finfo(float).eps) img = compute_color(u, v) return img def put_text(image, text, position=(10, 20)): image = cv2.resize(image.transpose([1, 2, 0]), (512, 512), interpolation=cv2.INTER_NEAREST) return cv2.putText(image, text, position, cv2.FONT_HERSHEY_SIMPLEX, 0.8, 0, thickness=2).transpose([2, 0, 1]) class TensorBoardLogger(object): def __init__(self, tb_log_dir, exp_string): """ Initialize summary writer """ self.exp_string = exp_string self.tb_log_dir = tb_log_dir self.val_img_dir = os.path.join(self.tb_log_dir, 'val_image') if CONFIG.local_rank == 0: util.make_dir(self.tb_log_dir) util.make_dir(self.val_img_dir) self.writer = SummaryWriter(self.tb_log_dir+'/' + self.exp_string) else: self.writer = None def scalar_summary(self, tag, value, step, phase='train'): if CONFIG.local_rank == 0: sum_name = '{}/{}'.format(phase.capitalize(), tag) self.writer.add_scalar(sum_name, value, step) def image_summary(self, image_set, step, phase='train', save_val=True): """ Record image in tensorboard The input image should be a numpy array with shape (C, H, W) like a torch tensor :param image_set: dict of images :param step: :param phase: :param save_val: save images in folder in validation or testing :return: """ if CONFIG.local_rank == 0: for tag, image_numpy in image_set.items(): sum_name = '{}/{}'.format(phase.capitalize(), tag) image_numpy = image_numpy.transpose([1, 2, 0]) image_numpy = cv2.resize(image_numpy, (360, 360), interpolation=cv2.INTER_NEAREST) if len(image_numpy.shape) == 2: image_numpy = image_numpy[None, :,:] else: image_numpy = image_numpy.transpose([2, 0, 1]) self.writer.add_image(sum_name, image_numpy, step) if (phase=='test') and save_val: tags = list(image_set.keys()) image_pack = self._reshape_rgb(image_set[tags[0]]) image_pack = cv2.resize(image_pack, (512, 512), interpolation=cv2.INTER_NEAREST) for tag in tags[1:]: image = self._reshape_rgb(image_set[tag]) image = cv2.resize(image, (512, 512), interpolation=cv2.INTER_NEAREST) image_pack = np.concatenate((image_pack, image), axis=1) cv2.imwrite(os.path.join(self.val_img_dir, 'val_{:d}'.format(step)+'.png'), image_pack) @staticmethod def _reshape_rgb(image): """ Transform RGB/L -> BGR for OpenCV """ if len(image.shape) == 3 and image.shape[0] == 3: image = image.transpose([1, 2, 0]) image = image[...,::-1] elif len(image.shape) == 3 and image.shape[0] == 1: image = image.transpose([1, 2, 0]) image = np.repeat(image, 3, axis=2) elif len(image.shape) == 2: # image = image.transpose([1,0]) image = np.stack((image, image, image), axis=2) else: raise ValueError('Image shape {} not supported to save'.format(image.shape)) return image def __del__(self): if self.writer is not None: self.writer.close() class MyLogger(logging.Logger): """ Only write log in the first subprocess """ def __init__(self, *args, **kwargs): super(MyLogger, self).__init__(*args, **kwargs) def _log(self, level, msg, args, exc_info=None, extra=None, stack_info=False): if CONFIG.local_rank == 0: super()._log(level, msg, args, exc_info, extra, stack_info) def get_logger(log_dir=None, tb_log_dir=None, logging_level="DEBUG"): """ Return a default build-in logger if log_file=None and tb_log_dir=None Return a build-in logger which dump stdout to log_file if log_file is assigned Return a build-in logger and tensorboard summary writer if tb_log_dir is assigned :param log_file: logging file dumped from stdout :param tb_log_dir: tensorboard dir :param logging_level: :return: Logger or [Logger, TensorBoardLogger] """ level = LEVELS[logging_level.upper()] exp_string = datetime.datetime.now().strftime("%Y-%m-%d-%H-%M-%S") logging.setLoggerClass(MyLogger) logger = logging.getLogger('Logger') logger.setLevel(level) # create formatter formatter = logging.Formatter('[%(asctime)s] %(levelname)s: %(message)s', datefmt='%m-%d %H:%M:%S') # create console handler ch = logging.StreamHandler() ch.setLevel(level) ch.setFormatter(formatter) # add the handlers to logger logger.addHandler(ch) # create file handler if log_dir is not None and CONFIG.local_rank == 0: log_file = os.path.join(log_dir, exp_string) fh = logging.FileHandler(log_file+'.log', mode='w') fh.setLevel(level) fh.setFormatter(formatter) logger.addHandler(fh) pprint(CONFIG, stream=fh.stream) # create tensorboard summary writer if tb_log_dir is not None: tb_logger = TensorBoardLogger(tb_log_dir=tb_log_dir, exp_string=exp_string) return logger, tb_logger else: return logger def normalize_image(image): """ normalize image array to 0~1 """ image_flat = torch.flatten(image, start_dim=1) return (image - image_flat.min(dim=1, keepdim=False)[0].view(3,1,1)) / ( image_flat.max(dim=1, keepdim=False)[0].view(3,1,1) - image_flat.min(dim=1, keepdim=False)[0].view(3,1,1) + 1e-8)