vismatch 1.1.1__py3-none-any.whl

vismatch/third_party/LISRD/lisrd/datasets/hpatches.py

@@ -0,0 +1,135 @@
+""" HPatches dataset. """
+
+import os
+import numpy as np
+import cv2
+import logging
+from pathlib import Path
+import torch
+import torch.nn.functional as F
+from torch.utils.data import Dataset
+
+from .base_dataset import BaseDataset
+from .utils.data_reader import resize_and_crop
+from ..utils.geometry_utils import select_k_best
+from ..utils.pytorch_utils import keypoints_to_grid
+
+
+class Hpatches(BaseDataset):
+    def __init__(self, config, device):
+        super().__init__(config, device)
+
+    def get_dataset(self, split):
+        assert split == 'test', 'Only test split supported'
+        return _Dataset(self._config)
+
+
+class _Dataset(Dataset):
+    def __init__(self, config):
+        self._config = config   
+        root_dir = Path(os.path.expanduser(config['data_path']))
+        folder_paths = [x for x in root_dir.iterdir() if x.is_dir()]
+        if len(folder_paths) == 0:
+            raise ValueError(
+                f'Could not find any image in folder: {root_dir}.')
+        logging.info(f'Found {len(folder_paths)} scenes in image folder.')
+        self._image0_paths = []
+        self._image1_paths = []
+        self._homographies = []
+        for path in folder_paths:
+            if (config['alteration'] != 'all'
+                and config['alteration'] != path.stem[0]):
+                continue
+            for i in range(2, 7):
+                self._image0_paths.append(str(Path(path, "1.ppm")))
+                self._image1_paths.append(str(Path(path, str(i) + '.ppm')))
+                self._homographies.append(
+                    np.loadtxt(str(Path(path, "H_1_" + str(i)))))
+    
+    def adapt_homography_to_preprocessing(self, H, img_shape0, img_shape1):
+        source_size0 = np.array(img_shape0, dtype=float)
+        source_size1 = np.array(img_shape1, dtype=float)
+        target_size = np.array(self._config['resize'], dtype=float)
+        
+        # Get the scaling factor in resize
+        scale0 = np.amax(target_size / source_size0)
+        scaling0 = np.diag([1. / scale0, 1. / scale0, 1.]).astype(float)
+        scale1 = np.amax(target_size / source_size1)
+        scaling1 = np.diag([scale1, scale1, 1.]).astype(float)
+
+        # Get the translation params in crop
+        pad_y0 = (source_size0[0] * scale0 - target_size[0]) / 2.
+        pad_x0 = (source_size0[1] * scale0 - target_size[1]) / 2.
+        translation0 = np.array([[1., 0., pad_x0],
+                                 [0., 1., pad_y0],
+                                 [0., 0., 1.]], dtype=float)
+        pad_y1 = (source_size1[0] * scale1 - target_size[0]) / 2.
+        pad_x1 = (source_size1[1] * scale1 - target_size[1]) / 2.
+        translation1 = np.array([[1., 0., -pad_x1],
+                                 [0., 1., -pad_y1],
+                                 [0., 0., 1.]], dtype=float)
+
+        return translation1 @ scaling1 @ H @ scaling0 @ translation0
+
+    def __getitem__(self, item):
+        img0_path = self._image0_paths[item]
+        img0 = cv2.imread(img0_path)
+        img0 = cv2.cvtColor(img0, cv2.COLOR_BGR2RGB)
+        img1_path = self._image1_paths[item]
+        img1 = cv2.imread(img1_path)
+        img1 = cv2.cvtColor(img1, cv2.COLOR_BGR2RGB)
+        H = self._homographies[item].astype(float)
+
+        if 'resize' in self._config:
+            H = self.adapt_homography_to_preprocessing(H, img0.shape[:2],
+                                                       img1.shape[:2])
+            img0 = resize_and_crop(img0, self._config['resize'])
+            img1 = resize_and_crop(img1, self._config['resize'])
+        img_size = img0.shape[:2]
+        H_inv = np.linalg.inv(H)
+
+        img0 = img0.astype(float) / 255.
+        img1 = img1.astype(float) / 255.
+
+        # Extract the keypoints and descriptors of each method
+        features = {}
+        for m in self._config['models_name']:
+            features[m] = {}
+            feat0 = np.load(img0_path + '.' + m)
+            feat1 = np.load(img1_path + '.' + m)
+
+            # Extract a fixed number of shared keypoints between the images
+            kp0, mask0 = select_k_best(
+                feat0['keypoints'][:, [1, 0]], feat0['scores'],
+                self._config['num_kp'], H, img_size, margin=3)
+            features[m]['keypoints0'] = kp0
+            kp1, mask1 = select_k_best(
+                feat1['keypoints'][:, [1, 0]], feat1['scores'],
+                self._config['num_kp'], H_inv, img_size, margin=3)
+            features[m]['keypoints1'] = kp1
+
+            # Extract the local descriptors
+            features[m]['descriptors0'] = feat0['descriptors'][mask0]
+            features[m]['descriptors1'] = feat1['descriptors'][mask1]
+
+            # Extract meta descriptors if they exist
+            if 'meta_descriptors' in feat0:
+                meta_desc0_t = torch.tensor(feat0['meta_descriptors'])
+                grid0 = keypoints_to_grid(torch.tensor(kp0),
+                                          img_size).repeat(4, 1, 1, 1)
+                features[m]['meta_descriptors0'] = F.normalize(F.grid_sample(
+                    meta_desc0_t,
+                    grid0).squeeze(3).permute(2, 0, 1), dim=2).numpy()
+                meta_desc1_t = torch.tensor(feat1['meta_descriptors'])
+                grid1 = keypoints_to_grid(torch.tensor(kp1),
+                                          img_size).repeat(4, 1, 1, 1)
+                features[m]['meta_descriptors1'] = F.normalize(F.grid_sample(
+                    meta_desc1_t,
+                    grid1).squeeze(3).permute(2, 0, 1), dim=2).numpy()
+        
+        return {'image0': img0, 'image1': img1, 'homography': H,
+                'img0_path': img0_path, 'img1_path': img1_path,
+                'features': features, 'img_size': img_size}
+
+    def __len__(self):
+        return len(self._homographies)

