IIR-Lab/aligned_utils.py

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