def rotations(self):
rots = [self.matrix]
for i in range(3):
base = rotate90(rots[-1])
for p in permutations(range(1,4)):
rots.append(replace(base,p))
return rots
def forward(self, images):
"""images.shape = (batch, time, channels, obs_height, obs_width)
output.shape = (batch, time, orientations, map_height, map_width)"""
(batch_sz, seq_length, channels, obs_h, obs_w) = images.shape
assert obs_w == obs_h and obs_w % 2 == 1
map_sz = self.map_size
state = None
# merge time dimension into batch dimension, and run net
images_ = rearrange(images, 'b t e h w -> (b t) e h w')
obs = self.cnn(images_)
obs = rearrange(obs, '(b t) e h w -> b t e h w', t=seq_length)
assert obs.shape[-3] == self.embedding_size
view_range = (obs.shape[-1] - 1) // 2
# create stack of rotated observations. shape = (batch, time, orientations, embedding, height, width)
rotated_obs = t.stack([rotate90(obs, times=i) for i in range(self.orientations)], dim=2)
# initial pose heatmap, one-hot at center of map and first orientation
pose = t.zeros((batch_sz, self.orientations, map_sz, map_sz), device=images.device)
pose[:, 0, (map_sz - 1) // 2, (map_sz - 1) // 2] = 1
(raw_poses, softmax_poses, maps) = ([], [], []) # pose tensors and maps through time
for step in range(seq_length - 1):
# registration: deconvolve pose estimate with rotated observations to get map-sized embedding
# output shape = (batch, embeding, height, width)
registered_obs = self.register(pose, rotated_obs[:, step, :, :, :, :])
# aggregate registered observations into a map, updating the map state
(map, state) = self.aggregate_rnn(registered_obs, state)
# localization: convolve map with rotated observations (from next frame) to obtain pose estimate
pose = self.localize(map, rotated_obs[:, step + 1, :, :, :, :], pad=not self.improved_padding)
if self.temperature != 1.0:
pose = pose * self.temperature
# return poses before the softmax
raw_poses.append(pose)
# softmax over spatial/orientation dimensions
pose = self.softmax_pose(pose)
# option to apply padding after softmax (so edges are ignored)
if self.improved_padding: pose = F.pad(pose, [view_range] * 4)
# return poses after the softmax, and the map
softmax_poses.append(pose)
maps.append(map)
# aggregate poses (before softmax) over time and apply padding if needed
raw_poses = t.stack(raw_poses, dim=1)
if self.improved_padding:
raw_poses = F.pad(raw_poses, [view_range] * 4)
return {'raw_poses': raw_poses, 'softmax_poses': t.stack(softmax_poses, dim=1), 'maps': t.stack(maps, dim=1)}
def is_max(self):
p = self.current
n = to_int(p)
for i in range(3):
p = rotate90(p)
for q in permutations(range(1, 4)):
if len(p) == len(self.current) and to_int(replace(p, q)) > n:
return False
return True
def is_max(self):
p = self.current
n = to_int(p)
for i in range(3):
p = rotate90(p)
for q in permutations(range(1,4)):
if len(p) == len(self.current) and to_int(replace(p,q)) > n:
return False
return True
def update_visiblility(square_patch, visible_patch, angle_ranges, u, v):
new_angle_range = []
for ang_sm, ang_lg, side in angle_ranges:
# rotate so that ang_sm and ang_lg is in range (-pi/4, pi/4)
square_patch = rotate90(square_patch, -side)
visible_patch = rotate90(visible_patch, -side)
rotate = (pi / 2) * side
# update strip visibility
square_patch[-1, :], visible_patch[:, -1, :], strip_ang_ranges\
= eval_strip_visibility(square_patch[-1, :], visible_patch[:, -1, :], ang_sm - rotate, ang_lg - rotate, u, v)
# rotate back to original angle
new_angle_range += [(s + rotate, l + rotate, side) for s, l in strip_ang_ranges]
square_patch = rotate90(square_patch, side)
visible_patch = rotate90(visible_patch, side)
return visible_patch, new_angle_range
def embed_view(patch, env_shape, ang90, h0, w0):
"""Embed a local view in an environment at the given pose"""
patch = rotate90(patch, ang90)
assert len(env_shape) == 2
image = torch.zeros((*patch.shape[:-2], *env_shape), dtype=patch.dtype)
h_env, w_env = env_shape
h_patch, w_patch = patch.shape[-2:]
image[..., max(0, h0):h_patch + h0, max(0, w0):w_patch + w0]\
= patch[..., max(0, -h0):h_env - h0, max(0, -w0):w_env - w0]
return image
def extract_view(env, x, y, ang, view_range): """Extract a local view from an environment at the given pose""" # get coordinates of window to extract xs = t.arange(x - view_range, x + view_range + 1, dtype=t.long) ys = t.arange(y - view_range, y + view_range + 1, dtype=t.long) # get coordinate 0 instead of going out of bounds (h, w) = (env.shape[-2], env.shape[-1]) (invalid_xs, invalid_ys) = ((xs < 0) | (xs >= w), (ys < 0) | (ys >= h)) # coords outside the env xs[invalid_xs] = 0 ys[invalid_ys] = 0 # extract view, and set to 0 observations that were out of bounds #view = env[..., ys, xs] # not equivalent to view = env[..., y1:y2, x1:x2] view = env.index_select(dim=-2, index=ys).index_select(dim=-1, index=xs) view[..., :, invalid_xs] = 0 view[..., invalid_ys, :] = 0 # rotate back. note only 90 degrees rotations are allowed return rotate90(view, (-ang) % 4)
def extract_view(env, h, w, ang90, view_range):
"""Extract a local view from an environment at the given pose"""
# get coordinates of window to extract
hs = torch.arange(h - view_range, h + view_range + 1, dtype=torch.long)
ws = torch.arange(w - view_range, w + view_range + 1, dtype=torch.long)
# get coordinate 0 instead of going out of bounds
h_env, w_env = env.shape[-2:]
invalid_hs, invalid_ws = (
(hs < 0) | (hs >= h_env), (ws < 0) | (ws >= w_env)
) # coords outside the env
hs[invalid_hs] = 0
ws[invalid_ws] = 0
# extract view, and set to 0 observations that were out of bounds
# view = env[..., xs, ys] # not equivalent to view = env[..., y1:y2, x1:x2]
view = env.index_select(dim=-2, index=hs).index_select(dim=-1, index=ws)
view[..., invalid_hs, :] = 0
view[..., :, invalid_ws] = 0
# rotate. note only 90 degrees rotations are allowed
return rotate90(view, ang90)