vismatch/third_party/LISRD/lisrd/datasets/mixed_dataset.py

@@ -0,0 +1,53 @@
+""" Compose multiple datasets in a single loader. """
+
+import os
+import numpy as np
+import torch
+import cv2
+from copy import deepcopy
+from pathlib import Path
+from torch.utils.data import Dataset
+
+from .base_dataset import BaseDataset
+from .utils.homographies import sample_homography
+from .utils.data_augmentation import photometric_augmentation
+from . import get_dataset
+
+
+class MixedDataset(BaseDataset):
+    def __init__(self, config, device):
+        super().__init__(config, device)
+        base_config = {'sizes': {'train': 30000, 'val': 500, 'test': 1000}}
+        base_config.update(self._config)
+        self._config = base_config.copy()
+
+        # Initialize the datasets
+        self._datasets = []
+        for i, d in enumerate(config['datasets']):
+            base_config['name'] = d
+            base_config['data_path'] = config['data_paths'][i]
+            base_config['photometric_augmentation']['enable'] = config[
+                'photo_aug'][i]
+            self._datasets.append(get_dataset(d)(deepcopy(base_config),
+                                                 device))
+
+        self._weights = config['weights']
+
+    def get_dataset(self, split):
+        return _Dataset(self._datasets, self._weights,
+                        split, self._config['sizes'][split])
+    
+
+class _Dataset(Dataset):
+    def __init__(self, datasets, weights, split, size):
+        self._datasets = [d.get_dataset(split) for d in datasets]
+        self._weights = weights
+        self._size = size
+
+    def __getitem__(self, item):
+        dataset = self._datasets[np.random.choice(
+            range(len(self._datasets)), p=self._weights)]
+        return dataset[np.random.randint(len(dataset))]
+
+    def __len__(self):
+        return self._size

vismatch/third_party/LISRD/lisrd/datasets/rdnim.py

