def generate_confmaps(example):
"""Local processing function for dataset mapping."""
# Pull out first instance as (n_nodes, 2) tensor.
example["points"] = tf.gather(example["instances"], 0, axis=0)
# Generate confidence maps.
example["confidence_maps"] = make_confmaps(example["points"],
xv=xv,
yv=yv,
sigma=self.sigma *
self.output_stride)
if self.with_offsets:
example["offsets"] = mask_offsets(
make_offsets(example["points"],
xv,
yv,
stride=self.output_stride),
example["confidence_maps"],
self.offsets_threshold,
)
if self.flatten_offsets:
shape = tf.shape(example["offsets"])
example["offsets"] = tf.reshape(example["offsets"],
[shape[0], shape[1], -1])
return example
def test_mask_offsets():
points = np.array([[1.0, 1.0]], "float32")
xv, yv = sleap.nn.data.confidence_maps.make_grid_vectors(4,
4,
output_stride=1)
off = offset_regression.make_offsets(points, xv, yv, stride=1)
cm = sleap.nn.data.confidence_maps.make_confmaps(points, xv, yv, sigma=1)
off_mask = offset_regression.mask_offsets(off, cm, threshold=0.2)
np.testing.assert_array_equal(off_mask[:3, :3], off[:3, :3])
np.testing.assert_array_equal(off_mask[3:, :], 0.0)
np.testing.assert_array_equal(off_mask[:, 3:], 0.0)
def generate_confmaps(example):
"""Local processing function for dataset mapping."""
example["instance_confidence_maps"] = make_confmaps(
example["center_instance"],
xv=xv,
yv=yv,
sigma=self.sigma * self.output_stride,
)
if self.with_offsets:
example["offsets"] = mask_offsets(
make_offsets(example["center_instance"],
xv,
yv,
stride=self.output_stride),
example["instance_confidence_maps"],
self.offsets_threshold,
)
if self.flatten_offsets:
shape = tf.shape(example["offsets"])
example["offsets"] = tf.reshape(example["offsets"],
[shape[0], shape[1], -1])
if self.all_instances:
if self.with_offsets:
cms, offsets = make_multi_confmaps_with_offsets(
example["all_instances"],
xv,
yv,
self.output_stride,
self.sigma * self.output_stride,
offsets_threshold=self.offsets_threshold,
)
example["all_instance_offsets"] = offsets
if self.flatten_offsets:
shape = tf.shape(example["all_instance_offsets"])
example["all_instance_offsets"] = tf.reshape(
example["all_instance_offsets"],
[shape[0], shape[1], -1])
else:
cms = make_multi_confmaps(
example["all_instances"],
xv=xv,
yv=yv,
sigma=self.sigma * self.output_stride,
)
example["all_instance_confidence_maps"] = cms
return example
def test_make_offsets():
points = np.array([[1.8, 2.3], [0.4, 3.1], [np.nan, np.nan]], "float32")
xv, yv = sleap.nn.data.confidence_maps.make_grid_vectors(4,
4,
output_stride=1)
off = offset_regression.make_offsets(points, xv, yv, stride=1)
assert off.shape == (4, 4, 3, 2)
XX, YY = tf.meshgrid(xv, yv)
np.testing.assert_allclose(XX + off[..., 0, 0], 1.8, atol=1e-6)
np.testing.assert_allclose(YY + off[..., 0, 1], 2.3, atol=1e-6)
np.testing.assert_allclose(XX + off[..., 1, 0], 0.4, atol=1e-6)
np.testing.assert_allclose(YY + off[..., 1, 1], 3.1, atol=1e-6)
np.testing.assert_array_equal(off[..., 2, :], 0)
def make_multi_confmaps_with_offsets(
instances: tf.Tensor,
xv: tf.Tensor,
yv: tf.Tensor,
stride: int,
sigma: float,
offsets_threshold: float,
) -> tf.Tensor:
"""Make confidence maps and offsets for multiple instances through reduction.
Args:
instances: A tensor of shape `(n_instances, n_nodes, 2)` and dtype `tf.float32`
containing instance points where the last axis corresponds to (x, y) pixel
coordinates on the image. This must be rank-3 even if a single instance is
present.
xv: Sampling grid vector for x-coordinates of shape `(grid_width,)` and dtype
`tf.float32`. This can be generated by
`sleap.nn.data.utils.make_grid_vectors`.
yv: Sampling grid vector for y-coordinates of shape `(grid_height,)` and dtype
`tf.float32`. This can be generated by
`sleap.nn.data.utils.make_grid_vectors`.
stride: Scaling factor for offset coordinates. The individual offset vectors
will be divided by this value. Useful for adjusting for strided sampling
grids so that the offsets point to the smaller grid coordinates.
sigma: Standard deviation of the 2D Gaussian distribution sampled to generate
confidence maps.
offsets_threshold: Minimum confidence map value below which offsets will be
replaced with zeros.
flatten_offsets: If `True`, the last two channels of the offset maps will be
flattened to produce rank-3 tensors. If `False`, the generated offset maps
will be rank-4 with shaape `(height, width, n_nodes, 2)`.
Returns:
A tuple of `(confmaps, offsets)`.
`confmaps` are confidence maps as a tensor of shape
`(grid_height, grid_width, n_nodes)` and dtype `tf.float32`.
Each channel will contain the elementwise maximum of the confidence maps
generated from all individual points for the associated node.
`offsets` are offset maps as a `tf.Tensor` of shape
`(grid_height, grid_width, n_nodes, 2)` and dtype `tf.float32`. The last axis
corresponds to the x- and y-offsets at each grid point for each node.
Notes:
The confidence maps and offsets are computed individually for each instance
and immediately max-reduced to avoid maintaining the entire set of all instance
maps. This enables memory-efficient generation of multi-instance maps for
examples with large numbers of instances.
See also: sleap.nn.data.make_grid_vectors, make_confmaps, make_multi_confmaps
"""
# Initialize output tensors.
grid_height = tf.shape(yv)[0]
grid_width = tf.shape(xv)[0]
n_nodes = tf.shape(instances)[1]
cms = tf.zeros((grid_height, grid_width, n_nodes), tf.float32)
offsets = tf.zeros((grid_height, grid_width, n_nodes, 2), tf.float32)
# Eliminate instances completely outside of image.
in_img = (instances > 0) & (instances < tf.reshape(
tf.stack([xv[-1], yv[-1]], axis=0), [1, 1, 2]))
in_img = tf.reduce_any(tf.reduce_all(in_img, axis=-1), axis=1)
in_img = tf.ensure_shape(in_img, [None])
instances = tf.boolean_mask(instances, in_img)
# Generate and reduce outputs by instance.
for points in instances:
cms_instance = make_confmaps(points, xv, yv, sigma=sigma)
cms = tf.maximum(cms, cms_instance)
offsets_instance = mask_offsets(
make_offsets(points, xv, yv, stride=stride),
cms_instance,
threshold=offsets_threshold,
)
offsets += offsets_instance
return cms, offsets