import torch from torch.utils.data.dataloader import _DataLoaderIter, DataLoader # need torch.__version__ == '1.1.0' # from torch.utils.data import DataLoader # if torch.__version__ == '1.5.1' or torch.__version__ == '1.4.0' use this # from torch.utils.data.dataloader import _MultiProcessingDataLoaderIter as _DataLoaderIter # from torch.utils.data.dataloader import DataLoader import torch.optim as optim import imageio import cv2 import numpy as np import os import torch.nn as nn import torch.nn.functional as F import torch.utils.model_zoo as model_zoo import torch from torch.nn.init import xavier_normal, kaiming_normal from torch.utils.data import Dataset import pickle import argparse import collections import random from shutil import rmtree import time import zipfile import png import array import warnings import shutil class tools(): class abstract_config(): name_filter_out_list = [] # some attributes should not appear in the file name. Write its name here def get_name(self, print_now=True): temp = dir(self) name_filter_out_list = self.name_filter_out_list + ['name_filter_out_list', 'get_name', 'update', 'update_ex_name', 'get_dict', 'check_length_of_file_path', 'check_length_of_file_name'] temp = list(filter(lambda x: True if x.find('__') < 0 and x not in name_filter_out_list else False, temp)) temp = sorted(temp) if print_now: norm_length = 50 print('=' * 10) print('{') for i in temp: temp_str = "'%s'" % i if len(temp_str) < norm_length: temp_str += ' ' * (norm_length - len(temp_str)) try: temp_str += ": '%s,%s', " % (getattr(self, i), '') except: continue # print((i, getattr(self, i))) print('\t' + temp_str) print('}') print('=' * 10) name = '' for i in temp: name += '%s|%s_' % (i, getattr(self, i)) return name @classmethod def check_length_of_file_name(cls, file_name): if len(file_name) >= 255: return False else: return True @classmethod def check_length_of_file_path(cls, filepath): if len(filepath) >= 4096: return False else: return True def update(self, data: dict): def dict_class(obj): temp = {} for name in dir(obj): value = getattr(obj, name) if not name.startswith('_'): temp[name] = value return temp s_dict = dict_class(self) t_key = list(data.keys()) for i in s_dict.keys(): if i in t_key: setattr(self, i, data[i]) print('set param ==== %s: %s' % (i, data[i])) def get_dict(self): def dict_class(obj): temp = {} for name in dir(obj): value = getattr(obj, name) if not name.startswith('_'): temp[name] = value return temp s_dict = dict_class(self) return s_dict def update_ex_name(self, ex_name: str): return ex_name class abstract_model(nn.Module): def save_model(self, save_path): torch.save(self.state_dict(), save_path) def load_model(self, load_path, if_relax=False, if_print=True): if if_print: print('loading protrained model from %s' % load_path) if if_relax: model_dict = self.state_dict() pretrained_dict = torch.load(load_path) # pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict} pretrained_dict_v2 = {} for k, v in pretrained_dict.items(): if k in model_dict: if v.shape == model_dict[k].shape: pretrained_dict_v2[k] = v model_dict.update(pretrained_dict_v2) self.load_state_dict(model_dict) else: self.load_state_dict(torch.load(load_path)) @classmethod def choose_gpu(cls, model, gpu_opt=None): # choose gpu if gpu_opt is None: # gpu=0 model = model.cuda() # torch.cuda.set_device(gpu) # model.cuda(gpu) # model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[gpu]) # print('torch.cuda.device_count() ',torch.cuda.device_count()) # model=torch.nn.parallel.DistributedDataParallel(model,device_ids=range(torch.cuda.device_count())) model = torch.nn.DataParallel(model, device_ids=list(range(torch.cuda.device_count()))) # multi gpu elif gpu_opt == 0: model = model.cuda() else: if type(gpu_opt) != int: raise ValueError('wrong gpu config, it show be int: %s' % (str(gpu_opt))) torch.cuda.set_device(gpu_opt) model = model.cuda(gpu_opt) return model @classmethod def save_model_gpu(cls, model, path): name_dataparallel = torch.nn.DataParallel.__name__ if type(model).__name__ == name_dataparallel: model = model.module model.save_model(path) class abs_test_model(): save_dir = '' def eval_forward(self, im1, im2, gt, *args): # do model forward and cache forward results return 0 def eval_save_result(self, save_name, predflow, *args, **kwargs): pass class data_prefetcher(): def __init__(self, dataset, gpu_opt=None, batch_size=1, shuffle=False, num_workers=0, pin_memory=False, drop_last=False): self.dataset = dataset loader = DataLoader(dataset=self.dataset, batch_size=batch_size, num_workers=num_workers, shuffle=shuffle, drop_last=drop_last, pin_memory=pin_memory) # self.loader = iter(loader) self.loader = _DataLoaderIter(loader) self.stream = torch.cuda.Stream() self.gpu_opt = gpu_opt self.batch_size = batch_size self.shuffle = shuffle self.num_workers = num_workers self.pin_memory = pin_memory self.drop_last = drop_last def build(self): loader = DataLoader(dataset=self.dataset, batch_size=self.batch_size, num_workers=self.num_workers, shuffle=self.shuffle, drop_last=self.drop_last, pin_memory=self.pin_memory) self.loader = _DataLoaderIter(loader) # self.loader = iter(loader) def next(self): try: # batch = next(self.loader) batch = self.loader.next() except StopIteration: self.build() return None # print('self.batch',type(self.batch)) # for i in range(len(self.batch)): # print('i',i,type(self.batch[i])) with torch.cuda.stream(self.stream): batch = tools.tensor_gpu(*batch, check_on=True, non_blocking=True, gpu_opt=self.gpu_opt) # self.next_img = self.next_img.cuda(non_blocking=True).float() # self.next_seg = self.next_seg.cuda(non_blocking=True).float() # self.next_weight = self.next_weight.cuda(non_blocking=True) # self.mask2 = self.mask2.cuda(non_blocking=True).float() # self.mask3 = self.mask3.cuda(non_blocking=True).float() # With Amp, it isn't necessary to manually convert data to half. # if args.fp16: # self.next_input = self.next_input.half() # else: # self.next_input = self.next_input.float() # self.next_input = self.next_input.sub_(self.mean).div_(self.std) return batch class DataProvider: def __init__(self, dataset, batch_size, shuffle=True, num_worker=4, drop_last=True, pin_memory=True): self.batch_size = batch_size self.dataset = dataset self.dataiter = None self.iteration = 0 # self.epoch = 0 # self.shuffle = shuffle self.pin_memory = pin_memory self.num_worker = num_worker self.drop_last = drop_last def build(self): dataloader = DataLoader(self.dataset, batch_size=self.batch_size, shuffle=self.shuffle, num_workers=self.num_worker, pin_memory=self.pin_memory, drop_last=self.drop_last) self.dataiter = _DataLoaderIter(dataloader) def next(self): if self.dataiter is None: self.build() try: batch = self.dataiter.next() self.iteration += 1 # if self.is_cuda: # batch = [batch[0].cuda(), batch[1].cuda(), batch[2].cuda()] return batch except StopIteration: # ??epoch???reload self.epoch += 1 self.build() self.iteration = 1 # reset and return the 1st batch batch = self.dataiter.next() # if self.is_cuda: # batch = [batch[0].cuda(), batch[1].cuda(), batch[2].cuda()] return batch # read/write something as npz class npz_saver(): @classmethod def save_npz(cls, files, npz_save_path): np.savez(npz_save_path, files=[files, 0]) @classmethod def load_npz(cls, npz_save_path): with np.load(npz_save_path) as fin: files = fin['files'] files = list(files) return files[0] # read/write something as pkl class pickle_saver(): @classmethod def save_pickle(cls, files, file_path): with open(file_path, 'wb') as data: pickle.dump(files, data) @classmethod def load_picke(cls, file_path): with open(file_path, 'rb') as data: data = pickle.load(data) return data class AverageMeter(): def __init__(self): self.reset() def reset(self): self.val = 0 self.avg = 0 self.sum = 0 self.count = 0 def update(self, val, num): self.val = val self.sum += val * num self.count += num self.avg = self.sum / self.count class Avg_meter_ls(): def __init__(self): self.data_ls = {} self.short_name_ls = {} def update(self, name, val, num, short_name=None): if name not in self.data_ls.keys(): self.data_ls[name] = tools.AverageMeter() if short_name is None: short_name = name self.short_name_ls[name] = short_name self.data_ls[name].update(val=val, num=num) def print_loss(self, name): a = ' %s %.4f(%.4f)' % (self.short_name_ls[name], self.data_ls[name].val, self.data_ls[name].avg) return a def print_all_losses(self): a = '' for i in sorted(self.data_ls.keys()): a += ' %s %.4f(%.4f)' % (self.short_name_ls[i], self.data_ls[i].val, self.data_ls[i].avg) return a def reset(self): for name in self.data_ls.keys(): self.data_ls[name].reset() # tik tok class time_clock(): def __init__(self): self.st = 0 self.en = 0 self.start_flag = False def start(self): self.reset() self.start_flag = True self.st = time.time() def reset(self): self.start_flag = False self.st = 0 self.en = 0 def end(self): self.en = time.time() def get_during(self): return self.en - self.st # boundary dilated warping class boundary_dilated_warp(): @classmethod def get_grid(cls, batch_size, H, W, start): xx = torch.arange(0, W).view(1, -1).repeat(H, 1) yy = torch.arange(0, H).view(-1, 1).repeat(1, W) xx = xx.view(1, 1, H, W).repeat(batch_size, 1, 1, 1) yy = yy.view(1, 1, H, W).repeat(batch_size, 1, 1, 1) ones = torch.ones_like(xx) grid = torch.cat((xx, yy, ones), 1).float() if torch.cuda.is_available(): grid = grid.cuda() # print("grid",grid.shape) # print("start", start) grid[:, :2, :, :] = grid[:, :2, :, :] + start # 加上patch在原图内的偏移量 return grid @classmethod def transformer(cls, I, vgrid, train=True): # I: Img, shape: batch_size, 1, full_h, full_w # vgrid: vgrid, target->source, shape: batch_size, 2, patch_h, patch_w # outsize: (patch_h, patch_w) def _repeat(x, n_repeats): rep = torch.ones([n_repeats, ]).unsqueeze(0) rep = rep.int() x = x.int() x = torch.matmul(x.reshape([-1, 1]), rep) return x.reshape([-1]) def _interpolate(im, x, y, out_size, scale_h): # x: x_grid_flat # y: y_grid_flat # out_size: same as im.size # scale_h: True if normalized # constants num_batch, num_channels, height, width = im.size() out_height, out_width = out_size[0], out_size[1] # zero = torch.zeros_like([],dtype='int32') zero = 0 max_y = height - 1 max_x = width - 1 if scale_h: # scale indices from [-1, 1] to [0, width or height] # print('--Inter- scale_h:', scale_h) x = (x + 1.0) * (height) / 2.0 y = (y + 1.0) * (width) / 2.0 # do sampling x0 = torch.floor(x).int() x1 = x0 + 1 y0 = torch.floor(y).int() y1 = y0 + 1 x0 = torch.clamp(x0, zero, max_x) # same as np.clip x1 = torch.clamp(x1, zero, max_x) y0 = torch.clamp(y0, zero, max_y) y1 = torch.clamp(y1, zero, max_y) dim1 = torch.from_numpy(np.array(width * height)) dim2 = torch.from_numpy(np.array(width)) base = _repeat(torch.arange(0, num_batch) * dim1, out_height * out_width) # 其实就是单纯标出batch中每个图的下标位置 # base = torch.arange(0,num_batch) * dim1 # base = base.reshape(-1, 1).repeat(1, out_height * out_width).reshape(-1).int() # 区别？expand不对数据进行拷贝 .reshape(-1,1).expand(-1,out_height * out_width).reshape(-1) if torch.cuda.is_available(): dim2 = dim2.cuda() dim1 = dim1.cuda() y0 = y0.cuda() y1 = y1.cuda() x0 = x0.cuda() x1 = x1.cuda() base = base.cuda() base_y0 = base + y0 * dim2 base_y1 = base + y1 * dim2 idx_a = base_y0 + x0 idx_b = base_y1 + x0 idx_c = base_y0 + x1 idx_d = base_y1 + x1 # use indices to lookup pixels in the flat image and restore # channels dim im = im.permute(0, 2, 3, 1) im_flat = im.reshape([-1, num_channels]).float() idx_a = idx_a.unsqueeze(-1).long() idx_a = idx_a.expand(out_height * out_width * num_batch, num_channels) Ia = torch.gather(im_flat, 0, idx_a) idx_b = idx_b.unsqueeze(-1).long() idx_b = idx_b.expand(out_height * out_width * num_batch, num_channels) Ib = torch.gather(im_flat, 0, idx_b) idx_c = idx_c.unsqueeze(-1).long() idx_c = idx_c.expand(out_height * out_width * num_batch, num_channels) Ic = torch.gather(im_flat, 0, idx_c) idx_d = idx_d.unsqueeze(-1).long() idx_d = idx_d.expand(out_height * out_width * num_batch, num_channels) Id = torch.gather(im_flat, 0, idx_d) # and finally calculate interpolated values x0_f = x0.float() x1_f = x1.float() y0_f = y0.float() y1_f = y1.float() wa = torch.unsqueeze(((x1_f - x) * (y1_f - y)), 1) wb = torch.unsqueeze(((x1_f - x) * (y - y0_f)), 1) wc = torch.unsqueeze(((x - x0_f) * (y1_f - y)), 1) wd = torch.unsqueeze(((x - x0_f) * (y - y0_f)), 1) output = wa * Ia + wb * Ib + wc * Ic + wd * Id return output def _transform(I, vgrid, scale_h): C_img = I.shape[1] B, C, H, W = vgrid.size() x_s_flat = vgrid[:, 0, ...].reshape([-1]) y_s_flat = vgrid[:, 1, ...].reshape([-1]) out_size = vgrid.shape[2:] input_transformed = _interpolate(I, x_s_flat, y_s_flat, out_size, scale_h) output = input_transformed.reshape([B, H, W, C_img]) return output # scale_h = True output = _transform(I, vgrid, scale_h=False) if train: output = output.permute(0, 3, 1, 2) return output @classmethod def warp_im(cls, I_nchw, flow_nchw, start_n211): batch_size, _, img_h, img_w = I_nchw.size() _, _, patch_size_h, patch_size_w = flow_nchw.size() patch_indices = cls.get_grid(batch_size, patch_size_h, patch_size_w, start_n211) vgrid = patch_indices[:, :2, ...] # grid_warp = vgrid - flow_nchw grid_warp = vgrid + flow_nchw pred_I2 = cls.transformer(I_nchw, grid_warp) return pred_I2 class occ_check_model(): def __init__(self, occ_type='for_back_check', occ_alpha_1=1.0, occ_alpha_2=0.05, sum_abs_or_squar=True, obj_out_all='all'): ''' :param occ_type: method to check occ mask: bidirection check, or froward warping check(not implemented) :param occ_alpha_1: threshold :param occ_alpha_2: threshold :param obj_out_all: occ mask for: (1) all occ area; (2) only moving object occ area; (3) only out-plane occ area. ''' self.occ_type_ls = ['for_back_check', 'forward_warp'] assert occ_type in self.occ_type_ls assert obj_out_all in ['obj', 'out', 'all'] self.occ_type = occ_type self.occ_alpha_1 = occ_alpha_1 self.occ_alpha_2 = occ_alpha_2 self.sum_abs_or_squar = True # found that false is not OK self.obj_out_all = obj_out_all def __call__(self, flow_f, flow_b, scale=1): # 输入进来是可使用的光流 if self.obj_out_all == 'all': if self.occ_type == 'for_back_check': occ_1, occ_2 = self._forward_backward_occ_check(flow_fw=flow_f, flow_bw=flow_b, scale=scale) elif self.occ_type == 'forward_warp': raise ValueError('not implemented') else: raise ValueError('occ type should be in %s, get %s' % (self.occ_type_ls, self.occ_type)) return occ_1, occ_2 elif self.obj_out_all == 'obj': if self.occ_type == 'for_back_check': occ_1, occ_2 = self._forward_backward_occ_check(flow_fw=flow_f, flow_bw=flow_b, scale=scale) elif self.occ_type == 'forward_warp': raise ValueError('not implemented') elif self.occ_type == 'for_back_check&forward_warp': raise ValueError('not implemented') else: raise ValueError('occ type should be in %s, get %s' % (self.occ_type_ls, self.occ_type)) out_occ_fw = self.torch_outgoing_occ_check(flow_f) out_occ_bw = self.torch_outgoing_occ_check(flow_b) obj_occ_fw = self.torch_get_obj_occ_check(occ_mask=occ_1, out_occ=out_occ_fw) obj_occ_bw = self.torch_get_obj_occ_check(occ_mask=occ_2, out_occ=out_occ_bw) return obj_occ_fw, obj_occ_bw elif self.obj_out_all == 'out': out_occ_fw = self.torch_outgoing_occ_check(flow_f) out_occ_bw = self.torch_outgoing_occ_check(flow_b) return out_occ_fw, out_occ_bw else: raise ValueError("obj_out_all should be in ['obj','out','all'], but get: %s" % self.obj_out_all) def _forward_backward_occ_check(self, flow_fw, flow_bw, scale=1): """ In this function, the parameter alpha needs to be improved """ def length_sq_v0(x): # torch.sum(x ** 2, dim=1, keepdim=True) # temp = torch.sum(x ** 2, dim=1, keepdim=True) # temp = torch.pow(temp, 0.5) return torch.sum(torch.pow(x ** 2, 0.5), dim=1, keepdim=True) # return temp def length_sq(x): # torch.sum(x ** 2, dim=1, keepdim=True) temp = torch.sum(x ** 2, dim=1, keepdim=True) temp = torch.pow(temp, 0.5) # return torch.sum(torch.pow(x ** 2, 0.5), dim=1, keepdim=True) return temp if self.sum_abs_or_squar: sum_func = length_sq_v0 else: sum_func = length_sq mag_sq = sum_func(flow_fw) + sum_func(flow_bw) flow_bw_warped = tools.torch_warp(flow_bw, flow_fw) # torch_warp(img,flow) flow_fw_warped = tools.torch_warp(flow_fw, flow_bw) flow_diff_fw = flow_fw + flow_bw_warped flow_diff_bw = flow_bw + flow_fw_warped occ_thresh = self.occ_alpha_1 * mag_sq + self.occ_alpha_2 / scale occ_fw = sum_func(flow_diff_fw) < occ_thresh # 0 means the occlusion region where the photo loss we should ignore occ_bw = sum_func(flow_diff_bw) < occ_thresh # if IF_DEBUG: # temp_ = sum_func(flow_diff_fw) # tools.check_tensor(data=temp_, name='check occlusion mask sum_func flow_diff_fw') # temp_ = sum_func(flow_diff_bw) # tools.check_tensor(data=temp_, name='check occlusion mask sum_func flow_diff_bw') # tools.check_tensor(data=mag_sq, name='check occlusion mask mag_sq') # tools.check_tensor(data=occ_thresh, name='check occlusion mask occ_thresh') return occ_fw.float(), occ_bw.float() def forward_backward_occ_check(self, flow_fw, flow_bw, alpha1, alpha2, obj_out_all='obj'): """ In this function, the parameter alpha needs to be improved """ def length_sq_v0(x): # torch.sum(x ** 2, dim=1, keepdim=True) # temp = torch.sum(x ** 2, dim=1, keepdim=True) # temp = torch.pow(temp, 0.5) return torch.sum(torch.pow(x ** 2, 0.5), dim=1, keepdim=True) # return temp def length_sq(x): # torch.sum(x ** 2, dim=1, keepdim=True) temp = torch.sum(x ** 2, dim=1, keepdim=True) temp = torch.pow(temp, 0.5) # return torch.sum(torch.pow(x ** 2, 0.5), dim=1, keepdim=True) return temp if self.sum_abs_or_squar: sum_func = length_sq_v0 else: sum_func = length_sq mag_sq = sum_func(flow_fw) + sum_func(flow_bw) flow_bw_warped = tools.torch_warp(flow_bw, flow_fw) # torch_warp(img,flow) flow_fw_warped = tools.torch_warp(flow_fw, flow_bw) flow_diff_fw = flow_fw + flow_bw_warped flow_diff_bw = flow_bw + flow_fw_warped occ_thresh = alpha1 * mag_sq + alpha2 occ_fw = sum_func(flow_diff_fw) < occ_thresh # 0 means the occlusion region where the photo loss we should ignore occ_bw = sum_func(flow_diff_bw) < occ_thresh occ_fw = occ_fw.float() occ_bw = occ_bw.float() # if IF_DEBUG: # temp_ = sum_func(flow_diff_fw) # tools.check_tensor(data=temp_, name='check occlusion mask sum_func flow_diff_fw') # temp_ = sum_func(flow_diff_bw) # tools.check_tensor(data=temp_, name='check occlusion mask sum_func flow_diff_bw') # tools.check_tensor(data=mag_sq, name='check occlusion mask mag_sq') # tools.check_tensor(data=occ_thresh, name='check occlusion mask occ_thresh') if obj_out_all == 'obj': out_occ_fw = self.torch_outgoing_occ_check(flow_fw) out_occ_bw = self.torch_outgoing_occ_check(flow_bw) occ_fw = self.torch_get_obj_occ_check(occ_mask=occ_fw, out_occ=out_occ_fw) occ_bw = self.torch_get_obj_occ_check(occ_mask=occ_bw, out_occ=out_occ_bw) return occ_fw, occ_bw def _forward_warp_occ_check(self, flow_bw): # TODO return 0 @classmethod def torch_outgoing_occ_check(cls, flow): B, C, H, W = flow.size() # mesh grid xx = torch.arange(0, W).view(1, -1).repeat(H, 1) yy = torch.arange(0, H).view(-1, 1).repeat(1, W) xx = xx.view(1, 1, H, W).repeat(B, 1, 1, 1).float() yy = yy.view(1, 1, H, W).repeat(B, 1, 1, 1).float() flow_x, flow_y = torch.split(flow, 1, 1) if flow.is_cuda: xx = xx.cuda() yy = yy.cuda() # tools.check_tensor(flow_x, 'flow_x') # tools.check_tensor(flow_y, 'flow_y') # tools.check_tensor(xx, 'xx') # tools.check_tensor(yy, 'yy') pos_x = xx + flow_x pos_y = yy + flow_y # tools.check_tensor(pos_x, 'pos_x') # tools.check_tensor(pos_y, 'pos_y') # print(' ') # check mask outgoing_mask = torch.ones_like(pos_x) outgoing_mask[pos_x > W - 1] = 0 outgoing_mask[pos_x < 0] = 0 outgoing_mask[pos_y > H - 1] = 0 outgoing_mask[pos_y < 0] = 0 return outgoing_mask.float() @classmethod def torch_get_obj_occ_check(cls, occ_mask, out_occ): outgoing_mask = torch.zeros_like(occ_mask) if occ_mask.is_cuda: outgoing_mask = outgoing_mask.cuda() outgoing_mask[occ_mask == 1] = 1 outgoing_mask[out_occ == 0] = 1 return outgoing_mask class lib_to_show_flow(): @classmethod def point_vec(cls, img, flow, valid=None): meshgrid = np.meshgrid(range(img.shape[1]), range(img.shape[0])) dispimg = cv2.resize(img, None, fx=4, fy=4) colorflow = tools.flow_to_image(flow).astype(int) if valid is None: valid = np.ones((img.shape[0], img.shape[1]), dtype=flow.dtype) for i in range(img.shape[1]): # x for j in range(img.shape[0]): # y # if flow[j, i, 2] != 1: continue if valid[j, i] != 1: continue if j % 10 != 0 or i % 10 != 0: continue xend = int((meshgrid[0][j, i] + flow[j, i, 0]) * 4) yend = int((meshgrid[1][j, i] + flow[j, i, 1]) * 4) leng = np.linalg.norm(flow[j, i, :2]) if leng < 1: continue dispimg = cv2.arrowedLine(dispimg, (meshgrid[0][j, i] * 4, meshgrid[1][j, i] * 4), (xend, yend), (int(colorflow[j, i, 0]), int(colorflow[j, i, 1]), int(colorflow[j, i, 2])), 3, tipLength=8 / leng, line_type=cv2.LINE_AA) return dispimg @classmethod def flow_error_image_np(cls, flow_pred, flow_gt, mask_occ, mask_noc=None, log_colors=True): """Visualize the error between two flows as 3-channel color image. Adapted from the KITTI C++ devkit. Args: flow_pred: prediction flow of shape [ height, width, 2]. flow_gt: ground truth mask_occ: flow validity mask of shape [num_batch, height, width, 1]. Equals 1 at (occluded and non-occluded) valid pixels. mask_noc: Is 1 only at valid pixels which are not occluded. """ # mask_noc = tf.ones(tf.shape(mask_occ)) if mask_noc is None else mask_noc mask_noc = np.ones(mask_occ.shape) if mask_noc is None else mask_noc diff_sq = (flow_pred - flow_gt) ** 2 # diff = tf.sqrt(tf.reduce_sum(diff_sq, [3], keep_dims=True)) diff = np.sqrt(np.sum(diff_sq, axis=2, keepdims=True)) if log_colors: height, width, _ = flow_pred.shape # num_batch, height, width, _ = tf.unstack(tf.shape(flow_1)) colormap = [ [0, 0.0625, 49, 54, 149], [0.0625, 0.125, 69, 117, 180], [0.125, 0.25, 116, 173, 209], [0.25, 0.5, 171, 217, 233], [0.5, 1, 224, 243, 248], [1, 2, 254, 224, 144], [2, 4, 253, 174, 97], [4, 8, 244, 109, 67], [8, 16, 215, 48, 39], [16, 1000000000.0, 165, 0, 38]] colormap = np.asarray(colormap, dtype=np.float32) colormap[:, 2:5] = colormap[:, 2:5] / 255 # mag = tf.sqrt(tf.reduce_sum(tf.square(flow_2), 3, keep_dims=True)) tempp = np.square(flow_gt) # temp = np.sum(tempp, axis=2, keep_dims=True) # mag = np.sqrt(temp) mag = np.sqrt(np.sum(tempp, axis=2, keepdims=True)) # error = tf.minimum(diff / 3, 20 * diff / mag) error = np.minimum(diff / 3, 20 * diff / (mag + 1e-7)) im = np.zeros([height, width, 3]) for i in range(colormap.shape[0]): colors = colormap[i, :] cond = np.logical_and(np.greater_equal(error, colors[0]), np.less(error, colors[1])) # temp=np.tile(cond, [1, 1, 3]) im = np.where(np.tile(cond, [1, 1, 3]), np.ones([height, width, 1]) * colors[2:5], im) # temp=np.cast(mask_noc, np.bool) # im = np.where(np.tile(np.cast(mask_noc, np.bool), [1, 1, 3]), im, im * 0.5) im = np.where(np.tile(mask_noc == 1, [1, 1, 3]), im, im * 0.5) im = im * mask_occ else: error = (np.minimum(diff, 5) / 5) * mask_occ im_r = error # errors in occluded areas will be red im_g = error * mask_noc im_b = error * mask_noc im = np.concatenate([im_r, im_g, im_b], axis=2) # im = np.concatenate(axis=2, values=[im_r, im_g, im_b]) return im[:, :, ::-1] class SP_transform(): @classmethod def denormalize_coords(cls, xx, yy, width, height): """ scale indices from [-1, 1] to [0, width/height] """ xx = 0.5 * (width - 1.0) * (xx.float() + 1.0) yy = 0.5 * (height - 1.0) * (yy.float() + 1.0) return xx, yy @classmethod def normalize_coords(cls, xx, yy, width, height): """ scale indices from [0, width/height] to [-1, 1] """ xx = (2.0 / (width - 1.0)) * xx.float() - 1.0 yy = (2.0 / (height - 1.0)) * yy.float() - 1.0 return xx, yy @classmethod def apply_transform_to_params(cls, theta0, theta_transform): a1 = theta0[:, 0] a2 = theta0[:, 1] a3 = theta0[:, 2] a4 = theta0[:, 3] a5 = theta0[:, 4] a6 = theta0[:, 5] # b1 = theta_transform[:, 0] b2 = theta_transform[:, 1] b3 = theta_transform[:, 2] b4 = theta_transform[:, 3] b5 = theta_transform[:, 4] b6 = theta_transform[:, 5] # c1 = a1 * b1 + a4 * b2 c2 = a2 * b1 + a5 * b2 c3 = b3 + a3 * b1 + a6 * b2 c4 = a1 * b4 + a4 * b5 c5 = a2 * b4 + a5 * b5 c6 = b6 + a3 * b4 + a6 * b5 # new_theta = torch.stack([c1, c2, c3, c4, c5, c6], dim=1) return new_theta class _IdentityParams(nn.Module): def __init__(self): super(tools.SP_transform._IdentityParams, self).__init__() self._batch_size = 0 self.register_buffer("_o", torch.FloatTensor()) self.register_buffer("_i", torch.FloatTensor()) def _update(self, batch_size): torch.zeros([batch_size, 1], out=self._o) torch.ones([batch_size, 1], out=self._i) return torch.cat([self._i, self._o, self._o, self._o, self._i, self._o], dim=1) def forward(self, batch_size): if self._batch_size != batch_size: self._identity_params = self._update(batch_size) self._batch_size = batch_size return self._identity_params class RandomMirror(nn.Module): def __init__(self, vertical=True, p=0.5): super(tools.SP_transform.RandomMirror, self).__init__() self._batch_size = 0 self._p = p self._vertical = vertical self.register_buffer("_mirror_probs", torch.FloatTensor()) def update_probs(self, batch_size): torch.ones([batch_size, 1], out=self._mirror_probs) self._mirror_probs *= self._p def forward(self, theta_list): batch_size = theta_list[0].size(0) if batch_size != self._batch_size: self.update_probs(batch_size) self._batch_size = batch_size # apply random sign to a1 a2 a3 (these are the guys responsible for x) sign = torch.sign(2.0 * torch.bernoulli(self._mirror_probs) - 1.0) i = torch.ones_like(sign) horizontal_mirror = torch.cat([sign, sign, sign, i, i, i], dim=1) theta_list = [theta * horizontal_mirror for theta in theta_list] # apply random sign to a4 a5 a6 (these are the guys responsible for y) if self._vertical: sign = torch.sign(2.0 * torch.bernoulli(self._mirror_probs) - 1.0) vertical_mirror = torch.cat([i, i, i, sign, sign, sign], dim=1) theta_list = [theta * vertical_mirror for theta in theta_list] return theta_list class RandomAffineFlow(nn.Module): def __init__(self, cfg, addnoise=True): super(tools.SP_transform.RandomAffineFlow, self).__init__() self.cfg = cfg self._interp2 = tools.Interpolation.Interp2(clamp=False) self._flow_interp2 = tools.Interpolation.Interp2(clamp=False) self._meshgrid = tools.Interpolation.Meshgrid() self._identity = tools.SP_transform._IdentityParams() self._random_mirror = tools.SP_transform.RandomMirror(cfg.vflip) if cfg.hflip else tools.SP_transform.RandomMirror(p=1) self._addnoise = addnoise self.register_buffer("_noise1", torch.FloatTensor()) self.register_buffer("_noise2", torch.FloatTensor()) self.register_buffer("_xbounds", torch.FloatTensor([-1, -1, 1, 1])) self.register_buffer("_ybounds", torch.FloatTensor([-1, 1, -1, 1])) self.register_buffer("_x", torch.IntTensor(1)) self.register_buffer("_y", torch.IntTensor(1)) def inverse_transform_coords(self, width, height, thetas, offset_x=None, offset_y=None): xx, yy = self._meshgrid(width=width, height=height) xx = torch.unsqueeze(xx, dim=0).float() yy = torch.unsqueeze(yy, dim=0).float() if offset_x is not None: xx = xx + offset_x if offset_y is not None: yy = yy + offset_y a1 = thetas[:, 0].contiguous().view(-1, 1, 1) a2 = thetas[:, 1].contiguous().view(-1, 1, 1) a3 = thetas[:, 2].contiguous().view(-1, 1, 1) a4 = thetas[:, 3].contiguous().view(-1, 1, 1) a5 = thetas[:, 4].contiguous().view(-1, 1, 1) a6 = thetas[:, 5].contiguous().view(-1, 1, 1) xx, yy = tools.SP_transform.normalize_coords(xx, yy, width=width, height=height) xq = a1 * xx + a2 * yy + a3 yq = a4 * xx + a5 * yy + a6 xq, yq = tools.SP_transform.denormalize_coords(xq, yq, width=width, height=height) return xq, yq def transform_coords(self, width, height, thetas): xx1, yy1 = self._meshgrid(width=width, height=height) xx, yy = tools.SP_transform.normalize_coords(xx1, yy1, width=width, height=height) def _unsqueeze12(u): return torch.unsqueeze(torch.unsqueeze(u, dim=1), dim=1) a1 = _unsqueeze12(thetas[:, 0]) a2 = _unsqueeze12(thetas[:, 1]) a3 = _unsqueeze12(thetas[:, 2]) a4 = _unsqueeze12(thetas[:, 3]) a5 = _unsqueeze12(thetas[:, 4]) a6 = _unsqueeze12(thetas[:, 5]) # z = a1 * a5 - a2 * a4 b1 = a5 / z b2 = - a2 / z b4 = - a4 / z b5 = a1 / z # xhat = xx - a3 yhat = yy - a6 xq = b1 * xhat + b2 * yhat yq = b4 * xhat + b5 * yhat xq, yq = tools.SP_transform.denormalize_coords(xq, yq, width=width, height=height) return xq, yq def find_invalid(self, width, height, thetas): x = self._xbounds y = self._ybounds # a1 = torch.unsqueeze(thetas[:, 0], dim=1) a2 = torch.unsqueeze(thetas[:, 1], dim=1) a3 = torch.unsqueeze(thetas[:, 2], dim=1) a4 = torch.unsqueeze(thetas[:, 3], dim=1) a5 = torch.unsqueeze(thetas[:, 4], dim=1) a6 = torch.unsqueeze(thetas[:, 5], dim=1) # z = a1 * a5 - a2 * a4 b1 = a5 / z b2 = - a2 / z b4 = - a4 / z b5 = a1 / z # xhat = x - a3 yhat = y - a6 xq = b1 * xhat + b2 * yhat yq = b4 * xhat + b5 * yhat xq, yq = tools.SP_transform.denormalize_coords(xq, yq, width=width, height=height) # invalid = ( (xq < 0) | (yq < 0) | (xq >= width) | (yq >= height) ).sum(dim=1, keepdim=True) > 0 return invalid def apply_random_transforms_to_params(self, theta0, max_translate, min_zoom, max_zoom, min_squeeze, max_squeeze, min_rotate, max_rotate, validate_size=None): max_translate *= 0.5 batch_size = theta0.size(0) height, width = validate_size # collect valid params here thetas = torch.zeros_like(theta0) zoom = theta0.new(batch_size, 1).zero_() squeeze = torch.zeros_like(zoom) tx = torch.zeros_like(zoom) ty = torch.zeros_like(zoom) phi = torch.zeros_like(zoom) invalid = torch.ones_like(zoom).byte() while invalid.sum() > 0: # random sampling zoom.uniform_(min_zoom, max_zoom) squeeze.uniform_(min_squeeze, max_squeeze) tx.uniform_(-max_translate, max_translate) ty.uniform_(-max_translate, max_translate) phi.uniform_(-min_rotate, max_rotate) # construct affine parameters sx = zoom * squeeze sy = zoom / squeeze sin_phi = torch.sin(phi) cos_phi = torch.cos(phi) b1 = cos_phi * sx b2 = sin_phi * sy b3 = tx b4 = - sin_phi * sx b5 = cos_phi * sy b6 = ty theta_transform = torch.cat([b1, b2, b3, b4, b5, b6], dim=1) theta_try = tools.SP_transform.apply_transform_to_params(theta0, theta_transform) thetas = invalid.float() * theta_try + (1 - invalid).float() * thetas # compute new invalid ones invalid = self.find_invalid(width=width, height=height, thetas=thetas) # here we should have good thetas within borders return thetas def transform_image(self, images, thetas): batch_size, channels, height, width = images.size() xq, yq = self.transform_coords(width=width, height=height, thetas=thetas) transformed = self._interp2(images, xq, yq) return transformed def transform_flow(self, flow, theta1, theta2): batch_size, channels, height, width = flow.size() u = flow[:, 0, :, :] v = flow[:, 1, :, :] # inverse transform coords x0, y0 = self.inverse_transform_coords( width=width, height=height, thetas=theta1) x1, y1 = self.inverse_transform_coords( width=width, height=height, thetas=theta2, offset_x=u, offset_y=v) # subtract and create new flow u = x1 - x0 v = y1 - y0 new_flow = torch.stack([u, v], dim=1) # transform coords xq, yq = self.transform_coords(width=width, height=height, thetas=theta1) # interp2 transformed = self._flow_interp2(new_flow, xq, yq) return transformed def forward(self, data): # 01234 flow 12 21 23 32 imgs = data['imgs'] flows_f = data['flows_f'] masks_f = data['masks_f'] batch_size, _, height, width = imgs[0].size() # identity = no transform theta0 = self._identity(batch_size) # global transform theta_list = [self.apply_random_transforms_to_params( theta0, max_translate=self.cfg.trans[0], min_zoom=self.cfg.zoom[0], max_zoom=self.cfg.zoom[1], min_squeeze=self.cfg.squeeze[0], max_squeeze=self.cfg.squeeze[1], min_rotate=self.cfg.rotate[0], max_rotate=self.cfg.rotate[1], validate_size=[height, width]) ] # relative transform for i in range(len(imgs) - 1): theta_list.append( self.apply_random_transforms_to_params( theta_list[-1], max_translate=self.cfg.trans[1], min_zoom=self.cfg.zoom[2], max_zoom=self.cfg.zoom[3], min_squeeze=self.cfg.squeeze[2], max_squeeze=self.cfg.squeeze[3], min_rotate=-self.cfg.rotate[2], max_rotate=self.cfg.rotate[2], validate_size=[height, width]) ) # random flip images theta_list = self._random_mirror(theta_list) # 01234 imgs = [self.transform_image(im, theta) for im, theta in zip(imgs, theta_list)] if len(imgs) > 2: theta_list = theta_list[1:-1] # 12 23 flows_f = [self.transform_flow(flo, theta1, theta2) for flo, theta1, theta2 in zip(flows_f, theta_list[:-1], theta_list[1:])] masks_f = [self.transform_image(mask, theta) for mask, theta in zip(masks_f, theta_list)] if self._addnoise: ''' im1 (3, 320, 1152) max 0.5885537 min -0.4305366 mean -0.040912468 im1 (3, 320, 1152) max 0.5885537 min -0.4305366 mean -0.03847942 im1 (3, 320, 1152) max 0.5885537 min -0.4187718 mean -0.011021424 ''' stddev = np.random.uniform(0.0, 0.04) for im in imgs: noise = torch.zeros_like(im) noise.normal_(std=stddev) im.add_(noise) im.clamp_(-1.0, 1.0) data['imgs'] = imgs data['flows_f'] = flows_f data['masks_f'] = masks_f return data @classmethod def demo(cls): import pickle import cv2 im0 = cv2.imread("/data/luokunming/Optical_Flow_all/projects/Forward-Warp-master/test/im0.png")[np.newaxis, :, :, :] im1 = cv2.imread("/data/luokunming/Optical_Flow_all/projects/Forward-Warp-master/test/im1.png")[np.newaxis, :, :, :] mask = np.ones((1, 1, im1.shape[1], im1.shape[2])) with open("/data/luokunming/Optical_Flow_all/projects/Forward-Warp-master/test/flow.pkl", "rb+") as f: flow = pickle.load(f) im0 = torch.FloatTensor(im0).permute(0, 3, 1, 2) im1 = torch.FloatTensor(im1).permute(0, 3, 1, 2) mask = torch.FloatTensor(mask) # .permute(0, 3, 1, 2) flow = torch.FloatTensor(flow) flow = flow.permute(0, 3, 1, 2) tools.check_tensor(im0, 'im0') tools.check_tensor(im1, 'im1') tools.check_tensor(flow, 'flow') tools.check_tensor(mask, 'mask') class config(): def __init__(self): self.add_noise = False self.hflip = True self.rotate = [-0.01, 0.01, -0.01, 0.01] self.squeeze = [1.0, 1.0, 1.0, 1.0] self.trans = [0.04, 0.005] self.vflip = True self.zoom = [1.0, 1.4, 0.99, 1.01] model = tools.SP_transform.RandomAffineFlow(config(), addnoise=False) input = {'imgs': [im0 / 255, im1 / 255], 'flows_f': [flow], 'masks_f': [mask]} data = model(input) imgs0, imgs1 = data['imgs'] flows_f = data['flows_f'][0] # show tools.check_tensor(imgs0, 'imgs0 out') def process_image(tens, ind=0): tens_, = tools.tensor_gpu(tens, check_on=False) im = tens_[ind, :, :, :] im = np.transpose(im, (1, 2, 0)) return im im0_ori = tools.Show_GIF.im_norm(process_image(im0)) im0_res = tools.Show_GIF.im_norm(process_image(imgs0)) flow_ori = tools.flow_to_image(process_image(flow)) flow_res = tools.flow_to_image(process_image(flows_f)) tools.cv2_show_dict(im0_ori=im0_ori, im0_res=im0_res, flow_ori=flow_ori, flow_res=flow_res) class txt_read_write(): @classmethod def read(cls, path): with open(path, "r") as f: data = f.readlines() return data @classmethod def write(cls, path, data_ls): file_write_obj = open(path, 'a') for i in data_ls: file_write_obj.writelines(i) file_write_obj.close() @classmethod def demo(cls): txt_path = r'E:\research\unsupervised_optical_flow\projects\Ric-master\Ric-master\data\MPI-Sintel\frame_0001_match.txt' data = tools.txt_read_write.read(txt_path) print(' ') write_txt_path = txt_path = r'E:\research\unsupervised_optical_flow\projects\Ric-master\Ric-master\data\MPI-Sintel\temp.txt' tools.txt_read_write.write(write_txt_path, data[:10]) print(' ') class KITTI_png_flow_read_write(): def demo(self): def write_kitti_png_file(flow_fn, flow_data, mask_data=None): ''' :param flow_fn: save file path, .png file :param flow_data: [H,W,2] :param mask_data: can be occlusion mask (0-1 mask) or None :return: ''' flow_img = np.zeros((flow_data.shape[0], flow_data.shape[1], 3), dtype=np.uint16) if mask_data is None: mask_data = np.ones([flow_data.shape[0], flow_data.shape[1]], dtype=np.uint16) flow_img[:, :, 2] = (flow_data[:, :, 0] * 64.0 + 2 ** 15).astype(np.uint16) flow_img[:, :, 1] = (flow_data[:, :, 1] * 64.0 + 2 ** 15).astype(np.uint16) flow_img[:, :, 0] = mask_data[:, :] cv2.imwrite(flow_fn, flow_img) def read_png_flow(fpath): """ Read KITTI optical flow, returns u,v,valid mask """ R = png.Reader(fpath) width, height, data, _ = R.asDirect() # This only worked with python2. # I = np.array(map(lambda x:x,data)).reshape((height,width,3)) gt = np.array([x for x in data]).reshape((height, width, 3)) flow = gt[:, :, 0:2] flow = (flow.astype('float64') - 2 ** 15) / 64.0 flow = flow.astype(np.float) mask = gt[:, :, 2:3] mask = np.uint8(mask) flow = np.transpose(flow, [2, 0, 1]) mask = np.transpose(mask, [2, 0, 1]) return flow, mask @classmethod def MSE(cls, img1, img2): img1gray = cv2.cvtColor(img1, cv2.COLOR_BGR2GRAY) img2gray = cv2.cvtColor(img2, cv2.COLOR_BGR2GRAY) # cv2.imwrite("L1_.jpg",img1gray-img2gray) rows, cols = img1gray.shape[:2] loss = 0.0 pixel_nums = 0 for row in range(30, rows - 30): for col in range(60, cols - 60): if img1gray[row][col] == 0 or img2gray[row][col] == 0: continue else: pixel_nums += 1 loss += np.square(np.abs(img1gray[row][col] - img2gray[row][col])) loss /= pixel_nums return loss @classmethod def torch_warp_mask(cls, x, flo): """ warp an image/tensor (im2) back to im1, according to the optical flow x: [B, C, H, W] (im2) flo: [B, 2, H, W] flow """ B, C, H, W = x.size() # mesh grid xx = torch.arange(0, W).view(1, -1).repeat(H, 1) yy = torch.arange(0, H).view(-1, 1).repeat(1, W) xx = xx.view(1, 1, H, W).repeat(B, 1, 1, 1) yy = yy.view(1, 1, H, W).repeat(B, 1, 1, 1) grid = torch.cat((xx, yy), 1).float() if x.is_cuda: grid = grid.cuda() # print(grid.shape,flo.shape,'...') vgrid = grid + flo # scale grid to [-1,1] vgrid[:, 0, :, :] = 2.0 * vgrid[:, 0, :, :] / max(W - 1, 1) - 1.0 vgrid[:, 1, :, :] = 2.0 * vgrid[:, 1, :, :] / max(H - 1, 1) - 1.0 vgrid = vgrid.permute(0, 2, 3, 1) # B H,W,C output = nn.functional.grid_sample(x, vgrid, padding_mode='zeros') mask = torch.autograd.Variable(torch.ones(x.size())) if x.is_cuda: mask = mask.cuda() mask = nn.functional.grid_sample(mask, vgrid, padding_mode='zeros') mask[mask < 0.9999] = 0 mask[mask > 0] = 1 output = output * mask # # nchw->>>nhwc # if x.is_cuda: # output = output.cpu() # output_im = output.numpy() # output_im = np.transpose(output_im, (0, 2, 3, 1)) # output_im = np.squeeze(output_im) return output, mask @classmethod def torch_warp(cls, x, flo): """ warp an image/tensor (im2) back to im1, according to the optical flow x: [B, C, H, W] (im2) flo: [B, 2, H, W] flow """ B, C, H, W = x.size() # mesh grid xx = torch.arange(0, W).view(1, -1).repeat(H, 1) yy = torch.arange(0, H).view(-1, 1).repeat(1, W) xx = xx.view(1, 1, H, W).repeat(B, 1, 1, 1) yy = yy.view(1, 1, H, W).repeat(B, 1, 1, 1) grid = torch.cat((xx, yy), 1).float() if x.is_cuda: grid = grid.cuda() # print(grid.shape,flo.shape,'...') vgrid = grid + flo # scale grid to [-1,1] vgrid[:, 0, :, :] = 2.0 * vgrid[:, 0, :, :] / max(W - 1, 1) - 1.0 vgrid[:, 1, :, :] = 2.0 * vgrid[:, 1, :, :] / max(H - 1, 1) - 1.0 vgrid = vgrid.permute(0, 2, 3, 1) # B H,W,C # tools.check_tensor(x, 'x') # tools.check_tensor(vgrid, 'vgrid') output = nn.functional.grid_sample(x, vgrid, padding_mode='zeros') # mask = torch.autograd.Variable(torch.ones(x.size())) # if x.is_cuda: # mask = mask.cuda() # mask = nn.functional.grid_sample(mask, vgrid, padding_mode='zeros') # # mask[mask < 0.9999] = 0 # mask[mask > 0] = 1 # output = output * mask # # nchw->>>nhwc # if x.is_cuda: # output = output.cpu() # output_im = output.numpy() # output_im = np.transpose(output_im, (0, 2, 3, 1)) # output_im = np.squeeze(output_im) return output @classmethod def weights_init(cls, m): classname = m.__class__.__name__ if classname.find('conv') != -1: # torch.nn.init.xavier_normal(m.weight) torch.nn.init.kaiming_normal(m.weight) torch.nn.init.constant(m.bias, 0) @classmethod def warp_cv2(cls, img_prev, flow): # calculate mat w = int(img_prev.shape[1]) h = int(img_prev.shape[0]) y_coords, x_coords = np.mgrid[0:h, 0:w] coords = np.float32(np.dstack([x_coords, y_coords])) pixel_map = coords + flow new_frame = cv2.remap(img_prev, pixel_map, None, cv2.INTER_LINEAR) return new_frame @classmethod def flow_to_image(cls, flow, display=False): """ :param flow: H,W,2 :param display: :return: H,W,3 """ def compute_color(u, v): def make_color_wheel(): """ Generate color wheel according Middlebury color code :return: Color wheel """ RY = 15 YG = 6 GC = 4 CB = 11 BM = 13 MR = 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 """ compute optical flow color map :param u: optical flow horizontal map :param v: optical flow vertical map :return: optical flow in color code """ [h, w] = u.shape img = np.zeros([h, w, 3]) nanIdx = np.isnan(u) | np.isnan(v) u[nanIdx] = 0 v[nanIdx] = 0 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(0, 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 UNKNOWN_FLOW_THRESH = 1e7 """ Convert flow into middlebury color code image :param flow: optical flow map :return: optical flow image in middlebury color """ u = flow[:, :, 0] v = flow[:, :, 1] maxu = -999. maxv = -999. minu = 999. minv = 999. idxUnknow = (abs(u) > UNKNOWN_FLOW_THRESH) | (abs(v) > UNKNOWN_FLOW_THRESH) u[idxUnknow] = 0 v[idxUnknow] = 0 maxu = max(maxu, np.max(u)) minu = min(minu, np.min(u)) maxv = max(maxv, np.max(v)) minv = min(minv, np.min(v)) rad = np.sqrt(u ** 2 + v ** 2) maxrad = max(-1, np.max(rad)) if display: print("max flow: %.4f\nflow range:\nu = %.3f .. %.3f\nv = %.3f .. %.3f" % (maxrad, minu, maxu, minv, maxv)) u = u / (maxrad + np.finfo(float).eps) v = v / (maxrad + np.finfo(float).eps) img = compute_color(u, v) idx = np.repeat(idxUnknow[:, :, np.newaxis], 3, axis=2) img[idx] = 0 return np.uint8(img) @classmethod def write_flow_png(cls, filename, uv, v=None, mask=None): if v is None: assert (uv.ndim == 3) assert (uv.shape[2] == 2) u = uv[:, :, 0] v = uv[:, :, 1] else: u = uv assert (u.shape == v.shape) height_img, width_img = u.shape if mask is None: valid_mask = np.ones([height_img, width_img], dtype=np.uint16) else: valid_mask = mask flow_u = np.clip((u * 64 + 2 ** 15), 0.0, 65535.0).astype(np.uint16) flow_v = np.clip((v * 64 + 2 ** 15), 0.0, 65535.0).astype(np.uint16) output = np.stack((flow_u, flow_v, valid_mask), axis=-1) with open(filename, 'wb') as f: # writer = png.Writer(width=width_img, height=height_img, bitdepth=16) # temp = np.reshape(output, (-1, width_img * 3)) # writer.write(f, temp) png_writer = png.Writer(width=width_img, height=height_img, bitdepth=16, compression=3, greyscale=False) # png_writer.write_array(f, output) temp = np.reshape(output, (-1, width_img * 3)) png_writer.write(f, temp) @classmethod def write_kitti_png_file(cls, flow_fn, flow_data, mask_data=None): flow_img = np.zeros((flow_data.shape[0], flow_data.shape[1], 3), dtype=np.uint16) if mask_data is None: mask_data = np.ones([flow_data.shape[0], flow_data.shape[1]], dtype=np.uint16) flow_img[:, :, 2] = (flow_data[:, :, 0] * 64.0 + 2 ** 15).astype(np.uint16) flow_img[:, :, 1] = (flow_data[:, :, 1] * 64.0 + 2 ** 15).astype(np.uint16) flow_img[:, :, 0] = mask_data[:, :] cv2.imwrite(flow_fn, flow_img) @classmethod def WriteKittiPngFile(cls, path, uv, mask=None): """ Write 16-bit .PNG file as specified by KITTI-2015 (flow). u, v are lists of float values mask is a list of floats, denoting the *valid* pixels. """ assert (uv.ndim == 3) assert (uv.shape[2] == 2) u = uv[:, :, 0] v = uv[:, :, 1] height, width = u.shape if mask is None: valid_mask = np.ones([height, width]) else: valid_mask = mask data = array.array('H', [0]) * width * height * 3 # for i, (u_, v_, mask_) in enumerate(zip(u, v, mask)): data[0] = int(u * 64.0 + 2 ** 15) data[1] = int(v * 64.0 + 2 ** 15) data[2] = int(valid_mask) # if mask_ > 0: # print(data[3*i], data[3*i+1],data[3*i+2]) with open(path, 'wb') as png_file: png_writer = png.Writer(width=width, height=height, bitdepth=16, compression=3, greyscale=False) png_writer.write_array(png_file, data) @classmethod def write_flow(cls, flow, filename): """ write optical flow in Middlebury .flo format :param flow: optical flow map :param filename: optical flow file path to be saved :return: None """ f = open(filename, 'wb') magic = np.array([202021.25], dtype=np.float32) (height, width) = flow.shape[0:2] w = np.array([width], dtype=np.int32) h = np.array([height], dtype=np.int32) magic.tofile(f) w.tofile(f) h.tofile(f) flow.tofile(f) f.close() @classmethod def read_flow(cls, filename): """ read optical flow from Middlebury .flo file :param filename: name of the flow file :return: optical flow data in matrix """ f = open(filename, 'rb') try: magic = np.fromfile(f, np.float32, count=1)[0] # For Python3.x except: magic = np.fromfile(f, np.float32, count=1) # For Python2.x data2d = None if 202021.25 != magic: print('Magic number incorrect. Invalid .flo file') else: w = np.fromfile(f, np.int32, count=1) h = np.fromfile(f, np.int32, count=1) # print("Reading %d x %d flo file" % (h, w)) data2d = np.fromfile(f, np.float32, count=2 * w * h) # reshape data into 3D array (columns, rows, channels) data2d = np.resize(data2d, (h[0], w[0], 2)) f.close() return data2d @classmethod def read_flo(cls, filename): with open(filename, 'rb') as f: magic = np.fromfile(f, np.float32, count=1) if 202021.25 != magic: print('Magic number incorrect. Invalid .flo file') else: w = np.fromfile(f, np.int32, count=1) h = np.fromfile(f, np.int32, count=1) data = np.fromfile(f, np.float32, count=int(2 * w * h)) # Reshape data into 3D array (columns, rows, bands) data2D = np.resize(data, (h[0], w[0], 2)) return data2D @classmethod def write_flo(cls, flow, filename): """ write optical flow in Middlebury .flo format :param flow: optical flow map :param filename: optical flow file path to be saved :return: None """ f = open(filename, 'wb') magic = np.array([202021.25], dtype=np.float32) (height, width) = flow.shape[0:2] w = np.array([width], dtype=np.int32) h = np.array([height], dtype=np.int32) magic.tofile(f) w.tofile(f) h.tofile(f) flow.tofile(f) f.close() @classmethod def tensor_gpu(cls, *args, check_on=True, gpu_opt=None, non_blocking=True): def check_on_gpu(tensor_): if type(gpu_opt) == int: tensor_g = tensor_.cuda(gpu_opt, non_blocking=non_blocking) else: tensor_g = tensor_.cuda() return tensor_g def check_off_gpu(tensor_): if tensor_.is_cuda: tensor_c = tensor_.cpu() else: tensor_c = tensor_ tensor_c = tensor_c.detach().numpy() # tensor_c = cv2.normalize(tensor_c.detach().numpy(), None, 0, 255, cv2.NORM_MINMAX, cv2.CV_8U) return tensor_c if torch.cuda.is_available(): if check_on: data_ls = [check_on_gpu(a) for a in args] else: data_ls = [check_off_gpu(a) for a in args] else: if check_on: data_ls = args else: # data_ls = args data_ls = [a.detach().numpy() for a in args] # data_ls = [cv2.normalize(a.detach().numpy(), None, 0, 255, cv2.NORM_MINMAX, cv2.CV_8U) for a in args] # data_ls = args return data_ls @classmethod def tryremove(cls, name, file=False): try: if file: os.remove(name) else: rmtree(name) except OSError: pass @classmethod def check_tensor(cls, data, name, print_data=False, print_in_txt=None): if data.is_cuda: temp = data.detach().cpu().numpy() else: temp = data.detach().numpy() a = len(name) name_ = name + ' ' * 100 name_ = name_[0:max(a, 10)] print_str = '%s, %s, %s, %s,%s,%s,%s,%s' % (name_, temp.shape, data.dtype, ' max:%.2f' % np.max(temp), ' min:%.2f' % np.min(temp), ' mean:%.2f' % np.mean(temp), ' sum:%.2f' % np.sum(temp), data.device) if print_in_txt is None: print(print_str) else: print(print_str, file=print_in_txt) if print_data: print(temp) return print_str @classmethod def extract_zip(cls, zip_path, extract_dir): print('unzip file: %s' % zip_path) with zipfile.ZipFile(zip_path, "r") as zip_ref: zip_ref.extractall(extract_dir) @classmethod def clear(cls): os.system("clear") @classmethod def random_flag(cls, threshold_0_1=0.5): a = random.random() < threshold_0_1 return a @classmethod def compute_model_size(cls, model, *args): from thop import profile flops, params = profile(model, inputs=args, verbose=False) print('flops: %.3f G, params: %.3f M' % (flops / 1000 / 1000 / 1000, params / 1000 / 1000)) @classmethod def im_norm(cls, img): eps = 1e-5 a = np.max(img) b = np.min(img) img = (img - b) / (a - b) img = img * 255 img = img.astype('uint8') return img @classmethod def check_dir(cls, path): if not os.path.exists(path): os.makedirs(path)