Merge pull request #1375 from s-gupta/cmp

Implementation for Cognitive Mapping and Planning paper.

cognitive_mapping_and_planning/src/map_utils.py

+# Copyright 2016 The TensorFlow Authors All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Various function to compute the ground truth map for training etc.
+"""
+import copy
+import skimage.morphology
+import numpy as np
+import scipy.ndimage
+import matplotlib.pyplot as plt
+import PIL
+
+import src.utils as utils
+import cv2
+
+def _get_xy_bounding_box(vertex, padding):
+  """Returns the xy bounding box of the environment."""
+  min_ = np.floor(np.min(vertex[:, :2], axis=0) - padding).astype(np.int)
+  max_ = np.ceil(np.max(vertex[:, :2], axis=0) + padding).astype(np.int)
+  return min_, max_
+
+def _project_to_map(map, vertex, wt=None, ignore_points_outside_map=False):
+  """Projects points to map, returns how many points are present at each
+  location."""
+  num_points = np.zeros((map.size[1], map.size[0]))
+  vertex_ = vertex[:, :2] - map.origin
+  vertex_ = np.round(vertex_ / map.resolution).astype(np.int)
+  if ignore_points_outside_map:
+    good_ind = np.all(np.array([vertex_[:,1] >= 0, vertex_[:,1] < map.size[1],
+                                vertex_[:,0] >= 0, vertex_[:,0] < map.size[0]]),
+                      axis=0)
+    vertex_ = vertex_[good_ind, :]
+    if wt is not None:
+      wt = wt[good_ind, :]
+  if wt is None:
+    np.add.at(num_points, (vertex_[:, 1], vertex_[:, 0]), 1)
+  else:
+    assert(wt.shape[0] == vertex.shape[0]), \
+      'number of weights should be same as vertices.'
+    np.add.at(num_points, (vertex_[:, 1], vertex_[:, 0]), wt)
+  return num_points
+
+def make_map(padding, resolution, vertex=None, sc=1.):
+  """Returns a map structure."""
+  min_, max_ = _get_xy_bounding_box(vertex*sc, padding=padding)
+  sz = np.ceil((max_ - min_ + 1) / resolution).astype(np.int32)
+  max_ = min_ + sz * resolution - 1
+  map = utils.Foo(origin=min_, size=sz, max=max_, resolution=resolution,
+                  padding=padding)
+  return map
+
+def _fill_holes(img, thresh):
+  """Fills holes less than thresh area (assumes 4 connectivity when computing
+  hole area."""
+  l, n = scipy.ndimage.label(np.logical_not(img))
+  img_ = img == True
+  cnts = np.bincount(l.reshape(-1))
+  for i, cnt in enumerate(cnts):
+    if cnt < thresh:
+      l[l == i] = -1
+  img_[l == -1] = True
+  return img_
+
+def compute_traversibility(map, robot_base, robot_height, robot_radius,
+                           valid_min, valid_max, num_point_threshold, shapess,
+                           sc=100., n_samples_per_face=200):
+  """Returns a bit map with pixels that are traversible or not as long as the
+  robot center is inside this volume we are good colisions can be detected by
+  doing a line search on things, or walking from current location to final
+  location in the bitmap, or doing bwlabel on the traversibility map."""
+
+  tt = utils.Timer()
+  tt.tic()
+  num_obstcale_points = np.zeros((map.size[1], map.size[0]))
+  num_points = np.zeros((map.size[1], map.size[0]))
+
+  for i, shapes in enumerate(shapess):
+    for j in range(shapes.get_number_of_meshes()):
+      p, face_areas, face_idx = shapes.sample_points_on_face_of_shape(
+          j, n_samples_per_face, sc)
+      wt = face_areas[face_idx]/n_samples_per_face
+
+      ind = np.all(np.concatenate(
+        (p[:, [2]] > robot_base,
+         p[:, [2]] < robot_base + robot_height), axis=1),axis=1)
+      num_obstcale_points += _project_to_map(map, p[ind, :], wt[ind])
+
+      ind = np.all(np.concatenate(
+        (p[:, [2]] > valid_min,
+         p[:, [2]] < valid_max), axis=1),axis=1)
+      num_points += _project_to_map(map, p[ind, :], wt[ind])
+
+  selem = skimage.morphology.disk(robot_radius / map.resolution)
+  obstacle_free = skimage.morphology.binary_dilation(
+      _fill_holes(num_obstcale_points > num_point_threshold, 20), selem) != True
+  valid_space = _fill_holes(num_points > num_point_threshold, 20)
+  traversible = np.all(np.concatenate((obstacle_free[...,np.newaxis],
+                                       valid_space[...,np.newaxis]), axis=2),
+                       axis=2)
+  # plt.imshow(np.concatenate((obstacle_free, valid_space, traversible), axis=1))
+  # plt.show()
+
+  map_out = copy.deepcopy(map)
+  map_out.num_obstcale_points = num_obstcale_points
+  map_out.num_points = num_points
+  map_out.traversible = traversible
+  map_out.obstacle_free = obstacle_free
+  map_out.valid_space = valid_space
+  tt.toc(log_at=1, log_str='src.map_utils.compute_traversibility: ')
+  return map_out
+
+
+def resize_maps(map, map_scales, resize_method):
+  scaled_maps = []
+  for i, sc in enumerate(map_scales):
+    if resize_method == 'antialiasing':
+      # Resize using open cv so that we can compute the size.
+      # Use PIL resize to use anti aliasing feature.
+      map_ = cv2.resize(map*1, None, None, fx=sc, fy=sc, interpolation=cv2.INTER_LINEAR)
+      w = map_.shape[1]; h = map_.shape[0]
+
+      map_img = PIL.Image.fromarray((map*255).astype(np.uint8))
+      map__img = map_img.resize((w,h), PIL.Image.ANTIALIAS)
+      map_ = np.asarray(map__img).astype(np.float32)
+      map_ = map_/255.
+      map_ = np.minimum(map_, 1.0)
+      map_ = np.maximum(map_, 0.0)
+    elif resize_method == 'linear_noantialiasing':
+      map_ = cv2.resize(map*1, None, None, fx=sc, fy=sc, interpolation=cv2.INTER_LINEAR)
+    else:
+      logging.error('Unknown resizing method')
+    scaled_maps.append(map_)
+  return scaled_maps
+
+
+def pick_largest_cc(traversible):
+  out = scipy.ndimage.label(traversible)[0]
+  cnt = np.bincount(out.reshape(-1))[1:]
+  return out == np.argmax(cnt) + 1
+
+def get_graph_origin_loc(rng, traversible):
+  """Erode the traversibility mask so that we get points in the bulk of the
+  graph, and not end up with a situation where the graph is localized in the
+  corner of a cramped room. Output Locs is in the coordinate frame of the
+  map."""
+
+  aa = pick_largest_cc(skimage.morphology.binary_erosion(traversible == True,
+                                                         selem=np.ones((15,15))))
+  y, x = np.where(aa > 0)
+  ind = rng.choice(y.size)
+  locs = np.array([x[ind], y[ind]])
+  locs = locs + rng.rand(*(locs.shape)) - 0.5
+  return locs
+
+
+def generate_egocentric_maps(scaled_maps, map_scales, map_crop_sizes, loc,
+                             x_axis, y_axis, theta):
+  maps = []
+  for i, (map_, sc, map_crop_size) in enumerate(zip(scaled_maps, map_scales, map_crop_sizes)):
+    maps_i = np.array(get_map_to_predict(loc*sc, x_axis, y_axis, map_,
+                                         map_crop_size,
+                                         interpolation=cv2.INTER_LINEAR)[0])
+    maps_i[np.isnan(maps_i)] = 0
+    maps.append(maps_i)
+  return maps
+
+def generate_goal_images(map_scales, map_crop_sizes, n_ori, goal_dist,
+                         goal_theta, rel_goal_orientation):
+  goal_dist = goal_dist[:,0]
+  goal_theta = goal_theta[:,0]
+  rel_goal_orientation = rel_goal_orientation[:,0]
+
+  goals = [];
+  # Generate the map images.
+  for i, (sc, map_crop_size) in enumerate(zip(map_scales, map_crop_sizes)):
+    goal_i = np.zeros((goal_dist.shape[0], map_crop_size, map_crop_size, n_ori),
+                      dtype=np.float32)
+    x = goal_dist*np.cos(goal_theta)*sc + (map_crop_size-1.)/2.
+    y = goal_dist*np.sin(goal_theta)*sc + (map_crop_size-1.)/2.
+
+    for j in range(goal_dist.shape[0]):
+      gc = rel_goal_orientation[j]
+      x0 = np.floor(x[j]).astype(np.int32); x1 = x0 + 1;
+      y0 = np.floor(y[j]).astype(np.int32); y1 = y0 + 1;
+      if x0 >= 0 and x0 <= map_crop_size-1:
+        if y0 >= 0 and y0 <= map_crop_size-1:
+          goal_i[j, y0, x0, gc] = (x1-x[j])*(y1-y[j])
+        if y1 >= 0 and y1 <= map_crop_size-1:
+          goal_i[j, y1, x0, gc] = (x1-x[j])*(y[j]-y0)
+
+      if x1 >= 0 and x1 <= map_crop_size-1:
+        if y0 >= 0 and y0 <= map_crop_size-1:
+          goal_i[j, y0, x1, gc] = (x[j]-x0)*(y1-y[j])
+        if y1 >= 0 and y1 <= map_crop_size-1:
+          goal_i[j, y1, x1, gc] = (x[j]-x0)*(y[j]-y0)
+
+    goals.append(goal_i)
+  return goals
+
+def get_map_to_predict(src_locs, src_x_axiss, src_y_axiss, map, map_size,
+                       interpolation=cv2.INTER_LINEAR):
+  fss = []
+  valids = []
+
+  center = (map_size-1.0)/2.0
+  dst_theta = np.pi/2.0
+  dst_loc = np.array([center, center])
+  dst_x_axis = np.array([np.cos(dst_theta), np.sin(dst_theta)])
+  dst_y_axis = np.array([np.cos(dst_theta+np.pi/2), np.sin(dst_theta+np.pi/2)])
+
+  def compute_points(center, x_axis, y_axis):
+    points = np.zeros((3,2),dtype=np.float32)
+    points[0,:] = center
+    points[1,:] = center + x_axis
+    points[2,:] = center + y_axis
+    return points
+
+  dst_points = compute_points(dst_loc, dst_x_axis, dst_y_axis)
+  for i in range(src_locs.shape[0]):
+    src_loc = src_locs[i,:]
+    src_x_axis = src_x_axiss[i,:]
+    src_y_axis = src_y_axiss[i,:]
+    src_points = compute_points(src_loc, src_x_axis, src_y_axis)
+    M = cv2.getAffineTransform(src_points, dst_points)
+
+    fs = cv2.warpAffine(map, M, (map_size, map_size), None, flags=interpolation,
+                        borderValue=np.NaN)
+    valid = np.invert(np.isnan(fs))
+    valids.append(valid)
+    fss.append(fs)
+  return fss, valids
+