regression masks from instance segmentation masks

tf_instance_regression.py

def _next_instance(mask, iid, instances):
  """Process single instance and add it to the mask
  Args:
    mask: source mask
    iid: instance id
    instances: instance segmentation mask
  Returns:
    updated mask
  """
  yx = tf.to_int32(tf.where(tf.equal(instances, iid)))
  y0, y1 = tf.reduce_min(yx[:,0]), tf.reduce_max(yx[:,0])
  x0, x1 = tf.reduce_min(yx[:,1]), tf.reduce_max(yx[:,1])
  # TODO: find more efficient way?
  dy0 = tf.scatter_nd(yx, (yx[:,0] - y0), image_size)
  dx0 = tf.scatter_nd(yx, (yx[:,1] - x0), image_size)
  dy1 = tf.scatter_nd(yx, (y1 - yx[:,0]), image_size)
  dx1 = tf.scatter_nd(yx, (x1 - yx[:,1]), image_size)
  return mask + tf.stack([dy0, dx0, dy1, dx1], axis=2)

def regression_masks(instances, cid, parallel_iterations=1):
  """Get regression masks for the given class id
    instances: [B,H,W] tensor of tf.uint16
    cid: class id
    parallel_iterations: number of parallel iterations
  Returns: 
    tensor [B,H,W,4] with regression displacements
  """
  def _reg_mask(instances):
    # TODO: make it for a list of classes?
    mask = (instances >= cid*1000) & (instances < (cid+1)*1000)
    ids, idx = tf.unique(tf.boolean_mask(instances, mask))
    reg_mask = tf.zeros(image_size + (4,), dtype=tf.int32)
    reg_mask = tf.foldl(lambda mask, iid: _next_instance(mask, iid, instances), 
                        ids, reg_mask,
                        parallel_iterations=parallel_iterations,
                        back_prop=False)
    return reg_mask
  return tf.map_fn(_reg_mask, instances, parallel_iterations=1, back_prop=False)