File size: 1,282 Bytes
404d2af 8b973ee 404d2af 8b973ee 404d2af 8b973ee |
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 AachenDataset(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": "AachenDayNight",
}
|