@@ -0,0 +1,117 @@
+""" Rotated Day-Night Image Matching dataset. """
+
+import os
+import numpy as np
+import cv2
+import torch
+import torch.nn.functional as F
+from pathlib import Path
+from torch.utils.data import Dataset
+
+from .base_dataset import BaseDataset
+from .utils.data_reader import read_timestamps
+from ..utils.geometry_utils import select_k_best
+from ..utils.pytorch_utils import keypoints_to_grid
+
+
+class Rdnim(BaseDataset):
+    def __init__(self, config, device):
+        super().__init__(config, device)
+        root_dir = Path(os.path.expanduser(config['data_path']))
+        ref = config['reference']
+
+        # Extract the timestamps
+        timestamp_files = [p for p in Path(root_dir, 'time_stamps').iterdir()]
+        timestamps = {}
+        for f in timestamp_files:
+            id = f.stem
+            timestamps[id] = read_timestamps(str(f))
+
+        # Extract the reference images paths
+        references = {}
+        seq_paths = [p for p in Path(root_dir, 'references').iterdir()]
+        for seq in seq_paths:
+            id = seq.stem
+            references[id] = str(Path(seq, ref + '.jpg'))
+
+        # Extract the images paths and homographies
+        seq_path = [p for p in Path(root_dir, 'images').iterdir()]
+        self._paths = {'test': []}
+        for seq in seq_path:
+            id = seq.stem
+            images_path = [x for x in seq.iterdir() if x.suffix == '.jpg']
+            for img in images_path:
+                timestamp = timestamps[id]['time'][
+                    timestamps[id]['name'].index(img.name)]
+                H = np.loadtxt(str(img)[:-4] + '.txt').astype(float)
+                self._paths['test'].append({
+                    'img': str(img), 'ref': str(references[id]),
+                    'H': H, 'timestamp': timestamp})
+
+    def get_dataset(self, split):
+        assert split == 'test', 'RDNIM only available in test mode.'
+        return _Dataset(self._paths[split], self._config)
+
+
+class _Dataset(Dataset):
+    def __init__(self, paths, config):
+        self._paths = paths
+        self._config = config
+
+    def __getitem__(self, item):
+        img0_path = self._paths[item]['ref']
+        img0 = cv2.cvtColor(cv2.imread(img0_path), cv2.COLOR_BGR2RGB)
+        img1_path = self._paths[item]['img']
+        img1 = cv2.cvtColor(cv2.imread(img1_path), cv2.COLOR_BGR2RGB)
+        img_size = img0.shape[:2]
+        
+        # Normalize the images in [0, 1]
+        img0 = img0.astype(float) / 255.
+        img1 = img1.astype(float) / 255.
+        
+        # Compute valid masks
+        H = self._paths[item]['H']
+        H_inv = np.linalg.inv(H)
+
+        # Extract the keypoints and descriptors of each method
+        features = {}
+        for m in self._config['models_name']:
+            features[m] = {}
+            feat0 = np.load(self._paths[item]['ref'] + '.' + m)
+            feat1 = np.load(self._paths[item]['img'] + '.' + m)
+
+            # Extract a fixed number of shared keypoints between the images
+            kp0, mask0 = select_k_best(
+                feat0['keypoints'][:, [1, 0]], feat0['scores'],
+                self._config['num_kp'], H, img_size, margin=3)
+            features[m]['keypoints0'] = kp0
+            kp1, mask1 = select_k_best(
+                feat1['keypoints'][:, [1, 0]], feat1['scores'],
+                self._config['num_kp'], H_inv, img_size, margin=3)
+            features[m]['keypoints1'] = kp1
+
+            # Extract the local descriptors
+            features[m]['descriptors0'] = feat0['descriptors'][mask0]
+            features[m]['descriptors1'] = feat1['descriptors'][mask1]
+
+            # Extract meta descriptors if they exist
+            if 'meta_descriptors' in feat0:
+                meta_desc0_t = torch.tensor(feat0['meta_descriptors'])
+                grid0 = keypoints_to_grid(torch.tensor(kp0),
+                                          img_size).repeat(4, 1, 1, 1)
+                features[m]['meta_descriptors0'] = F.normalize(F.grid_sample(
+                    meta_desc0_t,
+                    grid0).squeeze(3).permute(2, 0, 1), dim=2).numpy()
+                meta_desc1_t = torch.tensor(feat1['meta_descriptors'])
+                grid1 = keypoints_to_grid(torch.tensor(kp1),
+                                          img_size).repeat(4, 1, 1, 1)
+                features[m]['meta_descriptors1'] = F.normalize(F.grid_sample(
+                    meta_desc1_t,
+                    grid1).squeeze(3).permute(2, 0, 1), dim=2).numpy()
+                
+        return {'img0': img0, 'img1': img1, 'img_size': img_size,
+                'timestamp': self._paths[item]['timestamp'],
+                'features': features, 'homography': H}
+
+    def __len__(self):
+        return len(self._paths)

