vismatch 1.1.1__py3-none-any.whl

vismatch/third_party/imatch-toolbox/third_party/patch2pix/utils/datasets/dataset_megadepth.py

@@ -0,0 +1,141 @@
+import os
+from PIL import Image
+import numpy as np
+import torch.utils.data as data
+
+from utils.datasets.preprocess import get_tuple_transform_ops
+from utils.eval.measure import sampson_distance
+from utils.eval.geometry import pose2fund
+
+class ImMatchDatasetMega(data.Dataset):
+    '''Data wrapper for train image-matching with triplets'''
+    def __init__(self, data_root, match_file, scene_list=None, wt=480, ht=320, item_type='triplet'):
+        print('\nInitialize ImMatchDatasetMega...')        
+        self.dataset = 'MegaDepth_undistort'
+        self.data_root = os.path.join(data_root, self.dataset)
+        self.match_file = match_file
+        self.transform_ops = get_tuple_transform_ops(resize=(ht, wt), normalize=True)
+        self.wt, self.ht = wt, ht
+        self.item_type = item_type
+        
+        # Initialize data
+        self.ims = {}            # {scene: {im: imsize}}
+        self.pos_pair_pool = []  # [pair]
+        self.load_pairs(scene_list)
+        self.Fs = {}
+        self.Ks = {}
+        
+        
+    def load_im(self, im_ref, crop=None):
+        im = Image.open(im_ref)
+        if crop:
+            dw, dh = crop  
+            im = np.array(im)
+
+            # Crop from right and buttom to keep the target aspect ratio
+            h, w, _ = im.shape
+            im = im[0: h - int(dh), 0: w - int(dw)]
+            #print(h, w, im.shape)
+            im = Image.fromarray(im)
+        return im
+            
+    def load_pairs(self, scene_list=None):        
+        match_dict = np.load(self.match_file, allow_pickle=True).item()
+        self.scenes = scene_list if scene_list else match_dict.keys()
+        print('Loading data from {}'.format(self.match_file))                
+        
+        num_ims = 0
+        for sc in self.scenes:
+            self.pos_pair_pool += match_dict[sc]['pairs']
+            self.ims[sc] = match_dict[sc]['ims']
+            num_ims += len(match_dict[sc]['ims'])    
+        print('Loaded scenes {} ims: {} pos pairs:{}'.format(len(self.scenes), num_ims, len(self.pos_pair_pool)))            
+    
+    def get_fundmat(self, pair, im1, im2):        
+        def scale_intrinsic(K, wi, hi):
+            sx, sy = self.wt / wi, self.ht /  hi
+            sK = np.array([[sx, 0, 0],
+                           [0, sy, 0],
+                           [0, 0, 1]])
+            return sK.dot(K)
+        
+        pair_key = (pair.im1, pair.im2)                
+        if pair_key not in self.Fs:        
+            # Recompute camera intrinsic matrix due to the resize
+            K1 = scale_intrinsic(pair.K1, im1.width, im1.height)
+            K2 = scale_intrinsic(pair.K2, im2.width, im2.height)
+
+            # Calculate F
+            F = pose2fund(K1, K2, pair.R, pair.t)        
+            self.Fs[pair_key] = (F, K1, K2)
+
+            # Sanity check
+            # scale = np.array([[im1.width/self.wt, im1.height/self.ht, im2.width/self.wt, im2.height/self.ht]])
+            # matches = pair.sanity_matches * scale        
+            # dists = sampson_distance(matches[:, :2], matches[:,2:], F)
+            # print(np.mean(dists))
+        return self.Fs[pair_key]
+    
+    def __getitem__(self, index):
+        """
+        Batch dict:
+            - 'src_im': anchor image
+            - 'pos_im': positive image sample to the anchor
+            - 'neg_im': negative image sample to the anchor
+            - 'im_pair_refs': path of images (src, pos, neg)
+            - 'pair_data': namedtuple contains relative pose information between src and pos ims.
+        """
+        data_dict = {}
+        
+        # Load positive pair data
+        pair = self.pos_pair_pool[index]
+        im_src_ref = os.path.join(self.data_root, pair.im1)
+        im_pos_ref = os.path.join(self.data_root, pair.im2)
+        im_src = self.load_im(im_src_ref, crop=pair.crop1)
+        im_pos = self.load_im(im_pos_ref, crop=pair.crop2)
+        
+        # Select a negative image from other scences        
+        if self.item_type == 'triplet':
+            other_scenes = list(self.scenes)
+            other_scenes.remove(pair.im1.split('/')[0])
+            neg_scene = np.random.choice(other_scenes)
+            im_neg_data = np.random.choice(self.ims[neg_scene])
+            im_neg_ref = os.path.join(self.data_root, im_neg_data.name)
+            im_neg = self.load_im(im_neg_ref, crop=im_neg_data.crop)
+            im_neg = self.transform_ops([im_neg])[0]  
+            #print(im_neg.shape)
+        else:
+            im_neg = None
+            im_neg_ref = None
+            
+        
+        # Compute fundamental matrix before RESIZE
+        F, K1, K2 = self.get_fundmat(pair, im_src, im_pos) 
+        
+        # Process images
+        im_src, im_pos = self.transform_ops((im_src, im_pos))
+        #print(im_src.shape, im_pos.shape)
+        
+        # Wrap data item
+        data_dict = {'src_im': im_src, 
+                     'pos_im': im_pos, 
+                     'neg_im': im_neg, 
+                     'im_pair_refs': (im_src_ref, im_pos_ref, im_neg_ref),
+                     'F': F,
+                     'K1': K1,
+                     'K2': K2
+                     }
+
+        return data_dict
+    
+    def __len__(self):
+        return len(self.pos_pair_pool)
+    
+    def __repr__(self):
+        fmt_str = 'ImMatchDatasetMega scenes:{} data type:{}\n'.format(len(self.scenes), self.item_type)
+        fmt_str += 'Number of data pairs: {}\n'.format(self.__len__())
+        fmt_str += 'Image root location: {}\n'.format(self.data_root)
+        fmt_str += 'Match file: {}\n'.format(self.match_file)
+        fmt_str += 'Transforms: {}\n'.format(self.transform_ops.__repr__().replace('\n', '\n    '))
+        return fmt_str  
+

