File size: 7,927 Bytes
d4cbafd | 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 | import torch
import numpy as np
from dataset.homography_warper import get_rotation_matrix2d, warp_affine_crop
class Map(object):
def __init__(self, data, homography, description=None):
self.data = data
self.homography = homography
self.description = description
def as_image(self):
raise NotImplementedError
def get_cropped_maps(self, world_pts, patch_size, rotation=None, device='cpu'):
raise NotImplementedError
def to_map_points(self, scene_pts):
raise NotImplementedError
class GeometricMap(Map):
"""
A Geometric Map is a int tensor of shape [layers, x, y]. The homography must transform a point in scene
coordinates to the respective point in map coordinates.
:param data: Numpy array of shape [layers, x, y]
:param homography: Numpy array of shape [3, 3]
"""
def __init__(self, data, homography, description=None):
#assert isinstance(data.dtype, np.floating), "Geometric Maps must be float values."
super(GeometricMap, self).__init__(data, homography, description=description)
self._last_padding = None
self._last_padded_map = None
self._torch_map = None
def torch_map(self, device):
if self._torch_map is not None:
return self._torch_map
self._torch_map = torch.tensor(self.data, dtype=torch.uint8, device=device)
return self._torch_map
def as_image(self):
# We have to transpose x and y to rows and columns. Assumes origin is lower left for image
# Also we move the channels to the last dimension
return (np.transpose(self.data, (2, 1, 0))).astype(np.uint)
def get_padded_map(self, padding_x, padding_y, device):
if self._last_padding == (padding_x, padding_y):
return self._last_padded_map
else:
self._last_padding = (padding_x, padding_y)
self._last_padded_map = torch.full((self.data.shape[0],
self.data.shape[1] + 2 * padding_x,
self.data.shape[2] + 2 * padding_y),
False, dtype=torch.uint8)
self._last_padded_map[..., padding_x:-padding_x, padding_y:-padding_y] = self.torch_map(device)
return self._last_padded_map
@staticmethod
def batch_rotate(map_batched, centers, angles, out_height, out_width):
"""
As the input is a map and the warp_affine works on an image coordinate system we would have to
flip the y axis updown, negate the angles, and flip it back after transformation.
This, however, is the same as not flipping at and not negating the radian.
:param map_batched:
:param centers:
:param angles:
:param out_height:
:param out_width:
:return:
"""
M = get_rotation_matrix2d(centers, angles, torch.ones_like(angles))
rotated_map_batched = warp_affine_crop(map_batched, centers, M,
dsize=(out_height, out_width), padding_mode='zeros')
return rotated_map_batched
@classmethod
def get_cropped_maps_from_scene_map_batch(cls, maps, scene_pts, patch_size, rotation=None, device='cpu'):
"""
Returns rotated patches of each map around the transformed scene points.
___________________
| | |
| |ps[3] |
| | |
| | |
| o|__________|
| | ps[2] |
| | |
|_______|__________|
ps = patch_size
:param maps: List of GeometricMap objects [bs]
:param scene_pts: Scene points: [bs, 2]
:param patch_size: Extracted Patch size after rotation: [-x, -y, +x, +y]
:param rotation: Rotations in degrees: [bs]
:param device: Device on which the rotated tensors should be returned.
:return: Rotated and cropped tensor patches.
"""
batch_size = scene_pts.shape[0]
lat_size = 2 * np.max((patch_size[0], patch_size[2]))
long_size = 2 * np.max((patch_size[1], patch_size[3]))
assert lat_size % 2 == 0, "Patch width must be divisible by 2"
assert long_size % 2 == 0, "Patch length must be divisible by 2"
lat_size_half = lat_size // 2
long_size_half = long_size // 2
context_padding_x = int(np.ceil(np.sqrt(2) * lat_size))
context_padding_y = int(np.ceil(np.sqrt(2) * long_size))
centers = torch.tensor([s_map.to_map_points(scene_pts[np.newaxis, i]) for i, s_map in enumerate(maps)],
dtype=torch.long, device=device).squeeze(dim=1) \
+ torch.tensor([context_padding_x, context_padding_y], device=device, dtype=torch.long)
padded_map = [s_map.get_padded_map(context_padding_x, context_padding_y, device=device) for s_map in maps]
padded_map_batched = torch.stack([padded_map[i][...,
centers[i, 0] - context_padding_x: centers[i, 0] + context_padding_x,
centers[i, 1] - context_padding_y: centers[i, 1] + context_padding_y]
for i in range(centers.shape[0])], dim=0)
center_patches = torch.tensor([[context_padding_y, context_padding_x]],
dtype=torch.int,
device=device).repeat(batch_size, 1)
if rotation is not None:
angles = torch.Tensor(rotation)
else:
angles = torch.zeros(batch_size)
rotated_map_batched = cls.batch_rotate(padded_map_batched/255.,
center_patches.float(),
angles,
long_size,
lat_size)
del padded_map_batched
return rotated_map_batched[...,
long_size_half - patch_size[1]:(long_size_half + patch_size[3]),
lat_size_half - patch_size[0]:(lat_size_half + patch_size[2])]
def get_cropped_maps(self, scene_pts, patch_size, rotation=None, device='cpu'):
"""
Returns rotated patches of the map around the transformed scene points.
___________________
| | |
| |ps[3] |
| | |
| | |
| o|__________|
| | ps[2] |
| | |
|_______|__________|
ps = patch_size
:param scene_pts: Scene points: [bs, 2]
:param patch_size: Extracted Patch size after rotation: [-lat, -long, +lat, +long]
:param rotation: Rotations in degrees: [bs]
:param device: Device on which the rotated tensors should be returned.
:return: Rotated and cropped tensor patches.
"""
return self.get_cropped_maps_from_scene_map_batch([self]*scene_pts.shape[0], scene_pts,
patch_size, rotation=rotation, device=device)
def to_map_points(self, scene_pts):
org_shape = None
if len(scene_pts.shape) > 2:
org_shape = scene_pts.shape
scene_pts = scene_pts.reshape((-1, 2))
N, dims = scene_pts.shape
points_with_one = np.ones((dims + 1, N))
points_with_one[:dims] = scene_pts.T
map_points = (self.homography @ points_with_one).T[..., :dims]
if org_shape is not None:
map_points = map_points.reshape(org_shape)
return map_points
class ImageMap(Map): # TODO Implement for image maps -> watch flipped coordinate system
def __init__(self):
raise NotImplementedError |