# Copyright 2019-present NAVER Corp.
# CC BY-NC-SA 3.0
# Available only for non-commercial use
import pdb
from PIL import Image
import numpy as np
import torch
import torchvision.transforms as tvf
from .transforms import instanciate_transformation
from .transforms_tools import persp_apply
RGB_mean = [0.485, 0.456, 0.406]
RGB_std = [0.229, 0.224, 0.225]
norm_RGB = tvf.Compose([tvf.ToTensor(), tvf.Normalize(mean=RGB_mean, std=RGB_std)])
class PairLoader:
"""On-the-fly jittering of pairs of image with dense pixel ground-truth correspondences.
crop: random crop applied to both images
scale: random scaling applied to img2
distort: random ditorsion applied to img2
self[idx] returns a dictionary with keys: img1, img2, aflow, mask
- img1: cropped original
- img2: distorted cropped original
- aflow: 'absolute' optical flow = (x,y) position of each pixel from img1 in img2
- mask: (binary image) valid pixels of img1
"""
def __init__(
self,
dataset,
crop="",
scale="",
distort="",
norm=norm_RGB,
what="aflow mask",
idx_as_rng_seed=False,
):
assert hasattr(dataset, "npairs")
assert hasattr(dataset, "get_pair")
self.dataset = dataset
self.distort = instanciate_transformation(distort)
self.crop = instanciate_transformation(crop)
self.norm = instanciate_transformation(norm)
self.scale = instanciate_transformation(scale)
self.idx_as_rng_seed = idx_as_rng_seed # to remove randomness
self.what = what.split() if isinstance(what, str) else what
self.n_samples = 5 # number of random trials per image
def __len__(self):
assert len(self.dataset) == self.dataset.npairs, pdb.set_trace() # and not nimg
return len(self.dataset)
def __repr__(self):
fmt_str = "PairLoader\n"
fmt_str += repr(self.dataset)
fmt_str += " npairs: %d\n" % self.dataset.npairs
short_repr = (
lambda s: repr(s).strip().replace("\n", ", ")[14:-1].replace(" ", " ")
)
fmt_str += " Distort: %s\n" % short_repr(self.distort)
fmt_str += " Crop: %s\n" % short_repr(self.crop)
fmt_str += " Norm: %s\n" % short_repr(self.norm)
return fmt_str
def __getitem__(self, i):
# from time import time as now; t0 = now()
if self.idx_as_rng_seed:
import random
random.seed(i)
np.random.seed(i)
# Retrieve an image pair and their absolute flow
img_a, img_b, metadata = self.dataset.get_pair(i, self.what)
# aflow contains pixel coordinates indicating where each
# pixel from the left image ended up in the right image
# as (x,y) pairs, but its shape is (H,W,2)
aflow = np.float32(metadata["aflow"])
mask = metadata.get("mask", np.ones(aflow.shape[:2], np.uint8))
# apply transformations to the second image
img_b = {"img": img_b, "persp": (1, 0, 0, 0, 1, 0, 0, 0)}
if self.scale:
img_b = self.scale(img_b)
if self.distort:
img_b = self.distort(img_b)
# apply the same transformation to the flow
aflow[:] = persp_apply(img_b["persp"], aflow.reshape(-1, 2)).reshape(
aflow.shape
)
corres = None
if "corres" in metadata:
corres = np.float32(metadata["corres"])
corres[:, 1] = persp_apply(img_b["persp"], corres[:, 1])
# apply the same transformation to the homography
homography = None
if "homography" in metadata:
homography = np.float32(metadata["homography"])
# p_b = homography * p_a
persp = np.float32(img_b["persp"] + (1,)).reshape(3, 3)
homography = persp @ homography
# determine crop size
img_b = img_b["img"]
crop_size = self.crop({"imsize": (10000, 10000)})["imsize"]
output_size_a = min(img_a.size, crop_size)
output_size_b = min(img_b.size, crop_size)
img_a = np.array(img_a)
img_b = np.array(img_b)
ah, aw, p1 = img_a.shape
bh, bw, p2 = img_b.shape
assert p1 == 3
assert p2 == 3
assert aflow.shape == (ah, aw, 2)
assert mask.shape == (ah, aw)
# Let's start by computing the scale of the
# optical flow and applying a median filter:
dx = np.gradient(aflow[:, :, 0])
dy = np.gradient(aflow[:, :, 1])
scale = np.sqrt(np.clip(np.abs(dx[1] * dy[0] - dx[0] * dy[1]), 1e-16, 1e16))
accu2 = np.zeros((16, 16), bool)
Q = lambda x, w: np.int32(16 * (x - w.start) / (w.stop - w.start))
def window1(x, size, w):
l = x - int(0.5 + size / 2)
r = l + int(0.5 + size)
if l < 0:
l, r = (0, r - l)
if r > w:
l, r = (l + w - r, w)
if l < 0:
l, r = 0, w # larger than width
return slice(l, r)
def window(cx, cy, win_size, scale, img_shape):
return (
window1(cy, win_size[1] * scale, img_shape[0]),
window1(cx, win_size[0] * scale, img_shape[1]),
)
n_valid_pixel = mask.sum()
sample_w = mask / (1e-16 + n_valid_pixel)
def sample_valid_pixel():
n = np.random.choice(sample_w.size, p=sample_w.ravel())
y, x = np.unravel_index(n, sample_w.shape)
return x, y
# Find suitable left and right windows
trials = 0 # take the best out of few trials
best = -np.inf, None
for _ in range(50 * self.n_samples):
if trials >= self.n_samples:
break # finished!
# pick a random valid point from the first image
if n_valid_pixel == 0:
break
c1x, c1y = sample_valid_pixel()
# Find in which position the center of the left
# window ended up being placed in the right image
c2x, c2y = (aflow[c1y, c1x] + 0.5).astype(np.int32)
if not (0 <= c2x < bw and 0 <= c2y < bh):
continue
# Get the flow scale
sigma = scale[c1y, c1x]
# Determine sampling windows
if 0.2 < sigma < 1:
win1 = window(c1x, c1y, output_size_a, 1 / sigma, img_a.shape)
win2 = window(c2x, c2y, output_size_b, 1, img_b.shape)
elif 1 <= sigma < 5:
win1 = window(c1x, c1y, output_size_a, 1, img_a.shape)
win2 = window(c2x, c2y, output_size_b, sigma, img_b.shape)
else:
continue # bad scale
# compute a score based on the flow
x2, y2 = aflow[win1].reshape(-1, 2).T.astype(np.int32)
# Check the proportion of valid flow vectors
valid = (
(win2[1].start <= x2)
& (x2 < win2[1].stop)
& (win2[0].start <= y2)
& (y2 < win2[0].stop)
)
score1 = (valid * mask[win1].ravel()).mean()
# check the coverage of the second window
accu2[:] = False
accu2[Q(y2[valid], win2[0]), Q(x2[valid], win2[1])] = True
score2 = accu2.mean()
# Check how many hits we got
score = min(score1, score2)
trials += 1
if score > best[0]:
best = score, win1, win2
if None in best: # counldn't find a good window
img_a = np.zeros(output_size_a[::-1] + (3,), dtype=np.uint8)
img_b = np.zeros(output_size_b[::-1] + (3,), dtype=np.uint8)
aflow = np.nan * np.ones((2,) + output_size_a[::-1], dtype=np.float32)
homography = np.nan * np.ones((3, 3), dtype=np.float32)
else:
win1, win2 = best[1:]
img_a = img_a[win1]
img_b = img_b[win2]
aflow = aflow[win1] - np.float32([[[win2[1].start, win2[0].start]]])
mask = mask[win1]
aflow[~mask.view(bool)] = np.nan # mask bad pixels!
aflow = aflow.transpose(2, 0, 1) # --> (2,H,W)
if corres is not None:
corres[:, 0] -= (win1[1].start, win1[0].start)
corres[:, 1] -= (win2[1].start, win2[0].start)
if homography is not None:
trans1 = np.eye(3, dtype=np.float32)
trans1[:2, 2] = (win1[1].start, win1[0].start)
trans2 = np.eye(3, dtype=np.float32)
trans2[:2, 2] = (-win2[1].start, -win2[0].start)
homography = trans2 @ homography @ trans1
homography /= homography[2, 2]
# rescale if necessary
if img_a.shape[:2][::-1] != output_size_a:
sx, sy = (np.float32(output_size_a) - 1) / (
np.float32(img_a.shape[:2][::-1]) - 1
)
img_a = np.asarray(
Image.fromarray(img_a).resize(output_size_a, Image.ANTIALIAS)
)
mask = np.asarray(
Image.fromarray(mask).resize(output_size_a, Image.NEAREST)
)
afx = Image.fromarray(aflow[0]).resize(output_size_a, Image.NEAREST)
afy = Image.fromarray(aflow[1]).resize(output_size_a, Image.NEAREST)
aflow = np.stack((np.float32(afx), np.float32(afy)))
if corres is not None:
corres[:, 0] *= (sx, sy)
if homography is not None:
homography = homography @ np.diag(np.float32([1 / sx, 1 / sy, 1]))
homography /= homography[2, 2]
if img_b.shape[:2][::-1] != output_size_b:
sx, sy = (np.float32(output_size_b) - 1) / (
np.float32(img_b.shape[:2][::-1]) - 1
)
img_b = np.asarray(
Image.fromarray(img_b).resize(output_size_b, Image.ANTIALIAS)
)
aflow *= [[[sx]], [[sy]]]
if corres is not None:
corres[:, 1] *= (sx, sy)
if homography is not None:
homography = np.diag(np.float32([sx, sy, 1])) @ homography
homography /= homography[2, 2]
assert aflow.dtype == np.float32, pdb.set_trace()
assert homography is None or homography.dtype == np.float32, pdb.set_trace()
if "flow" in self.what:
H, W = img_a.shape[:2]
mgrid = np.mgrid[0:H, 0:W][::-1].astype(np.float32)
flow = aflow - mgrid
result = dict(img1=self.norm(img_a), img2=self.norm(img_b))
for what in self.what:
try:
result[what] = eval(what)
except NameError:
pass
return result
def threaded_loader(loader, iscuda, threads, batch_size=1, shuffle=True):
"""Get a data loader, given the dataset and some parameters.
Parameters
----------
loader : object[i] returns the i-th training example.
iscuda : bool
batch_size : int
threads : int
shuffle : int
Returns
-------
a multi-threaded pytorch loader.
"""
return torch.utils.data.DataLoader(
loader,
batch_size=batch_size,
shuffle=shuffle,
sampler=None,
num_workers=threads,
pin_memory=iscuda,
collate_fn=collate,
)
def collate(batch, _use_shared_memory=True):
"""Puts each data field into a tensor with outer dimension batch size.
Copied from https://github.com/pytorch in torch/utils/data/_utils/collate.py
"""
import re
error_msg = "batch must contain tensors, numbers, dicts or lists; found {}"
elem_type = type(batch[0])
if isinstance(batch[0], torch.Tensor):
out = None
if _use_shared_memory:
# If we're in a background process, concatenate directly into a
# shared memory tensor to avoid an extra copy
numel = sum([x.numel() for x in batch])
storage = batch[0].storage()._new_shared(numel)
out = batch[0].new(storage)
return torch.stack(batch, 0, out=out)
elif (
elem_type.__module__ == "numpy"
and elem_type.__name__ != "str_"
and elem_type.__name__ != "string_"
):
elem = batch[0]
assert elem_type.__name__ == "ndarray"
# array of string classes and object
if re.search("[SaUO]", elem.dtype.str) is not None:
raise TypeError(error_msg.format(elem.dtype))
batch = [torch.from_numpy(b) for b in batch]
try:
return torch.stack(batch, 0)
except RuntimeError:
return batch
elif batch[0] is None:
return list(batch)
elif isinstance(batch[0], int):
return torch.LongTensor(batch)
elif isinstance(batch[0], float):
return torch.DoubleTensor(batch)
elif isinstance(batch[0], str):
return batch
elif isinstance(batch[0], dict):
return {key: collate([d[key] for d in batch]) for key in batch[0]}
elif isinstance(batch[0], (tuple, list)):
transposed = zip(*batch)
return [collate(samples) for samples in transposed]
raise TypeError((error_msg.format(type(batch[0]))))
def tensor2img(tensor, model=None):
"""convert back a torch/numpy tensor to a PIL Image
by undoing the ToTensor() and Normalize() transforms.
"""
mean = norm_RGB.transforms[1].mean
std = norm_RGB.transforms[1].std
if isinstance(tensor, torch.Tensor):
tensor = tensor.detach().cpu().numpy()
res = np.uint8(np.clip(255 * ((tensor.transpose(1, 2, 0) * std) + mean), 0, 255))
from PIL import Image
return Image.fromarray(res)
if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser("Tool to debug/visualize the data loader")
parser.add_argument(
"dataloader", type=str, help="command to create the data loader"
)
args = parser.parse_args()
from datasets import *
auto_pairs = lambda db: SyntheticPairDataset(
db,
"RandomScale(256,1024,can_upscale=True)",
"RandomTilting(0.5), PixelNoise(25)",
)
loader = eval(args.dataloader)
print("Data loader =", loader)
from tools.viz import show_flow
for data in loader:
aflow = data["aflow"]
H, W = aflow.shape[-2:]
flow = (aflow - np.mgrid[:H, :W][::-1]).transpose(1, 2, 0)
show_flow(tensor2img(data["img1"]), tensor2img(data["img2"]), flow)