vismatch/third_party/imatch-toolbox/third_party/patch2pix/utils/datasets/preprocess.py

@@ -0,0 +1,184 @@
+import torch
+import random
+from torchvision import transforms
+import numpy as np
+from PIL import Image
+
+def load_im_tensor(im_path, device, imsize=None, with_gray=True):    
+    im = Image.open(im_path)
+
+    # Resize
+    wt, ht = wo, ho = im.width, im.height
+    if imsize and max(wo, ho) > imsize and imsize > 0:
+        scale = imsize / max(wo, ho)
+        ht, wt = int(round(ho * scale)), int(round(wo * scale))
+        im = im.resize((wt, ht), Image.BICUBIC)    
+    scale = (wo / wt, ho / ht)
+    
+    # Gray
+    gray = None
+    if with_gray:
+        im_gray = np.array(im.convert('L'))
+        gray = transforms.functional.to_tensor(im_gray).unsqueeze(0).to(device)
+        
+    # RGB  
+    im = transforms.functional.to_tensor(im)
+    im = transforms.functional.normalize(im , mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+    im = im.unsqueeze(0).to(device)
+    if with_gray:
+        return im, gray, scale
+    return im, scale
+
+def load_im_flexible(im_path, k_size=2, upsample=16, imsize=None, crop_square=False):
+    """Load image for matching on fly.
+    The image can be arbitrary size and will be processed to 
+    fulfil the format required for matching.    
+    """    
+    # Load original image
+    img = Image.open(im_path)
+    img = img.convert('RGB')
+    wo, ho = img.width, img.height
+    if not (imsize and imsize > 0):
+        imsize = max(wo, ho)
+    elif imsize > max(wo, ho):     # Disable upsampling
+        imsize = max(wo, ho)
+        
+    wt, ht = cal_rescale_size(image_size=imsize, w=wo, h=ho, k_size=k_size, 
+                              scale_factor=1./upsample, no_print=True)    
+
+    # Resize image and transform to tensor
+    ops = get_tuple_transform_ops(resize=None, normalize=True)
+    img = transforms.functional.resize(img, (ht, wt), Image.BICUBIC)
+    img = ops([img])[0]
+    
+    # Mainly for beauty plotting
+    if crop_square:    
+        _, h, w = img.shape
+        img = img[:, :w,:]
+        
+    scale = (wo / wt, ho / ht)    
+    return img, scale
+
+def crop_from_bottom_right(w, h, target_ratio=1.5, min_ratio=1.3, max_ratio=1.7):
+    ratio = w / h    
+    if ratio < min_ratio or ratio > max_ratio:
+        return None
+    if ratio == target_ratio:
+        return 0, 0
+    if ratio > target_ratio:
+        # Cut the width
+        dh = h % 2
+        ht = h - dh
+        dw = w - ht * target_ratio
+        wt = w - dw
+
+    if ratio < target_ratio:
+        # Cut the height
+        dw = w % 3
+        wt = w - dw    
+        dh = h - wt / target_ratio
+        ht = h - dh
+    return dw, dh
+
+def cal_rescale_size(image_size, w, h, k_size=2, scale_factor=1/16, no_print=False):
+    # Calculate target image size (lager side=image_size)
+    wt = int(np.floor(w/(max(w, h)/image_size)*scale_factor/k_size)/scale_factor*k_size)
+    ht = int(np.floor(h/(max(w, h)/image_size)*scale_factor/k_size)/scale_factor*k_size)
+    N = wt * ht * scale_factor * scale_factor / (k_size ** 2)
+    if not no_print:
+        print('Target size {} Original: (w={},h={}), Rescaled: (w={},h={}) , matches resolution: {}'.format(image_size, w, h, 
+                                                                                                        wt, ht, N))
+    return wt, ht
+
+def get_tuple_transform_ops(resize=None, normalize=True, unscale=False):
+    ops = []
+    if resize:
+        ops.append(TupleResize(resize))
+    if normalize: 
+        ops.append(TupleToTensorScaled())
+        ops.append(TupleNormalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]))  # Imagenet mean/std
+    else:
+        if unscale:
+            ops.append(TupleToTensorUnscaled())
+        else:
+            ops.append(TupleToTensorScaled())
+    return TupleCompose(ops)
+
+class ToTensorScaled(object):
+    '''Convert a RGB PIL Image to a CHW ordered Tensor, scale the range to [0, 1]'''
+    def __call__(self, im):
+        im = np.array(im, dtype=np.float32).transpose((2, 0, 1))
+        im /= 255.0 
+        return torch.from_numpy(im)
+
+    def __repr__(self):
+        return 'ToTensorScaled(./255)'
+    
+class TupleToTensorScaled(object):
+    def __init__(self):
+        self.to_tensor = ToTensorScaled()
+        
+    def __call__(self, im_tuple):
+        return [self.to_tensor(im) for im in im_tuple]
+
+    def __repr__(self):
+        return 'TupleToTensorScaled(./255)'
+    
+class ToTensorUnscaled(object):
+    '''Convert a RGB PIL Image to a CHW ordered Tensor'''
+    def __call__(self, im):    
+        return torch.from_numpy(np.array(im, dtype=np.float32).transpose((2, 0, 1)))
+
+    def __repr__(self):
+        return 'ToTensorUnscaled()'
+
+class TupleToTensorUnscaled(object):
+    '''Convert a RGB PIL Image to a CHW ordered Tensor'''
+    def __init__(self):
+        self.to_tensor = ToTensorUnscaled()
+
+    def __call__(self, im_tuple):
+        return [self.to_tensor(im) for im in im_tuple]
+
+    def __repr__(self):
+        return 'TupleToTensorUnscaled()'
+    
+class TupleResize(object):
+    def __init__(self, size):
+        self.size = size
+        self.resize = transforms.Resize(size, Image.BICUBIC)
+
+    def __call__(self, im_tuple):
+        return [self.resize(im) for im in im_tuple]
+    
+    def __repr__(self):
+        return 'TupleResize(size={})'.format(self.size)
+    
+class TupleNormalize(object):
+    def __init__(self, mean, std):
+        self.mean = mean
+        self.std = std
+        self.normalize = transforms.Normalize(mean=mean, std=std)
+
+    def __call__(self, im_tuple):
+        return [self.normalize(im) for im in im_tuple]
+
+    def __repr__(self):
+        return 'TupleNormalize(mean={}, std={})'.format(self.mean, self.std)
+
+class TupleCompose(object):
+    def __init__(self, transforms):
+        self.transforms = transforms
+
+    def __call__(self, im_tuple):
+        for t in self.transforms:
+            im_tuple = t(im_tuple)
+        return im_tuple
+
+    def __repr__(self):
+        format_string = self.__class__.__name__ + '('
+        for t in self.transforms:
+            format_string += '\n'
+            format_string += '    {0}'.format(t)
+        format_string += '\n)'
+        return format_string    