vismatch/third_party/LISRD/lisrd/datasets/utils/data_augmentation.py

@@ -0,0 +1,168 @@
+"""
+Common photometric transforms for data augmentation.
+"""
+import numpy as np
+import cv2
+from PIL import Image
+from torchvision import transforms as transforms
+
+
+available_augmentations = [
+    'additive_gaussian_noise',
+    'additive_speckle_noise',
+    'random_brightness',
+    'random_contrast',
+    'additive_shade',
+    'motion_blur'
+]
+
+
+class additive_gaussian_noise(object):
+    def __init__(self, stddev_range=[5, 95]):
+        self.stddev_range = stddev_range
+
+    def __call__(self, input_image):
+        # Get the noise stddev
+        stddev = np.random.uniform(self.stddev_range[0], self.stddev_range[1])
+        noise = np.random.normal(0., stddev, size=input_image.shape)
+        noisy_image = (input_image + noise).clip(0., 255.)
+
+        return noisy_image
+
+
+class additive_speckle_noise(object):
+    def __init__(self, prob_range=[0.0, 0.01]):
+        self.prob_range = prob_range
+
+    def __call__(self, input_image):
+        # Sample
+        prob = np.random.uniform(self.prob_range[0], self.prob_range[1])
+        sample = np.random.uniform(0., 1., size=input_image.shape)
+
+        # Get the mask
+        mask0 = sample <= prob
+        mask1 = sample >= (1 - prob)
+
+        # Mask the image (here we assume the image ranges from 0~255
+        noisy = input_image.copy()
+        noisy[mask0] = 0.
+        noisy[mask1] = 255.
+
+        return noisy
+
+
+class random_brightness(object):
+    def __init__(self, brightness=0.5):
+        self.brightness = brightness
+
+        # Initialize the transformer
+        self.transform = transforms.ColorJitter(brightness=self.brightness)
+
+    def __call__(self, input_image):
+        # Convert to PIL image
+        if isinstance(input_image, np.ndarray):
+            input_image = Image.fromarray(input_image.astype(np.uint8))
+
+        return np.array(self.transform(input_image))
+
+
+class random_contrast(object):
+    def __init__(self, strength_range=[0.5, 1.5]):
+        self.strength_range = strength_range
+
+    def __call__(self, input_image):
+        strength = np.random.uniform(self.strength_range[0],
+                                     self.strength_range[1])
+        contrasted_img = input_image.copy()
+        mean = np.mean(contrasted_img)
+        contrasted_img = (contrasted_img - mean) * strength + mean
+
+        return contrasted_img.clip(0, 255)
+
+
+class additive_shade(object):
+    def __init__(self, nb_ellipses=20, transparency_range=[-0.8, 0.8],
+                   kernel_size_range=[100, 150]):
+        self.nb_ellipses = nb_ellipses
+        self.transparency_range = transparency_range
+        self.kernel_size_range = kernel_size_range
+
+    def __call__(self, input_image):
+        min_dim = min(input_image.shape[:2]) / 4
+        mask = np.zeros(input_image.shape[:2], np.uint8)
+        for _ in range(self.nb_ellipses):
+            ax = int(max(np.random.rand() * min_dim, min_dim / 5))
+            ay = int(max(np.random.rand() * min_dim, min_dim / 5))
+            max_rad = max(ax, ay)
+            x = np.random.randint(max_rad, input_image.shape[1] - max_rad)
+            y = np.random.randint(max_rad, input_image.shape[0] - max_rad)
+            angle = np.random.rand() * 90
+            cv2.ellipse(mask, (x, y), (ax, ay), angle, 0, 360, 255, -1)
+
+        transparency = np.random.uniform(*self.transparency_range)
+        kernel_size = np.random.randint(*self.kernel_size_range)
+
+        # kernel_size has to be odd
+        if (kernel_size % 2) == 0:
+            kernel_size += 1
+        mask = cv2.GaussianBlur(mask.astype(np.float32),
+                                (kernel_size, kernel_size), 0)
+        if len(input_image.shape) == 2:
+            shaded = input_image[:, :, None] * (1 - transparency
+                                                * mask[..., np.newaxis] / 255.)
+        else:
+            shaded = input_image * (1 - transparency
+                                    * mask[..., np.newaxis] / 255.)
+        shaded = np.clip(shaded, 0, 255)
+
+        return np.reshape(shaded, input_image.shape)
+
+
+class motion_blur(object):
+    def __init__(self, max_kernel_size=10):
+        self.max_kernel_size = max_kernel_size
+
+    def __call__(self, input_image):
+        # Either vertical, horizontal or diagonal blur
+        mode = np.random.choice(['h', 'v', 'diag_down', 'diag_up'])
+        ksize = np.random.randint(
+            0, int(round((self.max_kernel_size + 1) / 2))) * 2 + 1
+        center = int((ksize - 1) / 2)
+        kernel = np.zeros((ksize, ksize))
+        if mode == 'h':
+            kernel[center, :] = 1.
+        elif mode == 'v':
+            kernel[:, center] = 1.
+        elif mode == 'diag_down':
+            kernel = np.eye(ksize)
+        elif mode == 'diag_up':
+            kernel = np.flip(np.eye(ksize), 0)
+        var = ksize * ksize / 16.
+        grid = np.repeat(np.arange(ksize)[:, np.newaxis], ksize, axis=-1)
+        gaussian = np.exp(-(np.square(grid - center) + np.square(grid.T - center)) / (2. * var))
+        kernel *= gaussian
+        kernel /= np.sum(kernel)
+        blurred = cv2.filter2D(input_image, -1, kernel)
+
+        return np.reshape(blurred, input_image.shape)
+
+
+def photometric_augmentation(input_img, config):
+    """ Process the input image through multiple transforms. """
+    if 'primitives' in config:
+        transforms = config['primitives']
+    else:
+        transforms = available_augmentations
+    
+    # Take a random subset of transforms
+    n_transforms = len(transforms)
+    n_used = np.random.randint(n_transforms)
+    transforms = np.random.choice(transforms, n_used, replace=False)
+
+    # Apply the transforms
+    transformed_img = input_img.copy()
+    for primitive in transforms:
+        transform = globals()[primitive](**config['params'][primitive])
+        transformed_img = transform(transformed_img)
+    
+    return transformed_img

