utils/logger.py

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)