vismatch/third_party/imatch-toolbox/third_party/patch2pix/utils/eval/geometry.py

@@ -0,0 +1,90 @@
+import numpy as np
+from transforms3d.quaternions import quat2mat, mat2quat
+
+
+# The skew-symmetric matrix of vector
+skew = lambda v: np.array([[0, -v[2], v[1]], [v[2], 0, -v[0]], [-v[1], v[0], 0]])
+
+# Essential matrix & fundamental matrix 
+ess2fund = lambda K1, K2, E: np.linalg.inv(K2).T @ E @ np.linalg.inv(K1)
+ess2fund_inv = lambda K1_inv, K2_inv, E: K2_inv.T @ E @ K1_inv
+fund2ess = lambda F, K2, K1: K2.T @ F @ K1
+
+# Camera relative pose to fundamental matrix
+pose2ess = lambda R, t: skew(t.reshape(3,)) @ R
+pose2fund = lambda K1, K2, R, t: np.linalg.inv(K2).T @ R @ K1.T @ skew((K1 @ R.T).dot(t.reshape(3,)))
+pose2fund_inv = lambda K1, K2_inv, R, t: K2_inv.T @ R @ K1.T @ skew((K1 @ R.T).dot(t))
+
+# Normalize fundamental matrix
+normF = lambda F: F / F[-1,-1] # Normalize F by the last value
+normalize = lambda A:  A / np.linalg.norm(A)
+
+def compose_projection_matrix(R, t):
+    """Construct projection matrix 
+    Args:
+        - R: rotation matrix, size (3,3);
+        - t: translation vector, size (3,);
+    Return:
+        - projection matrix [R|t], size (3,4)
+    """
+    return np.hstack([R, np.expand_dims(t, axis=1)])
+
+def matches2relapose_cv(p1, p2, K1, K2, rthres=1):
+    import cv2
+    # Move back to image center based coordinates
+    f1, f2,  = K1[0,0], K2[0, 0]   
+    pc1 = np.array([K1[:2, 2]])
+    pc2 = np.array([K2[:2, 2]])
+
+    # Rescale to im2 's focal setting
+    p1 = (p1 - pc1) * f2 / f1
+    p2 = (p2 - pc2) 
+    K = np.array([[f2, 0, 0],
+                 [0, f2, 0],
+                 [0, 0, 1]])      
+    E, inls = cv2.findEssentialMat(p1, p2, cameraMatrix=K, method=cv2.FM_RANSAC, threshold=rthres)            
+    inls = np.where(inls > 0)[0]
+    _, R, t, _ = cv2.recoverPose(E, p1[inls], p2[inls], K)    
+    return E, inls, R, t
+
+def matches2relapose_degensac(p1, p2, K1, K2, rthres=1):
+    import pydegensac
+    import cv2    
+    
+    # Move back to image center based coordinates
+    f1, f2  = K1[0,0], K2[0, 0]   
+    pc1 = np.array([K1[:2, 2]])
+    pc2 = np.array([K2[:2, 2]])
+
+    # Rescale to im2 's focal setting
+    p1 = (p1 - pc1) * f2 / f1
+    p2 = (p2 - pc2) 
+    K = np.array([[f2, 0, 0],
+                 [0, f2, 0],
+                 [0, 0, 1]])
+    K1 = K2 = K
+    
+    F, inls = pydegensac.findFundamentalMatrix(p1, p2, rthres)    
+    E = fund2ess(F, K1, K2)
+    inls = np.where(inls > 0)[0]    
+    _, R, t, _ = cv2.recoverPose(E, p1[inls], p2[inls], K)
+    return E, inls, R, t
+
+def abs2relapose(c1, c2, q1, q2):
+    """Calculate relative pose between two cameras
+    Args:
+    - c1: absolute position of the first camera
+    - c2: absolute position of the second camera
+    - q1: orientation quaternion of the first camera
+    - q2: orientation quaternion of the second camera
+    Return:
+    - (t12, q12): relative pose giving the transformation from the 1st camera to the 2nd camera coordinates, 
+                  t12 is translation, q12 is relative rotation quaternion 
+    """
+    r1 = quat2mat(q1)
+    r2 = quat2mat(q2)
+    r12 = r2.dot(r1.T)
+    q12 = mat2quat(r12)
+    t12 = r2.dot(c1 - c2)
+    return (t12, q12)
+          

