|
import torch |
|
import cv2 |
|
import numpy as np |
|
|
|
from torchvision import transforms |
|
from torch.utils.data import Dataset |
|
from pathlib import Path |
|
|
|
|
|
class PatchesDataset(Dataset): |
|
""" |
|
HPatches dataset class. |
|
# Note: output_shape = (output_width, output_height) |
|
# Note: this returns Pytorch tensors, resized to output_shape (if specified) |
|
# Note: the homography will be adjusted according to output_shape. |
|
|
|
Parameters |
|
---------- |
|
root_dir : str |
|
Path to the dataset |
|
use_color : bool |
|
Return color images or convert to grayscale. |
|
data_transform : Function |
|
Transformations applied to the sample |
|
output_shape: tuple |
|
If specified, the images and homographies will be resized to the desired shape. |
|
type: str |
|
Dataset subset to return from ['i', 'v', 'all']: |
|
i - illumination sequences |
|
v - viewpoint sequences |
|
all - all sequences |
|
""" |
|
def __init__(self, root_dir, use_color=True, data_transform=None, output_shape=None, type='all'): |
|
super().__init__() |
|
self.type = type |
|
self.root_dir = root_dir |
|
self.data_transform = data_transform |
|
self.output_shape = output_shape |
|
self.use_color = use_color |
|
base_path = Path(root_dir) |
|
folder_paths = [x for x in base_path.iterdir() if x.is_dir()] |
|
image_paths = [] |
|
warped_image_paths = [] |
|
homographies = [] |
|
for path in folder_paths: |
|
if self.type == 'i' and path.stem[0] != 'i': |
|
continue |
|
if self.type == 'v' and path.stem[0] != 'v': |
|
continue |
|
num_images = 5 |
|
file_ext = '.ppm' |
|
for i in range(2, 2 + num_images): |
|
image_paths.append(str(Path(path, "1" + file_ext))) |
|
warped_image_paths.append(str(Path(path, str(i) + file_ext))) |
|
homographies.append(np.loadtxt(str(Path(path, "H_1_" + str(i))))) |
|
self.files = {'image_paths': image_paths, 'warped_image_paths': warped_image_paths, 'homography': homographies} |
|
|
|
def scale_homography(self, homography, original_scale, new_scale, pre): |
|
scales = np.divide(new_scale, original_scale) |
|
if pre: |
|
s = np.diag(np.append(scales, 1.)) |
|
homography = np.matmul(s, homography) |
|
else: |
|
sinv = np.diag(np.append(1. / scales, 1.)) |
|
homography = np.matmul(homography, sinv) |
|
return homography |
|
|
|
def __len__(self): |
|
return len(self.files['image_paths']) |
|
|
|
def __getitem__(self, idx): |
|
|
|
def _read_image(path): |
|
img = cv2.imread(path, cv2.IMREAD_COLOR) |
|
if self.use_color: |
|
return img |
|
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) |
|
return gray |
|
|
|
image = _read_image(self.files['image_paths'][idx]) |
|
|
|
warped_image = _read_image(self.files['warped_image_paths'][idx]) |
|
homography = np.array(self.files['homography'][idx]) |
|
sample = {'image': image, 'warped_image': warped_image, 'homography': homography, 'index' : idx} |
|
|
|
|
|
if self.output_shape is not None: |
|
sample['homography'] = self.scale_homography(sample['homography'], |
|
sample['image'].shape[:2][::-1], |
|
self.output_shape, |
|
pre=False) |
|
sample['homography'] = self.scale_homography(sample['homography'], |
|
sample['warped_image'].shape[:2][::-1], |
|
self.output_shape, |
|
pre=True) |
|
|
|
for key in ['image', 'warped_image']: |
|
sample[key] = cv2.resize(sample[key], self.output_shape) |
|
if self.use_color is False: |
|
sample[key] = np.expand_dims(sample[key], axis=2) |
|
|
|
transform = transforms.ToTensor() |
|
|
|
for key in ['image', 'warped_image']: |
|
sample[key] = transform(sample[key]).type('torch.FloatTensor') |
|
return sample |
|
|
