import numpy as np import cv2 # from PIL import Image import os import glob from tqdm import tqdm from pathlib import Path import torch import torch.nn.functional as F # Parameters of the motion estimation algorithms def warp_flow(img, flow): ''' Applies to img the transformation described by flow. ''' #assert len(flow.shape) == 3 and flow.shape[-1] == 2 hf, wf = flow.shape[:2] # flow = -flow flow[:, :, 0] += np.arange(wf) flow[:, :, 1] += np.arange(hf)[:, np.newaxis] res = cv2.remap(img, flow, None, cv2.INTER_LINEAR) return res def estimate_invflow(img0, img1, me_algo): ''' Estimates inverse optical flow by using the me_algo algorithm. ''' # Create estimator object if me_algo == "DeepFlow": of_estim = cv2.optflow.createOptFlow_DeepFlow() else: raise Exception("Incorrect motion estimation algorithm") # Run flow estimation (inverse flow) flow = of_estim.calc(img1, img0, None) # flow = cv.calcOpticalFlowFarneback(prvs,next, None, 0.5, 3, 15, 3, 5, 1.2, 0) return flow def align_frames(img_to_align, img_source, mc_alg='DeepFlow'): ''' Applies to img_to_align a transformation which converts it into img_source. Args: img_to_align: HxWxC image img_source: HxWxC image mc_alg: selects between DeepFlow, SimpleFlow, and TVL1. DeepFlow runs by default. Returns: HxWxC aligned image ''' if img_to_align.ndim == 2: img0 = img_to_align img1 = img_source else: img0 = img_to_align[:, :, 1] img1 = img_source[:, :, 1] out_img = None # Align frames according to selection in mc_alg flow = estimate_invflow(img0, img1, mc_alg) #print(flow.astype(np.float32)) # rectifier out_img = warp_flow(img_to_align, flow.astype(np.float32)) return out_img, flow def SIFT(img1gray, img2gray): # if i == 0: sift = cv2.xfeatures2d.SIFT_create() # 创建sift方法 # sift = cv2.SURF_create() # 创建sift方法 # find the keypoints and descriptors with SIFT kp1, des1 = sift.detectAndCompute(img1gray, None) # 用sift找到图像中的关键点和描述子 kp2, des2 = sift.detectAndCompute(img2gray, None) # FLANN parameters FLANN_INDEX_KDTREE = 1 # FLANN使用的算法选择，有0，1等，具体多少算法不太清楚。 index_params = dict(algorithm=FLANN_INDEX_KDTREE, trees=5) search_params = dict(checks=10) flann = cv2.FlannBasedMatcher(index_params, search_params) # 这里创建FLANN匹配算法 matches = flann.knnMatch(des1, des2, k=2) # 这里是使用创建的FLANN算法对两张图上的描述子进行匹配，使用k近邻匹配，k=2即最近邻匹配 # 上述返回的matches是一种数据类型，这样一个类型中包含了matches.queryIdx .trainIdx 和 .distance，由于是knn，k=2，返回两个最相似的特征点。 # 而上面返回的特征点kp1和kp2也是一种类，包含了kp1.pt：关键点坐标 kp1.angle：关键点方向 kp1.response：关键点强度 kp1.size该点直径大小 # Need to draw only good matches, so create a mask matchesMask = [[0, 0] for i in range(len(matches))] # 为了去画匹配的情况，创建了一个掩膜 good = [] # ratio test as per Lowe's paper for i, (m, n) in enumerate(matches): if m.distance < 0.65*n.distance: good.append(m) matchesMask[i] = [1, 0] MIN_MATCH_COUNT = 9 print(len(good)) if len(good) > MIN_MATCH_COUNT: src_pts = np.float32([kp1[m.queryIdx].pt for m in good]).reshape(-1, 1, 2) dst_pts = np.float32([kp2[m.trainIdx].pt for m in good]).reshape(-1, 1, 2) M, mask = cv2.findHomography(src_pts, dst_pts, cv2.RANSAC, 3) else: print('error!!!!!!!!!!!!!!!!!!!!!!!!!!!') return # print(M) return M def match_colors(im_ref, im_q, im_test): im_ref_mean_re = im_ref.view(*im_ref.shape[:2], -1) im_q_mean_re = im_q.view(*im_q.shape[:2], -1) # Estimate color transformation matrix by minimizing the least squares error c_mat_all = [] for ir, iq in zip(im_ref_mean_re, im_q_mean_re): c = torch.linalg.lstsq(iq.t(), ir.t()) c = c.solution[:im_ref_mean_re.size(1)] c_mat_all.append(c) c_mat = torch.stack(c_mat_all, dim=0) # Apply the transformation to test image im_test_re = im_test.view(*im_test.shape[:2], -1) im_t_conv = torch.matmul(im_test_re.permute(0, 2, 1), c_mat).permute(0, 2, 1) im_t_conv = im_t_conv.view(im_test.shape) return im_t_conv def color_correction(gt, in_put, output, scale_factor=2): # ds_gt = F.interpolate(gt, scale_factor=1.0 / scale_factor, mode='bilinear', align_corners=False, recompute_scale_factor=True) output_cor = match_channel_colors(gt, in_put, output) return output_cor def match_channel_colors(im_ref, im_q, im_test): im_ref_reshape = im_ref.view(*im_ref.shape[:2], -1) im_q_reshape = im_q.view(*im_q.shape[:2], -1) im_test_reshape = im_test.view(*im_test.shape[:2], -1) # Estimate color transformation matrix by minimizing the least squares error im_t_conv_list = [] for i in range(im_ref.size(1)): c_mat_all = [] for ir_batch, iq_batch in zip(im_ref_reshape[:, i:i+1, :], im_q_reshape[:, i:i+1, :]): c = torch.linalg.lstsq(iq_batch.t(), ir_batch.t()) c = c.solution[:1] c_mat_all.append(c) c_mat = torch.stack(c_mat_all, dim=0) # Apply the transformation to test image im_t_conv = torch.matmul(im_test_reshape[:, i:i+1, :].permute(0, 2, 1), c_mat).permute(0, 2, 1) im_t_conv = im_t_conv.view(*im_t_conv.shape[:2], *im_test.shape[-2:]) im_t_conv_list.append(im_t_conv) im_t_conv = torch.cat(im_t_conv_list, dim=1) return im_t_conv def img2tensor(imgs, bgr2rgb=True, float32=True): """Numpy array to tensor. Args: imgs (list[ndarray] | ndarray): Input images. bgr2rgb (bool): Whether to change bgr to rgb. float32 (bool): Whether to change to float32. Returns: list[tensor] | tensor: Tensor images. If returned results only have one element, just return tensor. """ def _totensor(img, bgr2rgb, float32): if img.shape[2] == 3 and bgr2rgb: if img.dtype == 'float64': img = img.astype('float32') img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img = torch.from_numpy(img.transpose(2, 0, 1)) if float32: img = img.float() return img if isinstance(imgs, list): return [_totensor(img, bgr2rgb, float32) for img in imgs] else: return _totensor(imgs, bgr2rgb, float32) def tensor2img(tensor, rgb2bgr=True, out_type=np.uint8, min_max=(0, 1)): """Convert torch Tensors into image numpy arrays. After clamping to [min, max], values will be normalized to [0, 1]. Args: tensor (Tensor or list[Tensor]): Accept shapes: 1) 4D mini-batch Tensor of shape (B x 3/1 x H x W); 2) 3D Tensor of shape (3/1 x H x W); 3) 2D Tensor of shape (H x W). Tensor channel should be in RGB order. rgb2bgr (bool): Whether to change rgb to bgr. out_type (numpy type): output types. If ``np.uint8``, transform outputs to uint8 type with range [0, 255]; otherwise, float type with range [0, 1]. Default: ``np.uint8``. min_max (tuple[int]): min and max values for clamp. Returns: (Tensor or list): 3D ndarray of shape (H x W x C) OR 2D ndarray of shape (H x W). The channel order is BGR. """ if not (torch.is_tensor(tensor) or (isinstance(tensor, list) and all(torch.is_tensor(t) for t in tensor))): raise TypeError(f'tensor or list of tensors expected, got {type(tensor)}') if torch.is_tensor(tensor): tensor = [tensor] result = [] for _tensor in tensor: _tensor = _tensor.squeeze(0).float().detach().cpu().clamp_(*min_max) _tensor = (_tensor - min_max[0]) / (min_max[1] - min_max[0]) n_dim = _tensor.dim() if n_dim == 4: img_np = make_grid(_tensor, nrow=int(math.sqrt(_tensor.size(0))), normalize=False).numpy() img_np = img_np.transpose(1, 2, 0) if rgb2bgr: img_np = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR) elif n_dim == 3: img_np = _tensor.numpy() img_np = img_np.transpose(1, 2, 0) if img_np.shape[2] == 1: # gray image img_np = np.squeeze(img_np, axis=2) else: if rgb2bgr: img_np = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR) elif n_dim == 2: img_np = _tensor.numpy() else: raise TypeError(f'Only support 4D, 3D or 2D tensor. But received with dimension: {n_dim}') if out_type == np.uint8: # Unlike MATLAB, numpy.unit8() WILL NOT round by default. img_np = (img_np * 255.0).round() img_np = img_np.astype(out_type) result.append(img_np) if len(result) == 1: result = result[0] return result