Spaces:
Runtime error
Runtime error
| # TRI-VIDAR - Copyright 2022 Toyota Research Institute. All rights reserved. | |
| import json | |
| import os | |
| import random | |
| from collections import OrderedDict | |
| import numpy as np | |
| import torch | |
| from vidar.utils.decorators import iterate1 | |
| from vidar.utils.types import is_seq, is_tensor, is_dict, is_int | |
| def stack_sample(sample, lidar_sample=None, radar_sample=None): | |
| """ | |
| Stack samples from multiple cameras | |
| Parameters | |
| ---------- | |
| sample : list[Dict] | |
| List of camera samples | |
| lidar_sample : list[Dict] | |
| List of lidar samples | |
| radar_sample : list[Dict] | |
| List of radar samples | |
| Returns | |
| ------- | |
| stacked_sample: Dict | |
| Stacked sample | |
| """ | |
| # If there are no tensors, return empty list | |
| if len(sample) == 0: | |
| return None | |
| # If there is only one sensor don't do anything | |
| if len(sample) == 1: | |
| sample = sample[0] | |
| return sample | |
| # Otherwise, stack sample | |
| first_sample = sample[0] | |
| stacked_sample = {} | |
| for key, val in first_sample.items(): | |
| # Global keys (do not stack) | |
| if key in ['idx', 'dataset_idx']: | |
| stacked_sample[key] = first_sample[key] | |
| # Meta keys | |
| elif key in ['meta']: | |
| stacked_sample[key] = {} | |
| for key2 in first_sample[key].keys(): | |
| stacked_sample[key][key2] = {} | |
| for key3 in first_sample[key][key2].keys(): | |
| stacked_sample[key][key2][key3] = torch.stack( | |
| [torch.tensor(s[key][key2][key3]) for s in sample], 0) | |
| # Stack tensors | |
| elif is_tensor(val): | |
| stacked_sample[key] = torch.stack([s[key] for s in sample], 0) | |
| # Stack list | |
| elif is_seq(first_sample[key]): | |
| stacked_sample[key] = [] | |
| # Stack list of torch tensors | |
| if is_tensor(first_sample[key][0]): | |
| for i in range(len(first_sample[key])): | |
| stacked_sample[key].append( | |
| torch.stack([s[key][i] for s in sample], 0)) | |
| else: | |
| stacked_sample[key] = [s[key] for s in sample] | |
| # Repeat for dictionaries | |
| elif is_dict(first_sample[key]): | |
| stacked_sample[key] = stack_sample([s[key] for s in sample]) | |
| # Append lists | |
| else: | |
| stacked_sample[key] = [s[key] for s in sample] | |
| # Return stacked sample | |
| return stacked_sample | |
| def merge_sample(samples): | |
| """Merge information from multiple samples""" | |
| merged_sample = {} | |
| for sample in samples: | |
| for key, val in sample.items(): | |
| if key not in merged_sample: | |
| merged_sample[key] = val | |
| else: | |
| merged_sample[key] = merge_sample([merged_sample[key], val]) | |
| return merged_sample | |
| def parse_crop(cfg, shape): | |
| """Parse crop information to generate borders""" | |
| borders = None | |
| if cfg.has('crop_borders'): | |
| borders = parse_crop_borders(cfg.crop_borders, shape) | |
| if cfg.has('crop_random'): | |
| if borders is None: | |
| borders = [0, 0, shape[1], shape[0]] | |
| borders = parse_crop_random(borders, cfg.crop_random) | |
| return borders | |
| def parse_crop_borders(borders, shape): | |
| """ | |
| Calculate borders for cropping. | |
| Parameters | |
| ---------- | |
| borders : Tuple | |
| Border input for parsing. Can be one of the following forms: | |
| (int, int, int, int): y, height, x, width | |
| (int, int): y, x --> y, height = image_height - y, x, width = image_width - x | |
| Negative numbers are taken from image borders, according to the shape argument | |
| Float numbers for y and x are treated as percentage, according to the shape argument, | |
| and in this case height and width are centered at that point. | |
| shape : Tuple | |
| Image shape (image_height, image_width), used to determine negative crop boundaries | |
| Returns | |
| ------- | |
| borders : Tuple | |
| Parsed borders for cropping (left, top, right, bottom) | |
| """ | |
| # Return full image if there are no borders to crop | |
| if len(borders) == 0: | |
| return 0, 0, shape[1], shape[0] | |
| # Copy borders for modification | |
| borders = list(borders).copy() | |
| # If borders are 4-dimensional | |
| if len(borders) == 4: | |
| borders = [borders[2], borders[0], borders[3], borders[1]] | |
| if is_int(borders[0]): | |
| # If horizontal cropping is integer (regular cropping) | |
| borders[0] += shape[1] if borders[0] < 0 else 0 | |
| borders[2] += shape[1] if borders[2] <= 0 else borders[0] | |
| else: | |
| # If horizontal cropping is float (center cropping) | |
| center_w, half_w = borders[0] * shape[1], borders[2] / 2 | |
| borders[0] = int(center_w - half_w) | |
| borders[2] = int(center_w + half_w) | |
| if is_int(borders[1]): | |
| # If vertical cropping is integer (regular cropping) | |
| borders[1] += shape[0] if borders[1] < 0 else 0 | |
| borders[3] += shape[0] if borders[3] <= 0 else borders[1] | |
| else: | |
| # If vertical cropping is float (center cropping) | |
| center_h, half_h = borders[1] * shape[0], borders[3] / 2 | |
| borders[1] = int(center_h - half_h) | |
| borders[3] = int(center_h + half_h) | |
| # If borders are 2-dimensional | |
| elif len(borders) == 2: | |
| borders = [borders[1], borders[0]] | |
| if is_int(borders[0]): | |
| # If cropping is integer (regular cropping) | |
| borders = (max(0, borders[0]), | |
| max(0, borders[1]), | |
| shape[1] + min(0, borders[0]), | |
| shape[0] + min(0, borders[1])) | |
| else: | |
| # If cropping is float (center cropping) | |
| center_w, half_w = borders[0] * shape[1], borders[1] / 2 | |
| center_h, half_h = borders[0] * shape[0], borders[1] / 2 | |
| borders = (int(center_w - half_w), int(center_h - half_h), | |
| int(center_w + half_w), int(center_h + half_h)) | |
| # Otherwise, invalid | |
| else: | |
| raise NotImplementedError('Crop tuple must have 2 or 4 values.') | |
| # Assert that borders are valid | |
| assert 0 <= borders[0] < borders[2] <= shape[1] and \ | |
| 0 <= borders[1] < borders[3] <= shape[0], 'Crop borders {} are invalid'.format(borders) | |
| # Return updated borders | |
| return borders | |
| def parse_crop_random(borders, shape): | |
| """ | |
| Create borders for random cropping. | |
| Crops are generated anywhere in the image inside the borders | |
| Parameters | |
| ---------- | |
| borders : Tuple | |
| Area of the image where random cropping can happen (left, top, right, bottom) | |
| shape : Tuple | |
| Cropped output shape (height, width) | |
| Returns | |
| ------- | |
| borders : tuple | |
| Parsed borders for cropping (left, top, right, bottom) | |
| """ | |
| # Return full borders if there is no random crop | |
| if len(shape) == 0: | |
| return borders | |
| # Check if random crop is valid | |
| assert 0 < shape[1] <= borders[2] - borders[0] and \ | |
| 0 < shape[0] <= borders[3] - borders[1], 'Random crop must be smaller than the image' | |
| # Sample a crop | |
| x = random.randint(borders[0], borders[2] - shape[1]) | |
| y = random.randint(borders[1], borders[3] - shape[0]) | |
| # Return new borders | |
| return x, y, x + shape[1], y + shape[0] | |
| def invert_pose(pose): | |
| """ | |
| Inverts a transformation matrix (pose) | |
| Parameters | |
| ---------- | |
| pose : np.Array | |
| Input pose [4, 4] | |
| Returns | |
| ------- | |
| inv_pose : np.Array | |
| Inverted pose [4, 4] | |
| """ | |
| inv_pose = np.eye(4) | |
| inv_pose[:3, :3] = np.transpose(pose[:3, :3]) | |
| inv_pose[:3, -1] = - inv_pose[:3, :3] @ pose[:3, -1] | |
| # return np.linalg.inv(pose) | |
| return inv_pose | |
| def make_relative_pose(samples): | |
| """ | |
| Convert sample poses to relative frane of reference (based on the first target frame) | |
| Parameters | |
| ---------- | |
| samples : list[Dict] | |
| Input samples | |
| Returns | |
| ------- | |
| samples : list[Dict] | |
| Relative samples | |
| """ | |
| # Do nothing if there is no pose | |
| if 'pose' not in samples[0]: | |
| return samples | |
| # Get inverse current poses | |
| inv_pose = [invert_pose(samples[i]['pose'][0]) for i in range(len(samples))] | |
| # For each camera | |
| for i in range(len(samples)): | |
| # For each context | |
| for j in samples[0]['pose'].keys(): | |
| if j == 0: | |
| if i > 0: | |
| samples[i]['pose'][j] = \ | |
| samples[i]['pose'][j] @ inv_pose[0] | |
| else: | |
| samples[i]['pose'][j] = \ | |
| samples[i]['pose'][j] @ inv_pose[i] | |
| return samples | |
| def dummy_intrinsics(image): | |
| """ | |
| Return dummy intrinsics calculated based on image resolution | |
| Parameters | |
| ---------- | |
| image : PIL Image | |
| Image from which intrinsics will be calculated | |
| Returns | |
| ------- | |
| intrinsics : np.Array | |
| Image intrinsics (fx = cx = w/2, fy = cy = h/2) [3,3] | |
| """ | |
| w, h = [float(d) for d in image.size] | |
| return np.array([[w/2, 0., w/2. - 0.5], | |
| [0., h/2, h/2. - 0.5], | |
| [0., 0., 1.]]) | |
| def load_ontology(name, filter_list=None): | |
| """Loads ontology from file and optionally filters it""" | |
| filename = 'vidar/ontologies/{}.json'.format(name) | |
| if os.path.exists(filename): | |
| ontology = json.load(open(filename, 'r')) | |
| if filter_list is not None and len(filter_list) > 0: | |
| ontology = filter_ontology(ontology, filter_list) | |
| return ontology | |
| else: | |
| return None | |
| def save_ontology(ontology, name): | |
| """Save ontology to a JSON file""" | |
| if is_seq(ontology): | |
| ontology = ontology[0] | |
| for key in ontology.keys(): | |
| ontology[key]['color'] = ontology[key]['color'].tolist() | |
| json.dump(ontology, open('ontologies/{}.json'.format(name), 'w')) | |
| def filter_ontology(ontology, values): | |
| """Filter ontology to remove certain classes""" | |
| new_ontology = OrderedDict() | |
| for i, val in enumerate(values[1:]): | |
| new_ontology[i] = ontology[str(val)] | |
| return new_ontology | |
| def convert_ontology(semantic_id, ontology_convert): | |
| """Convert from one ontology to another""" | |
| if ontology_convert is None: | |
| return semantic_id | |
| else: | |
| semantic_id_convert = semantic_id.copy() | |
| for key, val in ontology_convert.items(): | |
| semantic_id_convert[semantic_id == key] = val | |
| return semantic_id_convert | |
| def initialize_ontology(base, ontology): | |
| """Initialize ontology and conversion table if necessary""" | |
| return load_ontology(base), None | |