File size: 10,631 Bytes
fc16538
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
# 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]


@iterate1
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