Skip to content

Instantly share code, notes, and snippets.

@psycharo-zz
Last active February 3, 2017 13:07
Show Gist options
  • Save psycharo-zz/218f6c9a715c1bc4771b3444a79c3c98 to your computer and use it in GitHub Desktop.
Save psycharo-zz/218f6c9a715c1bc4771b3444a79c3c98 to your computer and use it in GitHub Desktop.
regression masks from instance segmentation masks
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)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment