Spaces:
Running
Running
File size: 4,365 Bytes
499e141 |
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 |
# -*- coding: utf-8 -*-
# @Author : xuelun
import glob
import torch
import imagesize
import torch.nn.functional as F
from os.path import join
from torch.utils.data import Dataset
from datasets.utils import read_images
class RobotcarDataset(Dataset):
def __init__(self,
root_dir, # data root dit
npz_root, # data info, like, overlap, image_path, depth_path
seq_name, # current sequence
mode, # train or val or test
min_overlap_score,
max_overlap_score,
max_resize, # max edge after resize
df, # general is 8 for ResNet w/ pre 3-layers
padding, # padding image for batch training
augment_fn, # augmentation function
max_samples, # max sample in current sequence
**kwargs):
super().__init__()
self.root = join('zeb', seq_name)
paths = glob.glob(join(self.root, '*.txt'))
lines = []
for path in paths:
with open(path, 'r') as file:
scene_id = path.rpartition('/')[-1].rpartition('.')[0].split('_')[0]
line = file.readline().strip().split()
lines.append([scene_id] + line)
self.pairs = sorted(lines)
self.scale = 1 / df
self.df = df
self.max_resize = max_resize
self.padding = padding
def __len__(self):
return len(self.pairs)
def __getitem__(self, idx):
pair = self.pairs[idx]
scene_id = pair[0]
timestamp0 = pair[1]
timestamp1 = pair[2]
img_path0 = join(self.root, '{}_{}.png'.format(scene_id, timestamp0))
img_path1 = join(self.root, '{}_{}.png'.format(scene_id, timestamp1))
width0, height0 = imagesize.get(img_path0)
width1, height1 = imagesize.get(img_path1)
image0, color0, scale0, resize0, mask0 = read_images(
img_path0, self.max_resize, self.df, self.padding, None)
image1, color1, scale1, resize1, mask1 = read_images(
img_path1, self.max_resize, self.df, self.padding, None)
K0 = torch.tensor(list(map(float, pair[5:14])), dtype=torch.float).reshape(3, 3)
K1 = torch.tensor(list(map(float, pair[14:23])), dtype=torch.float).reshape(3, 3)
# read image size
imsize0 = torch.tensor([height0, width0], dtype=torch.long)
imsize1 = torch.tensor([height1, width1], dtype=torch.long)
resize0 = torch.tensor(resize0, dtype=torch.long)
resize1 = torch.tensor(resize1, dtype=torch.long)
T_0to1 = torch.tensor(list(map(float, pair[23:])), dtype=torch.float).reshape(4, 4)
data = {
# image 0
'image0': image0, # (1, 3, h, w)
'color0': color0, # (1, h, w)
'imsize0': imsize0, # (2) - 2:(h, w)
'resize0': resize0, # (2) - 2:(h, w)
# image 1
'image1': image1,
'color1': color1,
'imsize1': imsize1, # (2) - 2:[h, w]
'resize1': resize1, # (2) - 2:(h, w)
# image transform
'T_0to1': T_0to1, # (4, 4)
'K0': K0, # (3, 3)
'K1': K1,
# pair information
'scale0': scale0, # [scale_w, scale_h]
'scale1': scale1,
'dataset_name': 'Robotcar',
'scene_id': scene_id,
'pair_id': f'{idx}-{idx}',
'pair_names': (str(timestamp0),
str(timestamp1)),
'covisible0': float(pair[3]),
'covisible1': float(pair[4]),
}
if mask0 is not None: # img_padding is True
if self.scale:
# noinspection PyArgumentList
[ts_mask_0, ts_mask_1] = F.interpolate(torch.stack([mask0, mask1], dim=0)[None].float(),
scale_factor=self.scale,
mode='nearest',
recompute_scale_factor=False)[0].bool()
# noinspection PyUnboundLocalVariable
data.update({'mask0': ts_mask_0, 'mask1': ts_mask_1})
return data
|