File size: 4,211 Bytes
404d2af
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
8b973ee
 
 
 
 
 
 
 
 
404d2af
 
 
 
 
 
 
 
 
 
 
 
8b973ee
404d2af
8b973ee
404d2af
 
8b973ee
404d2af
 
 
 
8b973ee
 
 
 
 
404d2af
 
 
 
8b973ee
404d2af
 
8b973ee
404d2af
 
 
 
8b973ee
404d2af
 
 
 
 
 
 
 
 
8b973ee
404d2af
8b973ee
 
 
 
 
 
 
 
404d2af
 
 
8b973ee
 
 
 
 
 
 
 
 
 
 
 
404d2af
8b973ee
404d2af
 
 
 
 
 
8b973ee
 
404d2af
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