Spaces:
Running
Running
| # Copyright 2019-present NAVER Corp. | |
| # CC BY-NC-SA 3.0 | |
| # Available only for non-commercial use | |
| import os, pdb | |
| import numpy as np | |
| from PIL import Image | |
| from .dataset import Dataset, CatDataset | |
| from tools.transforms import instanciate_transformation | |
| from tools.transforms_tools import persp_apply | |
| class PairDataset(Dataset): | |
| """A dataset that serves image pairs with ground-truth pixel correspondences.""" | |
| def __init__(self): | |
| Dataset.__init__(self) | |
| self.npairs = 0 | |
| def get_filename(self, img_idx, root=None): | |
| if is_pair( | |
| img_idx | |
| ): # if img_idx is a pair of indices, we return a pair of filenames | |
| return tuple(Dataset.get_filename(self, i, root) for i in img_idx) | |
| return Dataset.get_filename(self, img_idx, root) | |
| def get_image(self, img_idx): | |
| if is_pair( | |
| img_idx | |
| ): # if img_idx is a pair of indices, we return a pair of images | |
| return tuple(Dataset.get_image(self, i) for i in img_idx) | |
| return Dataset.get_image(self, img_idx) | |
| def get_corres_filename(self, pair_idx): | |
| raise NotImplementedError() | |
| def get_homography_filename(self, pair_idx): | |
| raise NotImplementedError() | |
| def get_flow_filename(self, pair_idx): | |
| raise NotImplementedError() | |
| def get_mask_filename(self, pair_idx): | |
| raise NotImplementedError() | |
| def get_pair(self, idx, output=()): | |
| """returns (img1, img2, `metadata`) | |
| `metadata` is a dict() that can contain: | |
| flow: optical flow | |
| aflow: absolute flow | |
| corres: list of 2d-2d correspondences | |
| mask: boolean image of flow validity (in the first image) | |
| ... | |
| """ | |
| raise NotImplementedError() | |
| def get_paired_images(self): | |
| fns = set() | |
| for i in range(self.npairs): | |
| a, b = self.image_pairs[i] | |
| fns.add(self.get_filename(a)) | |
| fns.add(self.get_filename(b)) | |
| return fns | |
| def __len__(self): | |
| return self.npairs # size should correspond to the number of pairs, not images | |
| def __repr__(self): | |
| res = "Dataset: %s\n" % self.__class__.__name__ | |
| res += " %d images," % self.nimg | |
| res += " %d image pairs" % self.npairs | |
| res += "\n root: %s...\n" % self.root | |
| return res | |
| def _flow2png(flow, path): | |
| flow = np.clip(np.around(16 * flow), -(2**15), 2**15 - 1) | |
| bytes = np.int16(flow).view(np.uint8) | |
| Image.fromarray(bytes).save(path) | |
| return flow / 16 | |
| def _png2flow(path): | |
| try: | |
| flow = np.asarray(Image.open(path)).view(np.int16) | |
| return np.float32(flow) / 16 | |
| except: | |
| raise IOError("Error loading flow for %s" % path) | |
| class StillPairDataset(PairDataset): | |
| """A dataset of 'still' image pairs. | |
| By overloading a normal image dataset, it appends the get_pair(i) function | |
| that serves trivial image pairs (img1, img2) where img1 == img2 == get_image(i). | |
| """ | |
| def get_pair(self, pair_idx, output=()): | |
| if isinstance(output, str): | |
| output = output.split() | |
| img1, img2 = map(self.get_image, self.image_pairs[pair_idx]) | |
| W, H = img1.size | |
| sx = img2.size[0] / float(W) | |
| sy = img2.size[1] / float(H) | |
| meta = {} | |
| if "aflow" in output or "flow" in output: | |
| mgrid = np.mgrid[0:H, 0:W][::-1].transpose(1, 2, 0).astype(np.float32) | |
| meta["aflow"] = mgrid * (sx, sy) | |
| meta["flow"] = meta["aflow"] - mgrid | |
| if "mask" in output: | |
| meta["mask"] = np.ones((H, W), np.uint8) | |
| if "homography" in output: | |
| meta["homography"] = np.diag(np.float32([sx, sy, 1])) | |
| return img1, img2, meta | |
| class SyntheticPairDataset(PairDataset): | |
| """A synthetic generator of image pairs. | |
| Given a normal image dataset, it constructs pairs using random homographies & noise. | |
| """ | |
| def __init__(self, dataset, scale="", distort=""): | |
| self.attach_dataset(dataset) | |
| self.distort = instanciate_transformation(distort) | |
| self.scale = instanciate_transformation(scale) | |
| def attach_dataset(self, dataset): | |
| assert isinstance(dataset, Dataset) and not isinstance(dataset, PairDataset) | |
| self.dataset = dataset | |
| self.npairs = dataset.nimg | |
| self.get_image = dataset.get_image | |
| self.get_key = dataset.get_key | |
| self.get_filename = dataset.get_filename | |
| self.root = None | |
| def make_pair(self, img): | |
| return img, img | |
| def get_pair(self, i, output=("aflow")): | |
| """Procedure: | |
| This function applies a series of random transformations to one original image | |
| to form a synthetic image pairs with perfect ground-truth. | |
| """ | |
| if isinstance(output, str): | |
| output = output.split() | |
| original_img = self.dataset.get_image(i) | |
| scaled_image = self.scale(original_img) | |
| scaled_image, scaled_image2 = self.make_pair(scaled_image) | |
| scaled_and_distorted_image = self.distort( | |
| dict(img=scaled_image2, persp=(1, 0, 0, 0, 1, 0, 0, 0)) | |
| ) | |
| W, H = scaled_image.size | |
| trf = scaled_and_distorted_image["persp"] | |
| meta = dict() | |
| if "aflow" in output or "flow" in output: | |
| # compute optical flow | |
| xy = np.mgrid[0:H, 0:W][::-1].reshape(2, H * W).T | |
| aflow = np.float32(persp_apply(trf, xy).reshape(H, W, 2)) | |
| meta["flow"] = aflow - xy.reshape(H, W, 2) | |
| meta["aflow"] = aflow | |
| if "homography" in output: | |
| meta["homography"] = np.float32(trf + (1,)).reshape(3, 3) | |
| return scaled_image, scaled_and_distorted_image["img"], meta | |
| def __repr__(self): | |
| res = "Dataset: %s\n" % self.__class__.__name__ | |
| res += " %d images and pairs" % self.npairs | |
| res += "\n root: %s..." % self.dataset.root | |
| res += "\n Scale: %s" % (repr(self.scale).replace("\n", "")) | |
| res += "\n Distort: %s" % (repr(self.distort).replace("\n", "")) | |
| return res + "\n" | |
| class TransformedPairs(PairDataset): | |
| """Automatic data augmentation for pre-existing image pairs. | |
| Given an image pair dataset, it generates synthetically jittered pairs | |
| using random transformations (e.g. homographies & noise). | |
| """ | |
| def __init__(self, dataset, trf=""): | |
| self.attach_dataset(dataset) | |
| self.trf = instanciate_transformation(trf) | |
| def attach_dataset(self, dataset): | |
| assert isinstance(dataset, PairDataset) | |
| self.dataset = dataset | |
| self.nimg = dataset.nimg | |
| self.npairs = dataset.npairs | |
| self.get_image = dataset.get_image | |
| self.get_key = dataset.get_key | |
| self.get_filename = dataset.get_filename | |
| self.root = None | |
| def get_pair(self, i, output=""): | |
| """Procedure: | |
| This function applies a series of random transformations to one original image | |
| to form a synthetic image pairs with perfect ground-truth. | |
| """ | |
| img_a, img_b_, metadata = self.dataset.get_pair(i, output) | |
| img_b = self.trf({"img": img_b_, "persp": (1, 0, 0, 0, 1, 0, 0, 0)}) | |
| trf = img_b["persp"] | |
| if "aflow" in metadata or "flow" in metadata: | |
| aflow = metadata["aflow"] | |
| aflow[:] = persp_apply(trf, aflow.reshape(-1, 2)).reshape(aflow.shape) | |
| W, H = img_a.size | |
| flow = metadata["flow"] | |
| mgrid = np.mgrid[0:H, 0:W][::-1].transpose(1, 2, 0).astype(np.float32) | |
| flow[:] = aflow - mgrid | |
| if "corres" in metadata: | |
| corres = metadata["corres"] | |
| corres[:, 1] = persp_apply(trf, corres[:, 1]) | |
| if "homography" in metadata: | |
| # p_b = homography * p_a | |
| trf_ = np.float32(trf + (1,)).reshape(3, 3) | |
| metadata["homography"] = np.float32(trf_ @ metadata["homography"]) | |
| return img_a, img_b["img"], metadata | |
| def __repr__(self): | |
| res = "Transformed Pairs from %s\n" % type(self.dataset).__name__ | |
| res += " %d images and pairs" % self.npairs | |
| res += "\n root: %s..." % self.dataset.root | |
| res += "\n transform: %s" % (repr(self.trf).replace("\n", "")) | |
| return res + "\n" | |
| class CatPairDataset(CatDataset): | |
| """Concatenation of several pair datasets.""" | |
| def __init__(self, *datasets): | |
| CatDataset.__init__(self, *datasets) | |
| pair_offsets = [0] | |
| for db in datasets: | |
| pair_offsets.append(db.npairs) | |
| self.pair_offsets = np.cumsum(pair_offsets) | |
| self.npairs = self.pair_offsets[-1] | |
| def __len__(self): | |
| return self.npairs | |
| def __repr__(self): | |
| fmt_str = "CatPairDataset(" | |
| for db in self.datasets: | |
| fmt_str += str(db).replace("\n", " ") + ", " | |
| return fmt_str[:-2] + ")" | |
| def pair_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 pair_call(self, func, i, *args, **kwargs): | |
| b, j = self.pair_which(i) | |
| return getattr(self.datasets[b], func)(j, *args, **kwargs) | |
| def get_pair(self, i, output=()): | |
| b, i = self.pair_which(i) | |
| return self.datasets[b].get_pair(i, output) | |
| def get_flow_filename(self, pair_idx, *args, **kwargs): | |
| return self.pair_call("get_flow_filename", pair_idx, *args, **kwargs) | |
| def get_mask_filename(self, pair_idx, *args, **kwargs): | |
| return self.pair_call("get_mask_filename", pair_idx, *args, **kwargs) | |
| def get_corres_filename(self, pair_idx, *args, **kwargs): | |
| return self.pair_call("get_corres_filename", pair_idx, *args, **kwargs) | |
| def is_pair(x): | |
| if isinstance(x, (tuple, list)) and len(x) == 2: | |
| return True | |
| if isinstance(x, np.ndarray) and x.ndim == 1 and x.shape[0] == 2: | |
| return True | |
| return False | |