vismatch 1.1.1__py3-none-any.whl

vismatch/third_party/LiftFeat/dataset/coco_augmentor.py

@@ -0,0 +1,298 @@
+"""
+    "LiftFeat: 3D Geometry-Aware Local Feature Matching"
+    COCO_20k image augmentor
+"""
+
+import torch
+from torch import nn
+from torch.utils.data import Dataset
+import torch.utils.data as data
+from torchvision import transforms
+import torch.nn.functional as F
+
+import cv2
+import kornia
+import kornia.augmentation as K
+from kornia.geometry.transform import get_tps_transform as findTPS
+from kornia.geometry.transform import warp_points_tps, warp_image_tps
+
+import glob
+import random
+import tqdm
+
+import numpy as np
+import pdb
+import time
+
+random.seed(0)
+torch.manual_seed(0)
+
+def generateRandomTPS(shape,grid=(8,6),GLOBAL_MULTIPLIER=0.3,prob=0.5):
+    
+    h, w = shape
+    sh, sw = h/grid[0], w/grid[1]
+    src = torch.dstack(torch.meshgrid(torch.arange(0, h + sh , sh), torch.arange(0, w + sw , sw), indexing='ij'))
+
+    offsets = torch.rand(grid[0]+1, grid[1]+1, 2) - 0.5
+    offsets *= torch.tensor([ sh/2, sw/2 ]).view(1, 1, 2)  * min(0.97, 2.0 * GLOBAL_MULTIPLIER)
+    dst = src + offsets if np.random.uniform() < prob else src
+    
+    src, dst = src.view(1, -1, 2), dst.view(1, -1, 2)
+    src = (src / torch.tensor([h,w]).view(1,1,2) ) * 2 - 1.
+    dst = (dst / torch.tensor([h,w]).view(1,1,2) ) * 2 - 1.
+    weights, A = findTPS(dst, src)
+
+    return src, weights, A
+
+
+def generateRandomHomography(shape,GLOBAL_MULTIPLIER=0.3):
+    #Generate random in-plane rotation [-theta,+theta]
+    theta = np.radians(np.random.uniform(-30, 30))
+
+    #Generate random scale in both x and y
+    scale_x, scale_y = np.random.uniform(0.35, 1.2, 2)
+
+    #Generate random translation shift
+    tx , ty = -shape[1]/2.0 , -shape[0]/2.0 
+    txn, tyn = np.random.normal(0, 120.0*GLOBAL_MULTIPLIER, 2) 
+
+    c, s = np.cos(theta), np.sin(theta)
+
+    #Affine coeffs
+    sx , sy = np.random.normal(0,0.6*GLOBAL_MULTIPLIER,2)
+
+    #Projective coeffs
+    p1 , p2 = np.random.normal(0,0.006*GLOBAL_MULTIPLIER,2)
+
+
+    # Build Homography from parmeterizations
+    H_t = np.array(((1,0, tx), (0, 1, ty), (0,0,1))) #t
+    H_r = np.array(((c,-s, 0), (s, c, 0), (0,0,1))) #rotation,
+    H_a = np.array(((1,sy, 0), (sx, 1, 0), (0,0,1))) # affine
+    H_p = np.array(((1, 0, 0), (0 , 1, 0), (p1,p2,1))) # projective
+    H_s = np.array(((scale_x,0, 0), (0, scale_y, 0), (0,0,1))) #scale
+    H_b = np.array(((1.0,0,-tx +txn), (0, 1, -ty + tyn), (0,0,1))) #t_back,
+
+    #H = H_e * H_s * H_a * H_p
+    H = np.dot(np.dot(np.dot(np.dot(np.dot(H_b,H_s),H_p),H_a),H_r),H_t)
+
+    return H
+
+
+class COCOAugmentor(nn.Module):
+    
+    def __init__(self,device,load_dataset=True,
+                img_dir="/home/yepeng_liu/code_python/dataset/coco_20k",
+                warp_resolution=(1200, 900),
+                out_resolution=(400, 300),
+                sides_crop=0.2,
+                max_num_imgs=50,
+                num_test_imgs=10,
+                batch_size=1,
+                photometric=True,
+                geometric=True,
+                reload_step=1_000
+                ):
+        super(COCOAugmentor,self).__init__()
+        self.half=16
+        self.device=device
+
+        self.dims=warp_resolution
+        self.batch_size=batch_size
+        self.out_resolution=out_resolution 
+        self.sides_crop=sides_crop
+        self.max_num_imgs=max_num_imgs
+        self.num_test_imgs=num_test_imgs
+        self.dims_t=torch.tensor([int(self.dims[0]*(1. - self.sides_crop)) - int(self.dims[0]*self.sides_crop) -1,
+                                  int(self.dims[1]*(1. - self.sides_crop)) - int(self.dims[1]*self.sides_crop) -1]).float().to(device).view(1,1,2)
+        self.dims_s=torch.tensor([self.dims_t[0,0,0] / out_resolution[0],
+                                  self.dims_t[0,0,1] / out_resolution[1]]).float().to(device).view(1,1,2) 
+
+        self.all_imgs=glob.glob(img_dir+'/*.jpg')+glob.glob(img_dir+'/*.png')
+        
+        self.photometric=photometric
+        self.geometric=geometric
+        self.cnt=1
+        self.reload_step=reload_step
+        
+        list_augmentation=[
+            kornia.augmentation.ColorJitter(0.15,0.15,0.15,0.15,p=1.),
+            kornia.augmentation.RandomEqualize(p=0.4),
+            kornia.augmentation.RandomGaussianBlur(p=0.3,sigma=(2.0,2.0),kernel_size=(7,7))
+        ]
+        
+        if photometric is False:
+            list_augmentation = []
+
+        self.aug_list=kornia.augmentation.ImageSequential(*list_augmentation)
+        
+        if len(self.all_imgs)<10:
+            raise RuntimeError('Couldnt find enough images to train. Please check the path: ',img_dir)
+
+        if load_dataset:
+            print('[COCO]: ',len(self.all_imgs),' images for training..')
+            if len(self.all_imgs) - num_test_imgs < max_num_imgs:
+                raise RuntimeError('Error: test set overlaps with training set! Decrease number of test imgs')
+
+            self.load_imgs()
+
+            self.TPS = True
+
+
+    def load_imgs(self):
+        random.shuffle(self.all_imgs)
+        train = []
+        for p in tqdm.tqdm(self.all_imgs[:self.max_num_imgs],desc='loading train'):
+            im=cv2.imread(p)
+            halfH,halfW=im.shape[0]//2,im.shape[1]//2
+            if halfH>halfW:
+                im=np.rot90(im)
+                halfH,halfW=halfW,halfH
+
+            if im.shape[0]!=self.dims[1] or im.shape[1]!=self.dims[0]:
+                im = cv2.resize(im, self.dims)
+
+            train.append(np.copy(im))
+
+        self.train=train
+        self.test=[
+            cv2.resize(cv2.imread(p),self.dims)
+            for p in tqdm.tqdm(self.all_imgs[-self.num_test_imgs:],desc='loading test')
+        ]
+
+    def norm_pts_grid(self, x):
+        if len(x.size()) == 2:
+            return (x.view(1,-1,2) * self.dims_s / self.dims_t) * 2. - 1 
+        return (x * self.dims_s / self.dims_t) * 2. - 1
+
+    def denorm_pts_grid(self, x):
+        if len(x.size()) == 2:
+            return ((x.view(1,-1,2) + 1) / 2.) / self.dims_s * self.dims_t
+        return ((x+1) / 2.) / self.dims_s * self.dims_t
+
+    def rnd_kps(self, shape, n = 256):
+        h, w = shape
+        kps = torch.rand(size = (3,n)).to(self.device)
+        kps[0,:]*=w
+        kps[1,:]*=h
+        kps[2,:] = 1.0
+
+        return kps
+
+    def warp_points(self, H, pts):
+      scale = self.dims_s.view(-1,2)
+      offset = torch.tensor([int(self.dims[0]*self.sides_crop), int(self.dims[1]*self.sides_crop)], device = pts.device).float()
+      pts = pts*scale + offset
+      pts = torch.vstack( [pts.t(), torch.ones(1, pts.shape[0], device = pts.device)])
+      warped = torch.matmul(H, pts)
+      warped = warped / warped[2,...] 
+      warped = warped.t()[:, :2]
+      return (warped - offset) / scale
+
+    @torch.inference_mode()
+    def forward(self, x, difficulty = 0.3, TPS = False, prob_deformation = 0.5, test = False):
+        """
+            Perform augmentation to a batch of images.
+
+            input:
+                x -> torch.Tensor(B, C, H, W): rgb images
+                difficulty -> float: level of difficulty, 0.1 is medium, 0.3 is already pretty hard
+                tps -> bool: Wether to apply non-rigid deformations in images
+                prob_deformation -> float: probability to apply a deformation
+
+            return:
+                'output'    ->   torch.Tensor(B, C, H, W): rgb images
+                Tuple:
+                    'H'       ->   torch.Tensor(3,3): homography matrix 
+                    'mask'  ->     torch.Tensor(B, H, W): mask of valid pixels after warp
+                    (deformation only)
+                    src, weights, A are parameters from a TPS warp (all torch.Tensors)
+
+        """
+
+        if self.cnt % self.reload_step == 0:
+            self.load_imgs()
+
+        if self.geometric is False:
+            difficulty = 0.
+
+        with torch.no_grad():
+            x = (x/255.).to(self.device)
+            b, c, h, w = x.shape
+            shape = (h, w)
+              
+            ######## Geometric Transformations
+
+            H = torch.tensor(np.array([generateRandomHomography(shape,difficulty) for b in range(self.batch_size)]),dtype=torch.float32).to(self.device)
+            
+            output = kornia.geometry.transform.warp_perspective(x,H,dsize=shape,padding_mode='zeros')
+
+            #crop % of image boundaries each side to reduce invalid pixels after warps
+            low_h = int(h * self.sides_crop); low_w = int(w*self.sides_crop)
+            high_h = int(h*(1. - self.sides_crop)); high_w= int(w * (1. - self.sides_crop))
+            output = output[..., low_h:high_h, low_w:high_w]
+            x = x[..., low_h:high_h, low_w:high_w]
+
+            #apply TPS if desired:
+            if TPS:
+                src, weights, A = None, None, None
+                for b in range(self.batch_size):
+                    b_src, b_weights, b_A = generateRandomTPS(shape, (8,6), difficulty, prob = prob_deformation)
+                    b_src, b_weights, b_A = b_src.to(self.device), b_weights.to(self.device), b_A.to(self.device)
+
+                    if src is None:
+                        src, weights, A = b_src, b_weights, b_A
+                    else:
+                        src = torch.cat((b_src, src))
+                        weights = torch.cat((b_weights, weights))
+                        A = torch.cat((b_A, A))
+
+                output = warp_image_tps(output, src, weights, A)
+
+            output = F.interpolate(output, self.out_resolution[::-1], mode = 'nearest')
+            x = F.interpolate(x, self.out_resolution[::-1], mode = 'nearest')
+
+            mask = ~torch.all(output == 0, dim=1, keepdim=True)
+            mask = mask.expand(-1,3,-1,-1)
+
+            # Make-up invalid regions with texture from the batch
+            rv = 1 if not TPS else 2
+            output_shifted = torch.roll(x, rv, 0)
+            output[~mask] = output_shifted[~mask]
+            mask = mask[:, 0, :, :]
+
+            ######## Photometric Transformations
+            output = self.aug_list(output)
+
+            b, c, h, w = output.shape
+            #Correlated Gaussian Noise
+            if np.random.uniform() > 0.5 and self.photometric:
+                noise = F.interpolate(torch.randn_like(output)*(10/255), (h//2, w//2))
+                noise = F.interpolate(noise, (h, w), mode = 'bicubic')
+                output = torch.clip( output + noise, 0., 1.)
+
+            #Random shadows
+            if np.random.uniform() > 0.6 and self.photometric:
+                noise = torch.rand((b, 1, h//64, w//64), device = self.device) * 1.3
+                noise = torch.clip(noise, 0.25, 1.0)
+                noise = F.interpolate(noise, (h, w), mode = 'bicubic')
+                noise = noise.expand(-1, 3, -1, -1)
+                output *= noise
+                output = torch.clip( output, 0., 1.)
+
+            self.cnt+=1
+
+        if TPS:
+            return output, (H, src, weights, A, mask)
+        else:
+            return output, (H, mask)
+
+    def get_correspondences(self, kps_target, T):
+        H, H2, src, W, A = T
+        undeformed  = self.denorm_pts_grid(   
+                                        warp_points_tps(self.norm_pts_grid(kps_target),
+                                        src, W, A) ).view(-1,2)
+
+        warped_to_src = self.warp_points(H@torch.inverse(H2), undeformed)
+
+        return warped_to_src

vismatch/third_party/LiftFeat/dataset/coco_wrapper.py

@@ -0,0 +1,175 @@
+import torch
+import numpy as np
+import pdb
+
+debug_cnt = -1
+
+def make_batch(augmentor, difficulty = 0.3, train = True):
+    Hs = []
+    img_list = augmentor.train if train else augmentor.test
+    dev = augmentor.device
+    batch_images = []
+
+    with torch.no_grad(): # we dont require grads in the augmentation
+        for b in range(augmentor.batch_size):
+            rdidx = np.random.randint(len(img_list))
+            img = torch.tensor(img_list[rdidx], dtype=torch.float32).permute(2,0,1).to(augmentor.device).unsqueeze(0)
+            batch_images.append(img)
+
+        batch_images = torch.cat(batch_images)
+
+        p1, H1 = augmentor(batch_images, difficulty)
+        p2, H2 = augmentor(batch_images, difficulty, TPS = True, prob_deformation = 0.7)
+        # p2, H2 = augmentor(batch_images, difficulty, TPS = False, prob_deformation = 0.7)
+
+    return p1, p2, H1, H2
+
+
+def plot_corrs(p1, p2, src_pts, tgt_pts):
+    import matplotlib.pyplot as plt
+    p1 = p1.cpu()
+    p2 = p2.cpu()
+    src_pts = src_pts.cpu() ; tgt_pts = tgt_pts.cpu()
+    rnd_idx = np.random.randint(len(src_pts), size=200)
+    src_pts = src_pts[rnd_idx, ...]
+    tgt_pts = tgt_pts[rnd_idx, ...]
+
+    #Plot ground-truth correspondences
+    fig, ax = plt.subplots(1,2,figsize=(18, 12))
+    colors = np.random.uniform(size=(len(tgt_pts),3))
+    #Src image
+    img = p1
+    for i, p in enumerate(src_pts):
+        ax[0].scatter(p[0],p[1],color=colors[i])
+    ax[0].imshow(img.permute(1,2,0).numpy()[...,::-1])
+
+    #Target img
+    img2 = p2
+    for i, p in enumerate(tgt_pts):
+        ax[1].scatter(p[0],p[1],color=colors[i])
+    ax[1].imshow(img2.permute(1,2,0).numpy()[...,::-1])
+    plt.show()
+
+
+def get_corresponding_pts(p1, p2, H, H2, augmentor, h, w, crop = None):
+    '''
+        Get dense corresponding points
+    '''
+    global debug_cnt
+    negatives, positives = [], []
+
+    with torch.no_grad():
+        #real input res of samples
+        rh, rw = p1.shape[-2:]
+        ratio = torch.tensor([rw/w, rh/h], device = p1.device)
+
+        (H, mask1) = H
+        (H2, src, W, A, mask2) = H2
+
+        #Generate meshgrid of target pts
+        x, y = torch.meshgrid(torch.arange(w, device=p1.device), torch.arange(h, device=p1.device), indexing ='xy')
+        mesh = torch.cat([x.unsqueeze(-1), y.unsqueeze(-1)], dim=-1)
+        target_pts = mesh.view(-1, 2) * ratio
+
+        #Pack all transformations into T
+        for batch_idx in range(len(p1)):
+            with torch.no_grad():
+                T = (H[batch_idx], H2[batch_idx], 
+                    src[batch_idx].unsqueeze(0), W[batch_idx].unsqueeze(0), A[batch_idx].unsqueeze(0))
+                #We now warp the target points to src image
+                src_pts = (augmentor.get_correspondences(target_pts, T) ) #target to src 
+                tgt_pts = (target_pts)
+            
+                #Check out of bounds points
+                mask_valid = (src_pts[:, 0] >=0) & (src_pts[:, 1] >=0) & \
+                            (src_pts[:, 0] < rw) & (src_pts[:, 1] < rh)
+
+                negatives.append( tgt_pts[~mask_valid] )            
+                tgt_pts = tgt_pts[mask_valid]
+                src_pts = src_pts[mask_valid]
+
+
+                #Remove invalid pixels
+                mask_valid =    mask1[batch_idx, src_pts[:,1].long(), src_pts[:,0].long()]  & \
+                                mask2[batch_idx, tgt_pts[:,1].long(), tgt_pts[:,0].long()]
+                tgt_pts = tgt_pts[mask_valid]
+                src_pts = src_pts[mask_valid]
+
+                # limit nb of matches if desired
+                if crop is not None:
+                    rnd_idx = torch.randperm(len(src_pts), device=src_pts.device)[:crop]
+                    src_pts = src_pts[rnd_idx]
+                    tgt_pts = tgt_pts[rnd_idx]
+
+                if debug_cnt >=0 and debug_cnt < 4:
+                    plot_corrs(p1[batch_idx], p2[batch_idx], src_pts , tgt_pts )
+                    debug_cnt +=1
+
+                src_pts = (src_pts / ratio)
+                tgt_pts = (tgt_pts / ratio)
+
+                #Check out of bounds points
+                padto = 10 if crop is not None else 2
+                mask_valid1 = (src_pts[:, 0] >= (0 + padto)) & (src_pts[:, 1] >= (0 + padto)) & \
+                             (src_pts[:, 0] < (w - padto)) & (src_pts[:, 1] < (h - padto))
+                mask_valid2 = (tgt_pts[:, 0] >= (0 + padto)) & (tgt_pts[:, 1] >= (0 + padto)) & \
+                             (tgt_pts[:, 0] < (w - padto)) & (tgt_pts[:, 1] < (h - padto))
+                mask_valid = mask_valid1 & mask_valid2
+                tgt_pts = tgt_pts[mask_valid]
+                src_pts = src_pts[mask_valid]         
+
+                #Remove repeated correspondences
+                lut_mat = torch.ones((h, w, 4), device = src_pts.device, dtype = src_pts.dtype) * -1
+                # src_pts_np = src_pts.cpu().numpy()
+                # tgt_pts_np = tgt_pts.cpu().numpy()
+                try:
+                    lut_mat[src_pts[:,1].long(), src_pts[:,0].long()] = torch.cat([src_pts, tgt_pts], dim=1)
+                    mask_valid = torch.all(lut_mat >= 0, dim=-1)
+                    points = lut_mat[mask_valid]
+                    positives.append(points)
+                except:
+                    pdb.set_trace()
+                    print('..')
+
+    return negatives, positives
+
+
+def crop_patches(tensor, coords, size = 7):
+    '''
+        Crop [size x size] patches around 2D coordinates from a tensor.
+    '''
+    B, C, H, W = tensor.shape
+
+    x, y = coords[:, 0], coords[:, 1]
+    y = y.view(-1, 1, 1)
+    x = x.view(-1, 1, 1)
+    halfsize = size // 2
+    # Create meshgrid for indexing
+    x_offset, y_offset = torch.meshgrid(torch.arange(-halfsize, halfsize+1), torch.arange(-halfsize, halfsize+1), indexing='xy')
+    y_offset = y_offset.to(tensor.device)
+    x_offset = x_offset.to(tensor.device)
+
+    # Compute indices around each coordinate
+    y_indices = (y + y_offset.view(1, size, size)).squeeze(0) + halfsize
+    x_indices = (x + x_offset.view(1, size, size)).squeeze(0) + halfsize
+
+    # Handle out-of-boundary indices with padding
+    tensor_padded = torch.nn.functional.pad(tensor, (halfsize, halfsize, halfsize, halfsize), mode='constant')
+
+    # Index tensor to get patches
+    patches = tensor_padded[:, :, y_indices, x_indices] # [B, C, N, H, W]
+    return patches
+
+def subpix_softmax2d(heatmaps, temp = 0.25):
+    N, H, W = heatmaps.shape
+    heatmaps = torch.softmax(temp * heatmaps.view(-1, H*W), -1).view(-1, H, W)
+    x, y = torch.meshgrid(torch.arange(W, device =  heatmaps.device ), torch.arange(H, device =  heatmaps.device ), indexing = 'xy')
+    x = x - (W//2)
+    y = y - (H//2)
+    #pdb.set_trace()
+    coords_x = (x[None, ...] * heatmaps)
+    coords_y = (y[None, ...] * heatmaps)
+    coords = torch.cat([coords_x[..., None], coords_y[..., None]], -1).view(N, H*W, 2)
+    coords = coords.sum(1)
+
+    return coords

vismatch/third_party/LiftFeat/dataset/dataset_utils.py

@@ -0,0 +1,183 @@
+"""
+	"LiftFeat: 3D Geometry-Aware Local Feature Matching"
+
+    MegaDepth data handling was adapted from 
+    LoFTR official code: https://github.com/zju3dv/LoFTR/blob/master/src/datasets/megadepth.py
+"""
+
+import io
+import cv2
+import numpy as np
+import h5py
+import torch
+from numpy.linalg import inv
+
+
+try:
+    # for internel use only
+    from .client import MEGADEPTH_CLIENT, SCANNET_CLIENT
+except Exception:
+    MEGADEPTH_CLIENT = SCANNET_CLIENT = None
+
+# --- DATA IO ---
+
+def load_array_from_s3(
+    path, client, cv_type,
+    use_h5py=False,
+):
+    byte_str = client.Get(path)
+    try:
+        if not use_h5py:
+            raw_array = np.fromstring(byte_str, np.uint8)
+            data = cv2.imdecode(raw_array, cv_type)
+        else:
+            f = io.BytesIO(byte_str)
+            data = np.array(h5py.File(f, 'r')['/depth'])
+    except Exception as ex:
+        print(f"==> Data loading failure: {path}")
+        raise ex
+
+    assert data is not None
+    return data
+
+
+def imread_gray(path, augment_fn=None, client=SCANNET_CLIENT):
+    cv_type = cv2.IMREAD_GRAYSCALE if augment_fn is None \
+                else cv2.IMREAD_COLOR
+    if str(path).startswith('s3://'):
+        image = load_array_from_s3(str(path), client, cv_type)
+    else:
+        image = cv2.imread(str(path), 1)
+
+    if augment_fn is not None:
+        image = cv2.imread(str(path), cv2.IMREAD_COLOR)
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        image = augment_fn(image)
+        image = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
+    return image  # (h, w)
+
+
+def get_resized_wh(w, h, resize=None):
+    if resize is not None:  # resize the longer edge
+        scale = resize / max(h, w)
+        w_new, h_new = int(round(w*scale)), int(round(h*scale))
+    else:
+        w_new, h_new = w, h
+    return w_new, h_new
+
+
+def get_divisible_wh(w, h, df=None):
+    if df is not None:
+        w_new, h_new = map(lambda x: int(x // df * df), [w, h])
+    else:
+        w_new, h_new = w, h
+    return w_new, h_new
+
+
+def pad_bottom_right(inp, pad_size, ret_mask=False):
+    assert isinstance(pad_size, int) and pad_size >= max(inp.shape[-2:]), f"{pad_size} < {max(inp.shape[-2:])}"
+    mask = None
+    if inp.ndim == 2:
+        padded = np.zeros((pad_size, pad_size), dtype=inp.dtype)
+        padded[:inp.shape[0], :inp.shape[1]] = inp
+        if ret_mask:
+            mask = np.zeros((pad_size, pad_size), dtype=bool)
+            mask[:inp.shape[0], :inp.shape[1]] = True
+    elif inp.ndim == 3:
+        padded = np.zeros((inp.shape[0], pad_size, pad_size), dtype=inp.dtype)
+        padded[:, :inp.shape[1], :inp.shape[2]] = inp
+        if ret_mask:
+            mask = np.zeros((inp.shape[0], pad_size, pad_size), dtype=bool)
+            mask[:, :inp.shape[1], :inp.shape[2]] = True
+    else:
+        raise NotImplementedError()
+    return padded, mask
+
+
+# --- MEGADEPTH ---
+
+def fix_path_from_d2net(path):
+    if not path:
+        return None
+
+    path = path.replace('Undistorted_SfM/', '')
+    path = path.replace('images', 'dense0/imgs')
+    path = path.replace('phoenix/S6/zl548/MegaDepth_v1/', '')
+
+    return path
+
+def read_megadepth_gray(path, resize=None, df=None, padding=False, augment_fn=None):
+    """
+    Args:
+        resize (int, optional): the longer edge of resized images. None for no resize.
+        padding (bool): If set to 'True', zero-pad resized images to squared size.
+        augment_fn (callable, optional): augments images with pre-defined visual effects
+    Returns:
+        image (torch.tensor): (1, h, w)
+        mask (torch.tensor): (h, w)
+        scale (torch.tensor): [w/w_new, h/h_new]        
+    """
+    # read image
+    image = imread_gray(path, augment_fn, client=MEGADEPTH_CLIENT)
+
+    # resize image
+    w, h = image.shape[1], image.shape[0]
+
+    if resize is not None:
+        if len(resize) == 2:
+            w_new, h_new = resize
+        else:
+            resize = resize[0]
+            w_new, h_new = get_resized_wh(w, h, resize)
+            w_new, h_new = get_divisible_wh(w_new, h_new, df)
+
+
+        image = cv2.resize(image, (w_new, h_new))
+        scale = torch.tensor([w/w_new, h/h_new], dtype=torch.float)
+
+        if padding:  # padding
+            pad_to = max(h_new, w_new)
+            image, mask = pad_bottom_right(image, pad_to, ret_mask=True)
+        else:
+            mask = None
+    else:
+        scale=torch.tensor([1.0,1.0],dtype=torch.float)
+        
+        if padding:
+            pad_to=max(w,h)
+            image,mask=pad_bottom_right(image,pad_to,ret_mask=True)
+        else:
+            mask=None
+
+    #image = torch.from_numpy(image).float()[None] / 255  # (h, w) -> (1, h, w) and normalized
+    image_t = torch.from_numpy(image).float().permute(2,0,1) / 255  # (h, w) -> (1, h, w) and normalized
+    mask = torch.from_numpy(mask) if mask is not None else None
+
+    return image, image_t, mask, scale
+
+
+def read_megadepth_depth(path, pad_to=None):
+
+    if str(path).startswith('s3://'):
+        depth = load_array_from_s3(path, MEGADEPTH_CLIENT, None, use_h5py=True)
+    else:
+        depth = np.array(h5py.File(path, 'r')['depth'])
+    if pad_to is not None:
+        depth, _ = pad_bottom_right(depth, pad_to, ret_mask=False)
+    depth = torch.from_numpy(depth).float()  # (h, w)
+    return depth
+
+
+def imread_bgr(path, augment_fn=None, client=SCANNET_CLIENT):
+    cv_type = cv2.IMREAD_GRAYSCALE if augment_fn is None else cv2.IMREAD_COLOR
+    if str(path).startswith('s3://'):
+        image = load_array_from_s3(str(path), client, cv_type)
+    else:
+        image = cv2.imread(str(path), 1)
+
+    if augment_fn is not None:
+        image = cv2.imread(str(path), cv2.IMREAD_COLOR)
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        image = augment_fn(image)
+        image = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
+    return image  # (h, w)