Spaces:
Running
Running
| 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 | |