Spaces:
Running
Running
File size: 4,211 Bytes
a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 |
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.0))
homography = np.matmul(s, homography)
else:
sinv = np.diag(np.append(1.0 / scales, 1.0))
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,
}
# Apply transformations
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
|