vismatch/third_party/LISRD/lisrd/datasets/utils/data_reader.py

@@ -0,0 +1,48 @@
+import csv
+import numpy as np
+import cv2
+
+
+def resize_and_crop(image, img_size):
+    """ Resize an image to the given img_size by first rescaling it
+        and then applying a central crop to fit the given dimension. """
+    source_size = np.array(image.shape[:2], dtype=float)
+    target_size = np.array(img_size, dtype=float)
+
+    # Scale
+    scale = np.amax(target_size / source_size)
+    inter_size = np.round(source_size * scale).astype(int)
+    image = cv2.resize(image, (inter_size[1], inter_size[0]))
+
+    # Central crop
+    pad = np.round((source_size * scale - target_size) / 2.).astype(int)
+    image = image[pad[0]:(pad[0] + int(target_size[0])),
+                    pad[1]:(pad[1] + int(target_size[1])), :]
+    
+    return image
+
+
+def read_timestamps(text_file):
+    """
+    Read a text file containing the timestamps of images
+    and return a dictionary matching the name of the image
+    to its timestamp.
+    """
+    timestamps = {'name': [], 'date': [], 'hour': [],
+                  'minute': [], 'time': []}
+    with open(text_file, 'r') as csvfile:
+        reader = csv.reader(csvfile, delimiter=' ')
+        for row in reader:
+            timestamps['name'].append(row[0])
+            timestamps['date'].append(row[1])
+            hour = int(row[2])
+            timestamps['hour'].append(hour)
+            minute = int(row[3])
+            timestamps['minute'].append(minute)
+            timestamps['time'].append(hour + minute / 60.)
+    return timestamps
+
+
+def ascii_to_string(s):
+    """ Convert the array s of ascii values into the corresponding string. """
+    return ''.join(chr(i) for i in s)