fffiloni/Video-Matting-Anything
53
1import os2import cv23import torch4import logging5import datetime6import numpy as np7from pprint import pprint8from utils import util9from utils.config import CONFIG10from tensorboardX import SummaryWriter11 12 13LEVELS = {14 "DEBUG": logging.DEBUG,15 "INFO": logging.INFO,16 "WARNING": logging.WARNING,17 "ERROR": logging.ERROR,18 "CRITICAL": logging.CRITICAL,19}20 21 22def make_color_wheel():23 # from https://github.com/JiahuiYu/generative_inpainting/blob/master/inpaint_ops.py24 RY, YG, GC, CB, BM, MR = (15, 6, 4, 11, 13, 6)25 ncols = RY + YG + GC + CB + BM + MR26 colorwheel = np.zeros([ncols, 3])27 col = 028 # RY29 colorwheel[0:RY, 0] = 25530 colorwheel[0:RY, 1] = np.transpose(np.floor(255*np.arange(0, RY) / RY))31 col += RY32 # YG33 colorwheel[col:col+YG, 0] = 255 - np.transpose(np.floor(255*np.arange(0, YG) / YG))34 colorwheel[col:col+YG, 1] = 25535 col += YG36 # GC37 colorwheel[col:col+GC, 1] = 25538 colorwheel[col:col+GC, 2] = np.transpose(np.floor(255*np.arange(0, GC) / GC))39 col += GC40 # CB41 colorwheel[col:col+CB, 1] = 255 - np.transpose(np.floor(255*np.arange(0, CB) / CB))42 colorwheel[col:col+CB, 2] = 25543 col += CB44 # BM45 colorwheel[col:col+BM, 2] = 25546 colorwheel[col:col+BM, 0] = np.transpose(np.floor(255*np.arange(0, BM) / BM))47 col += + BM48 # MR49 colorwheel[col:col+MR, 2] = 255 - np.transpose(np.floor(255 * np.arange(0, MR) / MR))50 colorwheel[col:col+MR, 0] = 25551 return colorwheel52 53 54COLORWHEEL = make_color_wheel()55 56 57def compute_color(u,v):58 # from https://github.com/JiahuiYu/generative_inpainting/blob/master/inpaint_ops.py59 h, w = u.shape60 img = np.zeros([h, w, 3])61 nanIdx = np.isnan(u) | np.isnan(v)62 u[nanIdx] = 063 v[nanIdx] = 064 colorwheel = COLORWHEEL65 # colorwheel = make_color_wheel()66 ncols = np.size(colorwheel, 0)67 rad = np.sqrt(u**2+v**2)68 a = np.arctan2(-v, -u) / np.pi69 fk = (a+1) / 2 * (ncols - 1) + 170 k0 = np.floor(fk).astype(int)71 k1 = k0 + 172 k1[k1 == ncols+1] = 173 f = fk - k074 for i in range(np.size(colorwheel,1)):75 tmp = colorwheel[:, i]76 col0 = tmp[k0-1] / 25577 col1 = tmp[k1-1] / 25578 col = (1-f) * col0 + f * col179 idx = rad <= 180 col[idx] = 1-rad[idx]*(1-col[idx])81 notidx = np.logical_not(idx)82 col[notidx] *= 0.7583 img[:, :, i] = np.uint8(np.floor(255 * col*(1-nanIdx)))84 return img85 86def flow_to_image(flow):87 # part from https://github.com/JiahuiYu/generative_inpainting/blob/master/inpaint_ops.py88 maxrad = -189 u = flow[0, :, :]90 v = flow[1, :, :]91 rad = np.sqrt(u ** 2 + v ** 2)92 maxrad = max(maxrad, np.max(rad))93 u = u/(maxrad + np.finfo(float).eps)94 v = v/(maxrad + np.finfo(float).eps)95 img = compute_color(u, v)96 97 return img98 99 100def put_text(image, text, position=(10, 20)):101 image = cv2.resize(image.transpose([1, 2, 0]), (512, 512), interpolation=cv2.INTER_NEAREST)102 return cv2.putText(image, text, position, cv2.FONT_HERSHEY_SIMPLEX, 0.8, 0, thickness=2).transpose([2, 0, 1])103 104 105class TensorBoardLogger(object):106 def __init__(self, tb_log_dir, exp_string):107 """108 Initialize summary writer109 """110 self.exp_string = exp_string111 self.tb_log_dir = tb_log_dir112 self.val_img_dir = os.path.join(self.tb_log_dir, 'val_image')113 114 if CONFIG.local_rank == 0:115 util.make_dir(self.tb_log_dir)116 util.make_dir(self.val_img_dir)117 118 self.writer = SummaryWriter(self.tb_log_dir+'/' + self.exp_string)119 else:120 self.writer = None121 122 def scalar_summary(self, tag, value, step, phase='train'):123 if CONFIG.local_rank == 0:124 sum_name = '{}/{}'.format(phase.capitalize(), tag)125 self.writer.add_scalar(sum_name, value, step)126 127 def image_summary(self, image_set, step, phase='train', save_val=True):128 """129 Record image in tensorboard130 The input image should be a numpy array with shape (C, H, W) like a torch tensor131 :param image_set: dict of images132 :param step:133 :param phase:134 :param save_val: save images in folder in validation or testing135 :return:136 """137 if CONFIG.local_rank == 0:138 for tag, image_numpy in image_set.items():139 sum_name = '{}/{}'.format(phase.capitalize(), tag)140 image_numpy = image_numpy.transpose([1, 2, 0])141 142 image_numpy = cv2.resize(image_numpy, (360, 360), interpolation=cv2.INTER_NEAREST)143 144 if len(image_numpy.shape) == 2:145 image_numpy = image_numpy[None, :,:]146 else:147 image_numpy = image_numpy.transpose([2, 0, 1])148 self.writer.add_image(sum_name, image_numpy, step)149 150 if (phase=='test') and save_val:151 tags = list(image_set.keys())152 image_pack = self._reshape_rgb(image_set[tags[0]])153 image_pack = cv2.resize(image_pack, (512, 512), interpolation=cv2.INTER_NEAREST)154 155 for tag in tags[1:]:156 image = self._reshape_rgb(image_set[tag])157 image = cv2.resize(image, (512, 512), interpolation=cv2.INTER_NEAREST)158 image_pack = np.concatenate((image_pack, image), axis=1)159 160 cv2.imwrite(os.path.join(self.val_img_dir, 'val_{:d}'.format(step)+'.png'), image_pack)161 162 @staticmethod163 def _reshape_rgb(image):164 """165 Transform RGB/L -> BGR for OpenCV166 """167 if len(image.shape) == 3 and image.shape[0] == 3:168 image = image.transpose([1, 2, 0])169 image = image[...,::-1]170 elif len(image.shape) == 3 and image.shape[0] == 1:171 image = image.transpose([1, 2, 0])172 image = np.repeat(image, 3, axis=2)173 elif len(image.shape) == 2:174 # image = image.transpose([1,0])175 image = np.stack((image, image, image), axis=2)176 else:177 raise ValueError('Image shape {} not supported to save'.format(image.shape))178 return image179 180 def __del__(self):181 if self.writer is not None:182 self.writer.close()183 184 185class MyLogger(logging.Logger):186 """187 Only write log in the first subprocess188 """189 def __init__(self, *args, **kwargs):190 super(MyLogger, self).__init__(*args, **kwargs)191 192 def _log(self, level, msg, args, exc_info=None, extra=None, stack_info=False):193 if CONFIG.local_rank == 0:194 super()._log(level, msg, args, exc_info, extra, stack_info)195 196 197def get_logger(log_dir=None, tb_log_dir=None, logging_level="DEBUG"):198 """199 Return a default build-in logger if log_file=None and tb_log_dir=None200 Return a build-in logger which dump stdout to log_file if log_file is assigned201 Return a build-in logger and tensorboard summary writer if tb_log_dir is assigned202 :param log_file: logging file dumped from stdout203 :param tb_log_dir: tensorboard dir204 :param logging_level:205 :return: Logger or [Logger, TensorBoardLogger]206 """207 level = LEVELS[logging_level.upper()]208 exp_string = datetime.datetime.now().strftime("%Y-%m-%d-%H-%M-%S")209 210 logging.setLoggerClass(MyLogger)211 logger = logging.getLogger('Logger')212 logger.setLevel(level)213 # create formatter214 formatter = logging.Formatter('[%(asctime)s] %(levelname)s: %(message)s', datefmt='%m-%d %H:%M:%S')215 216 # create console handler217 ch = logging.StreamHandler()218 ch.setLevel(level)219 ch.setFormatter(formatter)220 # add the handlers to logger221 logger.addHandler(ch)222 223 # create file handler224 if log_dir is not None and CONFIG.local_rank == 0:225 log_file = os.path.join(log_dir, exp_string)226 fh = logging.FileHandler(log_file+'.log', mode='w')227 fh.setLevel(level)228 fh.setFormatter(formatter)229 logger.addHandler(fh)230 pprint(CONFIG, stream=fh.stream)231 232 # create tensorboard summary writer233 if tb_log_dir is not None:234 tb_logger = TensorBoardLogger(tb_log_dir=tb_log_dir, exp_string=exp_string)235 return logger, tb_logger236 else:237 return logger238 239 240def normalize_image(image):241 """242 normalize image array to 0~1243 """244 image_flat = torch.flatten(image, start_dim=1)245 return (image - image_flat.min(dim=1, keepdim=False)[0].view(3,1,1)) / (246 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)247