1+
2+ import re
3+ from collections import defaultdict
4+ import os
5+
6+ from PIL import Image
7+ from torch .utils .data import Dataset
8+ import numpy as np
9+ from packnet_sfm .utils .image import load_image
10+
11+ ########################################################################################################################
12+ #### FUNCTIONS
13+ ########################################################################################################################
14+
15+ def dummy_calibration (image ):
16+ w , h = [float (d ) for d in image .size ]
17+ return np .array ([[1000. , 0. , w / 2. - 0.5 ],
18+ [0. , 1000. , h / 2. - 0.5 ],
19+ [0. , 0. , 1. ]])
20+
21+ def get_idx (filename ):
22+ return int (re .search (r'\d+' , filename ).group ())
23+
24+ ########################################################################################################################
25+ #### DATASET
26+ ########################################################################################################################
27+
28+ class EUROCDataset (Dataset ):
29+ def __init__ (self , root_dir , split , data_transform = None ,
30+ forward_context = 0 , back_context = 0 , strides = (1 ,),
31+ depth_type = None , cameras = [], ** kwargs ):
32+ super ().__init__ ()
33+ # Asserts
34+ # assert depth_type is None or depth_type == '', \
35+ # 'ImageDataset currently does not support depth types'
36+ assert len (strides ) == 1 and strides [0 ] == 1 , \
37+ 'ImageDataset currently only supports stride of 1.'
38+
39+ self .root_dir = root_dir
40+ self .split = split
41+
42+ self .cameras = cameras
43+
44+ self .backward_context = back_context
45+ self .forward_context = forward_context
46+ self .has_context = self .backward_context + self .forward_context > 0
47+ self .strides = 1
48+
49+ self .files = []
50+ self .file_tree = defaultdict (list )
51+ self .read_files (root_dir )
52+
53+ self .depth_type = depth_type
54+ self .with_depth = depth_type is not '' and depth_type is not None
55+
56+ # print('file tree')
57+ # print(self.file_tree)
58+
59+ for k , v in self .file_tree .items ():
60+ file_set = set (self .file_tree [k ])
61+ files = [fname for fname in sorted (v ) if self ._has_context (fname , file_set )]
62+ # files = [fname for fname in sorted(v)]
63+ self .files .extend ([[k , fname ] for fname in files ])
64+
65+ self .data_transform = data_transform
66+
67+ def read_files (self , directory , ext = ('.png' , '.jpg' , '.jpeg' ), skip_empty = True ):
68+ files = defaultdict (list )
69+ for entry in os .scandir (directory ):
70+ relpath = os .path .relpath (entry .path , directory )
71+ if entry .is_dir ():
72+ d_files = self .read_files (entry .path , ext = ext , skip_empty = skip_empty )
73+ if skip_empty and not len (d_files ):
74+ continue
75+ # files[relpath] = d_files[entry.path]
76+ self .file_tree [entry .path ] = d_files [entry .path ]
77+ elif entry .is_file ():
78+ if ext is None or entry .path .lower ().endswith (tuple (ext )):
79+ files [directory ].append (relpath )
80+ return files
81+
82+ def __len__ (self ):
83+ return len (self .files )
84+
85+ def _change_idx (self , idx , filename ):
86+ _ , ext = os .path .splitext (os .path .basename (filename ))
87+ return self .split .format (idx ) + ext
88+
89+ def _has_context (self , filename , file_set ):
90+ context_paths = self ._get_context_file_paths (filename , file_set )
91+
92+ return all ([f in file_set for f in context_paths ])
93+
94+ def _get_context_file_paths (self , filename , file_set ):
95+ fidx = get_idx (filename )
96+ # a hacky way to deal with two different strides in euroc dataset
97+ idxs = [- self .backward_context , - self .forward_context , self .backward_context , self .forward_context ]
98+ potential_files = [self ._change_idx (fidx + i , filename ) for i in idxs ]
99+ valid_paths = [fname for fname in potential_files if fname in file_set ]
100+
101+ # return [self._change_idx(fidx + i, filename) for i in idxs]
102+ return valid_paths
103+
104+ def _read_rgb_context_files (self , session , filename ):
105+ file_set = set (self .file_tree [session ])
106+ context_paths = self ._get_context_file_paths (filename , file_set )
107+
108+ return [self ._read_rgb_file (session , filename ) for filename in context_paths ]
109+
110+
111+ def _read_rgb_file (self , session , filename ):
112+
113+ gray_image = load_image (os .path .join (self .root_dir , session , filename ))
114+ gray_image_np = np .array (gray_image )
115+ rgb_image_np = np .stack ([gray_image_np for _ in range (3 )], axis = 2 )
116+ rgb_image = Image .fromarray (rgb_image_np )
117+
118+ return rgb_image
119+
120+ def _read_npy_depth (self , session , depth_filename ):
121+ depth_file_path = os .path .join (self .root_dir , session , '../../depth_maps' , depth_filename )
122+ return np .load (depth_file_path )
123+
124+ def _read_depth (self , session , depth_filename ):
125+ """Get the depth map from a file."""
126+ if self .depth_type in ['vicon' ]:
127+ return self ._read_npy_depth (session , depth_filename )
128+ else :
129+ raise NotImplementedError (
130+ 'Depth type {} not implemented' .format (self .depth_type ))
131+
132+ def _has_depth (self , session , depth_filename ):
133+ depth_file_path = os .path .join (self .root_dir , session , '../../depth_maps' , depth_filename )
134+
135+ return os .path .isfile (depth_file_path )
136+
137+ def __getitem__ (self , idx ):
138+ session , filename = self .files [idx ]
139+ image = self ._read_rgb_file (session , filename )
140+
141+ intrinsic_type = 'euroc' if len (self .cameras ) == 0 else 'euroc_{}' .format (self .cameras [0 ])
142+
143+ sample = {
144+ 'idx' : idx ,
145+ 'filename' : '%s_%s' % (session , os .path .splitext (filename )[0 ]),
146+ 'rgb' : image ,
147+ 'intrinsics' : dummy_calibration (image ),
148+ 'intrinsic_type' : intrinsic_type
149+ }
150+
151+ if self .has_context :
152+ sample ['rgb_context' ] = \
153+ self ._read_rgb_context_files (session , filename )
154+
155+ depth_filename = filename .split ('.' )[0 ] + 'depth.npy'
156+ if self .with_depth :
157+ if self ._has_depth (session , depth_filename ):
158+ sample ['depth' ] = self ._read_depth (session , depth_filename )
159+
160+
161+ if self .data_transform :
162+ sample = self .data_transform (sample )
163+
164+ return sample
165+
166+ ########################################################################################################################
167+ if __name__ == "__main__" :
168+ data_dir = '/data/datasets/euroc/V1_01_easy_has_depth/mav0'
169+ euroc_dataset = EUROCDataset (root_dir = data_dir ,split = '{:09}' ,depth_type = 'vicon' ,
170+ forward_context = 249999872 ,back_context = 250000128 ,
171+ cameras = ['cam0' ])
172+ print (len (euroc_dataset ))
173+ print (euroc_dataset [0 ].keys ())
174+ print (euroc_dataset [0 ]['rgb' ])
175+ print (euroc_dataset [0 ]['rgb_context' ])
176+ print (euroc_dataset [0 ]['intrinsics' ].shape )
177+ print (euroc_dataset [0 ]['intrinsic_type' ])
178+ print (euroc_dataset [0 ]['depth' ].shape )
0 commit comments