Spaces:
Runtime error
Runtime error
File size: 5,529 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 |
# TRI-VIDAR - Copyright 2022 Toyota Research Institute. All rights reserved.
import torch
from pytorch3d.transforms.rotation_conversions import \
matrix_to_euler_angles, euler_angles_to_matrix
from vidar.utils.data import keys_in
from vidar.utils.decorators import iterate1, iterate12
from vidar.utils.types import is_tensor, is_list, is_seq
def flip_lr_fn(tensor):
"""Function to flip a tensor from left to right"""
return torch.flip(tensor, [-1])
def flip_flow_lr_fn(flow):
"""Function to flip a flow tensor from left to right"""
flow_flip = torch.flip(flow, [3])
flow_flip[:, :1, :, :] *= -1
return flow_flip.contiguous()
def flip_intrinsics_lr_fn(K, shape):
"""Function to flip a 3x3 intrinsic matrix from left to right"""
K = K.clone()
K[:, 0, 2] = shape[-1] - K[:, 0, 2]
return K
def flip_pose_lr_fn(T):
"""Function to flip a 4x4 transformation matrix from left to right"""
rot = T[:, :3, :3]
axis = matrix_to_euler_angles(rot, convention='XYZ')
axis[:, [1, 2]] = axis[:, [1, 2]] * -1
rot = euler_angles_to_matrix(axis, convention='XYZ')
T[:, :3, :3] = rot
T[:, 0, -1] = - T[:, 0, -1]
return T
@iterate1
def flip_lr(tensor, flip=True):
"""Flip a tensor from left to right"""
# Not flipping option
if not flip:
return tensor
# If it's a list, repeat
if is_list(tensor):
return [flip_lr(t) for t in tensor]
# Return flipped tensor
if tensor.dim() == 5:
return torch.stack([flip_lr_fn(tensor[:, i])
for i in range(tensor.shape[1])], 1)
else:
return flip_lr_fn(tensor)
@iterate1
def flip_flow_lr(flow, flip=True):
"""Flip a flow tensor from left to right"""
# Not flipping option
if not flip:
return flow
# If it's a list, repeat
if is_list(flow):
return [flip_flow_lr(f) for f in flow]
# Flip flow and invert first dimension
if flow.dim() == 5:
return torch.stack([flip_flow_lr_fn(flow[:, i])
for i in range(flow.shape[1])], 1)
else:
return flip_flow_lr_fn(flow)
@iterate12
def flip_intrinsics_lr(K, shape, flip=True):
"""Flip a 3x3 camera intrinsic matrix from left to right"""
# Not flipping option
if not flip:
return K
# If shape is a tensor, use it's dimensions
if is_tensor(shape):
shape = shape.shape
# Flip horizontal information (first row)
if K.dim() == 4:
return torch.stack([flip_intrinsics_lr_fn(K[:, i], shape)
for i in range(K.shape[1])], 1)
else:
return flip_intrinsics_lr_fn(K, shape)
def flip_pose_lr(pose, flip=True):
"""Flip a 4x4 transformation matrix from left to right"""
# Not flipping option
if not flip:
return pose
# Repeat for all pose keys
for key in pose.keys():
# Get pose key
if key == 0:
if pose[key].dim() == 3:
continue
elif pose[key].dim() == 4:
T = pose[key][:, 1:].clone()
else:
raise ValueError('Invalid pose dimension')
else:
T = pose[key].clone()
# Flip pose
if T.dim() == 4:
T = torch.stack([flip_pose_lr_fn(T[:, i])
for i in range(T.shape[1])], 1)
else:
T = flip_pose_lr_fn(T)
# Store flipped value back
if key == 0:
pose[key][:, 1:] = T
else:
pose[key] = T
# Return flipped pose
return pose
def flip_batch(batch, flip=True):
"""Flip a batch from left to right"""
# Not flipping option
if not flip:
return batch
# If it's a list, repeat
if is_seq(batch):
return [flip_batch(b) for b in batch]
# Flip batch
flipped_batch = {}
# Keys to not flip
for key in keys_in(batch, ['idx', 'filename', 'splitname']):
flipped_batch[key] = batch[key]
# Tensor flipping
for key in keys_in(batch, ['rgb', 'mask', 'input_depth', 'depth', 'semantic']):
flipped_batch[key] = flip_lr(batch[key])
# Intrinsics flipping
for key in keys_in(batch, ['intrinsics']):
flipped_batch[key] = flip_intrinsics_lr(batch[key], batch['rgb'])
# Pose flipping
for key in keys_in(batch, ['pose']):
flipped_batch[key] = flip_pose_lr(batch[key])
return flipped_batch
def flip_predictions(predictions, flip=True):
"""Flip predictions from left to right"""
# Not flipping option
if not flip:
return predictions
# Flip predictions
flipped_predictions = {}
for key in predictions.keys():
if key.startswith('depth'):
flipped_predictions[key] = flip_lr(predictions[key])
if key.startswith('pose'):
flipped_predictions[key] = flip_pose_lr(predictions[key])
# Return flipped predictions
return flipped_predictions
def flip_output(output, flip=True):
"""Flip output from left to right"""
# Not flipping option
if not flip:
return output
# If it's a list, repeat
if is_seq(output):
return [flip_output(b) for b in output]
# Flip output
flipped_output = {}
# Do not flip loss and metrics
for key in keys_in(output, ['loss', 'metrics']):
flipped_output[key] = output[key]
# Flip predictions
flipped_output['predictions'] = flip_predictions(output['predictions'])
# Return flipped output
return flipped_output
|