# Copyright 2022-present NAVER Corp.
# CC BY-NC-SA 4.0
# Available only for non-commercial use
from pdb import set_trace as bb
import os, os.path as osp
from tqdm import tqdm
from PIL import Image
import numpy as np
import torch
from .image_set import ImageSet
from .transforms import instanciate_transforms
from .utils import DatasetWithRng
invh = np.linalg.inv
class ImagePairs (DatasetWithRng):
""" Base class for a dataset that serves image pairs.
"""
imgs = None # regular image dataset
pairs = [] # list of (idx1, idx2), ...
def __init__(self, image_set, pairs, trf=None, **rng):
assert image_set and pairs, 'empty images or pairs'
super().__init__(**rng)
self.imgs = image_set
self.pairs = pairs
self.trf = instanciate_transforms(trf, rng=self.rng)
def __len__(self):
return len(self.pairs)
def __getitem__(self, idx):
transform = self.trf or (lambda x:x)
pair = tuple(map(transform, self._load_pair(idx)))
return pair, {}
def _load_pair(self, idx):
i,j = self.pairs[idx]
img1 = self.imgs.get_image(i)
return (img1, img1) if i == j else (img1, self.imgs.get_image(j))
def __repr__(self):
return f'{self.__class__.__name__}({len(self)} pairs from {self.imgs})'
class StillImagePairs (ImagePairs):
""" A dataset of 'still' image pairs used for debugging purposes.
"""
def __init__(self, image_set, pairs=None, **rng):
if isinstance(image_set, ImagePairs):
super().__init__(image_set.imgs, pairs or image_set.pairs, **rng)
else:
super().__init__(image_set, pairs or [(i,i) for i in range(len(image_set))], **rng)
def __getitem__(self, idx):
img1, img2 = self._load_pair(idx)
sx, sy = img2.size / np.float32(img1.size)
return (img1, img2), dict(homography=np.diag(np.float32([sx, sy, 1])))
class SyntheticImagePairs (StillImagePairs):
""" A synthetic generator of image pairs.
Given a normal image dataset, it constructs pairs using random homographies & noise.
scale: prior image scaling.
distort: distortion applied independently to (img1,img2) if sym=True else just img2
sym: (bool) see above.
"""
def __init__(self, image_set, scale='', distort='', sym=False, **rng):
super().__init__(image_set, **rng)
self.symmetric = sym
self.scale = instanciate_transforms(scale, rng=self.rng)
self.distort = instanciate_transforms(distort, rng=self.rng)
def __getitem__(self, idx):
(img1, img2), gt = super().__getitem__(idx)
img1 = dict(img=img1, homography=np.eye(3,dtype=np.float32))
if img1['img'] is img2:
img1 = self.scale(img1)
img2 = self.distort(dict(img1))
if self.symmetric: img1 = self.distort(img1)
else:
if self.symmetric: img1 = self.distort(self.scale(img1))
img2 = self.distort(self.scale(dict(img=img2, **gt)))
return (img1['img'], img2['img']), dict(homography=img2['homography'] @ invh(img1['homography']))
def __repr__(self):
format = lambda s: ','.join(l.strip() for l in repr(s).splitlines() if l).replace(',','',1)
return f"{self.__class__.__name__}({len(self)} images, scale={format(self.scale)}, distort={format(self.distort)})"
class CatImagePairs (DatasetWithRng):
""" Concatenation of several ImagePairs datasets
"""
def __init__(self, *pair_datasets, seed=torch.initial_seed()):
assert all(isinstance(db, ImagePairs) for db in pair_datasets)
self.pair_datasets = pair_datasets
DatasetWithRng.__init__(self, seed=seed) # init last
self._init()
def _init(self):
self._pair_offsets = np.cumsum([0] + [len(db) for db in self.pair_datasets])
self.npairs = self._pair_offsets[-1]
def __len__(self):
return self.npairs
def __repr__(self):
fmt_str = f"{type(self).__name__}({len(self)} pairs,"
for i,db in enumerate(self.pair_datasets):
npairs = self._pair_offsets[i+1] - self._pair_offsets[i]
fmt_str += f'\n\t{npairs} from '+str(db).replace("\n"," ") + ','
return fmt_str[:-1] + ')'
def __getitem__(self, idx):
b, i = self._which(idx)
return self.pair_datasets[b].__getitem__(i)
def _which(self, i):
pos = np.searchsorted(self._pair_offsets, i, side='right')-1
assert pos < self.npairs, 'Bad pair index %d >= %d' % (i, self.npairs)
return pos, i - self._pair_offsets[pos]
def _call(self, func, i, *args, **kwargs):
b, j = self._which(i)
return getattr(self.pair_datasets[b], func)(j, *args, **kwargs)
def init_worker(self, tid):
for db in self.pair_datasets:
db.init_worker(tid)
class BalancedCatImagePairs (CatImagePairs):
""" Balanced concatenation of several ImagePairs datasets
"""
def __init__(self, npairs=0, *pair_datasets, **kw):
assert isinstance(npairs, int) and npairs >= 0, 'BalancedCatImagePairs(npairs != int)'
assert len(pair_datasets) > 0, 'no dataset provided'
if len(pair_datasets) >= 3 and isinstance(pair_datasets[1], int):
assert len(pair_datasets) % 2 == 1
pair_datasets = [npairs] + list(pair_datasets)
npairs, pair_datasets = pair_datasets[0::2], pair_datasets[1::2]
assert all(isinstance(n, int) for n in npairs)
self._pair_offsets = np.cumsum([0]+npairs)
self.npairs = self._pair_offsets[-1]
else:
self.npairs = npairs or max(len(db) for db in pair_datasets)
self._pair_offsets = np.linspace(0, self.npairs, len(pair_datasets)+1).astype(int)
CatImagePairs.__init__(self, *pair_datasets, **kw)
def set_epoch(self, epoch):
DatasetWithRng.init_worker(self, epoch) # random seed only depends on the epoch
self._init() # reset permutations for this epoch
def init_worker(self, tid):
CatImagePairs.init_worker(self, tid)
def _init(self):
self._perms = []
for i,db in enumerate(self.pair_datasets):
assert len(db), 'cannot balance if there is an empty dataset'
avail = self._pair_offsets[i+1] - self._pair_offsets[i]
idxs = np.arange(len(db))
while len(idxs) < avail:
idxs = np.r_[idxs,idxs]
if self.seed: # if not seed, then no shuffle
self.rng.shuffle(idxs[(avail//len(db))*len(db):])
self._perms.append( idxs[:avail] )
# print(self._perms)
def _which(self, i):
pos, idx = super()._which(i)
return pos, self._perms[pos][idx]
class UnsupervisedPairs (ImagePairs):
""" Unsupervised image pairs obtained from SfM
"""
def __init__(self, img_set, pair_file_path):
assert isinstance(img_set, ImageSet), bb()
self.pair_list = self._parse_pair_list(pair_file_path)
self.corres_dir = osp.join(osp.split(pair_file_path)[0], 'corres')
tag_to_idx = {n:i for i,n in enumerate(img_set.imgs)}
img_indices = lambda pair: tuple([tag_to_idx[n] for n in pair])
super().__init__(img_set, [img_indices(pair) for pair in self.pair_list])
def __repr__(self):
return f"{type(self).__name__}({len(self)} pairs from {self.imgs})"
def _parse_pair_list(self, pair_file_path):
res = []
for row in open(pair_file_path).read().splitlines():
row = row.split()
if len(row) != 2: raise IOError()
res.append((row[0], row[1]))
return res
def get_corres_path(self, pair_idx):
img1, img2 = [osp.basename(self.imgs.imgs[i]) for i in self.pairs[pair_idx]]
return osp.join(self.corres_dir, f'{img1}_{img2}.npy')
def get_corres(self, pair_idx):
return np.load(self.get_corres_path(pair_idx))
def __getitem__(self, idx):
img1, img2 = self._load_pair(idx)
return (img1, img2), dict(corres=self.get_corres(idx))
if __name__ == '__main__':
from datasets import *
from tools.viz import show_random_pairs
db = BalancedCatImagePairs(
3125, SyntheticImagePairs(RandomWebImages(0,52),distort='RandomTilting(0.5)'),
4875, SyntheticImagePairs(SfM120k_Images(),distort='RandomTilting(0.5)'),
8000, SfM120k_Pairs())
show_random_pairs(db)