vismatch/third_party/imatch-toolbox/third_party/patch2pix/utils/eval/measure.py

@@ -0,0 +1,161 @@
+import numpy as np
+
+get_statis = lambda arr: 'Size={} Min={:.2f} Max={:.2f} Mean={:.2f} Median={:.2f}'.format(
+                                arr.shape, np.min(arr), np.max(arr), np.mean(arr), np.median(arr))
+
+def expand_homo_ones(arr2d, axis=1):
+    """Raise 2D array to homogenous coordinates
+    Args:
+        - arr2d: (N, 2) or (2, N)
+        - axis: the axis to append the ones
+    """    
+    if axis == 0:
+        ones = np.ones((1, arr2d.shape[1]))
+    else:
+        ones = np.ones((arr2d.shape[0], 1))      
+    return np.concatenate([arr2d, ones], axis=axis)
+
+def sampson_distance(pts1, pts2, F, homos=True, eps=1e-8):
+    """Calculate symmetric epipolar distance between 2 sets of points
+    Args:
+        - pts1, pts2: points correspondences in the two images, 
+          each has shape of (num_points, 2)
+        - F: fundamental matrix that fulfills x2^T*F*x1=0, 
+          where x1 and x2 are the correspondence points in the 1st and 2nd image 
+    Return:
+        A vector of (num_points,), containing root-squared epipolar distances
+          
+    """
+    
+    # Homogenous coordinates
+    if homos:
+        pts1 = expand_homo_ones(pts1, axis=1) #if pts1.shape[1] == 2 else pts1
+        pts2 = expand_homo_ones(pts2, axis=1) #if pts2.shape[1] == 2 else pts2
+    
+    # l2=F*x1, l1=F^T*x2
+    l2 = np.dot(F, pts1.T) # 3,N
+    l1 = np.dot(F.T, pts2.T)
+    dd = np.sum(l2.T * pts2, 1)  # Distance from pts2 to l2   
+    d = dd ** 2 / (eps + l1[0, :] ** 2 + l1[1, :] ** 2 + l2[0, :] ** 2 + l2[1, :] ** 2)   
+    return d
+
+
+def symmetric_epipolar_distance(pts1, pts2, F, homos=True, sqrt=False):
+    """Calculate symmetric epipolar distance between 2 sets of points
+    Args:
+        - pts1, pts2: points correspondences in the two images, 
+          each has shape of (num_points, 2)
+        - F: fundamental matrix that fulfills x2^T*F*x1=0, 
+          where x1 and x2 are the correspondence points in the 1st and 2nd image 
+    Return:
+        A vector of (num_points,), containing root-squared epipolar distances
+          
+    """
+    
+    # Homogenous coordinates
+    if homos:
+        pts1 = expand_homo_ones(pts1, axis=1) #if pts1.shape[1] == 2 else pts1
+        pts2 = expand_homo_ones(pts2, axis=1) #if pts2.shape[1] == 2 else pts2
+    
+    # l2=F*x1, l1=F^T*x2
+    l2 = np.dot(F, pts1.T) # 3,N
+    l1 = np.dot(F.T, pts2.T)
+    dd = np.sum(l2.T * pts2, 1)  # Distance from pts2 to l2
+    
+    if sqrt:
+        # The one following DFM and find correspondence paper
+        d = np.abs(dd) * (1.0 / np.sqrt(l1[0, :] ** 2 + l1[1, :] ** 2) + 1.0 / np.sqrt(l2[0, :] ** 2 + l2[1, :] ** 2))
+    else: 
+        # Original one as in MVG Hartley.
+        d = dd ** 2 * (1.0 / (l1[0, :] ** 2 + l1[1, :] ** 2) + 1.0 /(l2[0, :] ** 2 + l2[1, :] ** 2))
+    return d
+
+def cal_vec_angle_error(label, pred, eps=1e-14):
+    if len(label.shape) == 1:
+        label = np.expand_dims(label, axis=0)
+    if len(pred.shape) == 1:
+        pred = np.expand_dims(pred, axis=0)
+
+    v1 = pred / (np.linalg.norm(pred, axis=1, keepdims=True) + eps)
+    v2 = label / (np.linalg.norm(label, axis=1, keepdims=True) + eps)
+    d = np.sum(np.multiply(v1,v2), axis=1, keepdims=True) 
+    d = np.clip(d, a_min=-1, a_max=1)
+    error = np.degrees(np.arccos(d))    
+    return error.squeeze()
+
+def cal_quat_angle_error(label, pred, eps=1e-14):
+    if len(label.shape) == 1:
+        label = np.expand_dims(label, axis=0)
+    if len(pred.shape) == 1:
+        pred = np.expand_dims(pred, axis=0)
+    q1 = pred / (np.linalg.norm(pred, axis=1, keepdims=True) + eps)
+    q2 = label / (np.linalg.norm(label, axis=1, keepdims=True) + eps)
+    d = np.abs(np.sum(np.multiply(q1,q2), axis=1, keepdims=True))
+    d = np.clip(d, a_min=-1, a_max=1)
+    error = 2 * np.degrees(np.arccos(d))
+    return error.squeeze()
+
+def cal_rot_angle_error(Rgt, Rpred):
+    # Identical to quaternion angular error
+    return np.rad2deg(np.arccos((np.trace(Rpred.T.dot(Rgt)) - 1) / 2))
+
+def eval_matches_relapose(matches, K1, K2, q_, t_, cv_thres=1.0):
+    from utils.eval.geometry import  matches2relapose_cv
+    from transforms3d.quaternions import mat2quat
+    
+    p1 = matches[:,:2]
+    p2 = matches[:,2:4]        
+    E, inls, R, t = matches2relapose_cv(p1, p2, K1, K2, rthres=cv_thres)
+    
+    # Calculate relative angle errors
+    terr = cal_vec_angle_error(t.squeeze(), t_)
+    qerr = cal_quat_angle_error(mat2quat(R), q_)
+    return terr, qerr, inls
+
+def check_inliers_distr(inlier_dists, 
+                        bins = [0, 1e-2, 1, 5, 10, 25, 50, 100, 400, 2500, 1e5],
+                        tag='', return_ratios=False):
+    if not inlier_dists:
+        if return_ratios:
+            return None, ''
+        return ''
+    inlier_ratios = []
+    Npts = []
+    for dists in inlier_dists:
+        N = len(dists)
+        if N == 0:
+            continue
+        Npts.append(N)
+        hists = np.histogram(dists, bins)[0]
+        inlier_ratios.append(hists / N)
+
+    ratio_print = '{} Sample:{} N(mean/max/min):{:.0f}/{:.0f}/{:.0f}\nRatios(%):'.format(tag, len(inlier_dists), np.mean(Npts), 
+                                                                                         np.max(Npts), np.min(Npts))
+
+    ratios = []
+    for val, low, high in zip(np.mean(inlier_ratios, axis=0), bins[0:-1], bins[1::]):
+        ratio_print = '{} [{},{})={:.2f}'.format(ratio_print, low, high, 100*val)
+        ratios.append(100*val)
+    if return_ratios:
+            return ratios, ratio_print
+    return ratio_print
+
+def check_data_hist(data_list, bins, tag='', return_hist=False):
+    if not data_list:
+        return ''
+    hists = []
+    means = []
+    for data in data_list:
+        if len(data) == 0:
+            continue
+        nums = np.histogram(data, bins)[0]
+        hists.append(nums / len(data))
+        means.append(np.mean(data))
+
+    hist_print = f'{tag} mean={np.mean(means):.2f}'
+    mean_hists = np.mean(hists, axis=0)
+    for val, low, high in zip(mean_hists, bins[0:-1], bins[1::]):
+        hist_print += ' [{},{})={:.2f}'.format(low, high, 100 * val)
+    if return_hist:
+        return mean_hists, hist_print
+    return hist_print