Spaces:
Runtime error
Runtime error
| # TRI-VIDAR - Copyright 2022 Toyota Research Institute. All rights reserved. | |
| import os | |
| from glob import glob | |
| import numpy as np | |
| from vidar.utils.data import make_list | |
| from vidar.utils.types import is_list | |
| class FolderTree: | |
| """ | |
| Creates a dataset tree folder structure for file loading. | |
| Parameters | |
| ---------- | |
| path : String | |
| Path where dataset is stored | |
| prefix : String | |
| Optional prefix for each filename | |
| suffix : String | |
| Optional prefix for each filename | |
| sub_folders : list[String] | |
| Optional list of sub_folders located inside each data folder where data is stored | |
| deep : Int | |
| How deep in the folder structure we should go | |
| single_folder : Bool | |
| Whether the folder structure should be ignored, and a single folder is used | |
| nested : bool | |
| If true, go one folder deeper to find scenes | |
| stride: Int | |
| Stride for context generation | |
| context : Tuple | |
| Which context should be used | |
| """ | |
| def __init__(self, path, prefix='', suffix='', sub_folders=('',), deep=1, | |
| single_folder=False, nested=False, stride=1, context=()): | |
| # Store context information | |
| self.context = list(context) | |
| if 0 not in self.context: | |
| self.context.append(0) | |
| self.num_context = 0 if len(self.context) == 0 else max(self.context) - min(self.context) | |
| self.with_context = self.num_context > 0 | |
| self.min_context = 0 if not self.with_context else min(self.context) | |
| self.stride = stride | |
| self.pad_numbers = False | |
| # Initialize empty folder tree | |
| self.folder_tree = [] | |
| # If we are providing a file list, treat each line as a scene | |
| if is_list(path): | |
| self.folder_tree = [[file] for file in path] | |
| # If we are providing a folder | |
| else: | |
| # Get folders | |
| string = '*' + '/*' * (deep - 1) | |
| folders = glob(os.path.join(path, string)) | |
| folders.sort() | |
| # If nesting, go one folder deeper in order to find the scenes | |
| if nested: | |
| upd_folders = [] | |
| for folder in folders: | |
| new_folders = glob(os.path.join(folder, '*')) | |
| upd_folders.extend(new_folders) | |
| folders = upd_folders | |
| folders.sort() | |
| if single_folder: | |
| # Use current folder as the only one | |
| self.folder_tree.append(folders) | |
| else: | |
| # Populate folder tree | |
| for folder in folders: | |
| # For each sub-folder | |
| for sub_folder in make_list(sub_folders): | |
| # Get and sort files in each folder | |
| files = glob(os.path.join(folder, sub_folder, '{}*{}'.format(prefix, suffix))) | |
| if self.pad_numbers: | |
| for i in range(len(files)): | |
| pref, suf = files[i].split('/')[:-1], files[i].split('/')[-1] | |
| num, ext = suf.split('.') | |
| files[i] = '/'.join(pref) + ('/%010d' % int(num)) + '.' + ext | |
| files.sort() | |
| if self.pad_numbers: | |
| for i in range(len(files)): | |
| pref, suf = files[i].split('/')[:-1], files[i].split('/')[-1] | |
| num, ext = suf.split('.') | |
| files[i] = '/'.join(pref) + ('/%d' % int(num)) + '.' + ext | |
| if self.stride > 1: | |
| files = files[::self.stride] | |
| # Only store if there are more images than context | |
| if len(files) > self.num_context: | |
| self.folder_tree.append(files) | |
| # Get size of each folder | |
| self.slices = [len(folder) for folder in self.folder_tree] | |
| # Compensate for context size | |
| if self.with_context: | |
| self.slices = [s - self.num_context for s in self.slices] | |
| # Create cumulative size and get total | |
| self.slices = [0] + list(np.cumsum(self.slices)) | |
| self.total = self.slices[-1] | |
| def __len__(self): | |
| """Dataset size""" | |
| return self.total | |
| def get_idxs(self, idx): | |
| """Get folder and file indexes given dataset index""" | |
| idx1 = np.searchsorted(self.slices, idx, side='right') - 1 | |
| idx2 = idx - self.slices[idx1] | |
| return idx1, idx2 | |
| def get_item(self, idx, return_loc=False): | |
| """Return filename item given index""" | |
| idx1, idx2 = self.get_idxs(idx) | |
| item = {0: self.folder_tree[idx1][idx2 - self.min_context]} | |
| if return_loc: | |
| return item, idx2 - self.min_context | |
| else: | |
| return item | |
| def get_context(self, idx): | |
| """Return forward context given index.""" | |
| idx1, idx2 = self.get_idxs(idx) | |
| return {ctx: self.folder_tree[idx1][idx2 - self.min_context + ctx] for ctx in self.context} | |