File size: 1,272 Bytes
a80d6bb
 
 
 
 
 
 
 
 
 
 
c74a070
a80d6bb
 
 
 
 
 
 
 
c74a070
a80d6bb
 
c74a070
 
 
 
 
 
 
 
 
 
 
 
 
 
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
import os
from torch.utils.data import Dataset

from src.utils.dataset import read_img_gray


class InLocDataset(Dataset):
    def __init__(self, img_path, match_list_path, img_resize=None, down_factor=16):
        self.img_path = img_path
        self.img_resize = img_resize
        self.down_factor = down_factor
        with open(match_list_path, "r") as f:
            self.raw_pairs = f.readlines()
        print("number of matching pairs: ", len(self.raw_pairs))

    def __len__(self):
        return len(self.raw_pairs)

    def __getitem__(self, idx):
        raw_pair = self.raw_pairs[idx]
        image_name0, image_name1 = raw_pair.strip("\n").split(" ")
        path_img0 = os.path.join(self.img_path, image_name0)
        path_img1 = os.path.join(self.img_path, image_name1)
        img0, scale0 = read_img_gray(
            path_img0, resize=self.img_resize, down_factor=self.down_factor
        )
        img1, scale1 = read_img_gray(
            path_img1, resize=self.img_resize, down_factor=self.down_factor
        )
        return {
            "image0": img0,
            "image1": img1,
            "scale0": scale0,
            "scale1": scale1,
            "pair_names": (image_name0, image_name1),
            "dataset_name": "InLoc",
        }