def test_make_confmaps():
xv, yv = make_grid_vectors(image_height=4, image_width=5, output_stride=1)
points = tf.cast([[0.5, 1.], [3, 3.5], [2., 2.]], tf.float32)
cm = make_confmaps(points, xv, yv, sigma=1.)
assert cm.dtype == tf.float32
assert cm.shape == (4, 5, 3)
np.testing.assert_allclose(cm, [[
[0.535, 0., 0.018],
[0.535, 0., 0.082],
[0.197, 0.001, 0.135],
[0.027, 0.002, 0.082],
[0.001, 0.001, 0.018],
],
[
[0.882, 0., 0.082],
[0.882, 0.006, 0.368],
[0.325, 0.027, 0.607],
[0.044, 0.044, 0.368],
[0.002, 0.027, 0.082],
],
[
[0.535, 0.004, 0.135],
[0.535, 0.044, 0.607],
[0.197, 0.197, 1.],
[0.027, 0.325, 0.607],
[0.001, 0.197, 0.135],
],
[
[0.119, 0.01, 0.082],
[0.119, 0.119, 0.368],
[0.044, 0.535, 0.607],
[0.006, 0.882, 0.368],
[0., 0.535, 0.082],
]],
atol=1e-3)
# Grid aligned peak
points = tf.cast([[2, 3]], tf.float32)
cm = make_confmaps(points, xv, yv, sigma=1.)
assert cm.shape == (4, 5, 1)
assert cm[3, 2] == 1.0
# Output stride
xv, yv = make_grid_vectors(image_height=8, image_width=8, output_stride=2)
points = tf.cast([[2, 4]], tf.float32)
cm = make_confmaps(points, xv, yv, sigma=1.)
assert cm.shape == (4, 4, 1)
assert cm[2, 1] == 1.0
# Missing points
xv, yv = make_grid_vectors(image_height=8, image_width=8, output_stride=2)
points = tf.cast([[2, 4]], tf.float32)
cm = make_confmaps(points, xv, yv, sigma=1.)
points_with_nan = tf.cast([[2, 4], [np.nan, np.nan]], tf.float32)
cm_with_nan = make_confmaps(points_with_nan, xv, yv, sigma=1.)
assert cm_with_nan.shape == (4, 4, 2)
assert cm_with_nan.dtype == tf.float32
np.testing.assert_array_equal(cm_with_nan[:, :, 0], cm[:, :, 0])
assert (cm_with_nan[:, :, 1].numpy() == 0).all()
def test_make_pafs():
xv, yv = make_grid_vectors(image_height=3, image_width=3, output_stride=1)
edge_source = tf.cast([[1, 0.5], [0, 0]], tf.float32)
edge_destination = tf.cast([[1, 1.5], [2, 2]], tf.float32)
sigma = 1.0
pafs = edge_maps.make_pafs(
xv=xv,
yv=yv,
edge_source=edge_source,
edge_destination=edge_destination,
sigma=sigma,
)
np.testing.assert_allclose(
pafs,
[
[
[[0.0, 0.458], [0.707, 0.707]],
[[0.0, 0.969], [0.624, 0.624]],
[[0.0, 0.458], [0.096, 0.096]],
],
[
[[0.0, 0.607], [0.624, 0.624]],
[[0.0, 1.0], [0.707, 0.707]],
[[0.0, 0.607], [0.624, 0.624]],
],
[
[[0.0, 0.458], [0.096, 0.096]],
[[0.0, 0.969], [0.624, 0.624]],
[[0.0, 0.458], [0.707, 0.707]],
],
],
atol=1e-3,
)
def test_distance_to_edge():
xv, yv = make_grid_vectors(image_height=3, image_width=3, output_stride=1)
edge_source = tf.cast([[1, 0.5], [0, 0]], tf.float32)
edge_destination = tf.cast([[1, 1.5], [2, 2]], tf.float32)
sigma = 1.0
sampling_grid = tf.stack(tf.meshgrid(xv, yv), axis=-1) # (height, width, 2)
distances = edge_maps.distance_to_edge(sampling_grid,
edge_source=edge_source, edge_destination=edge_destination)
np.testing.assert_allclose(
distances,
[
[[1.25, 0. ],
[0.25, 0.5 ],
[1.25, 2. ]],
[[1. , 0.5 ],
[0. , 0. ],
[1. , 0.5 ]],
[[1.25, 2. ],
[0.25, 0.5 ],
[1.25, 0. ]]
],
atol=1e-3
)
def transform_dataset(self, input_ds: tf.data.Dataset) -> tf.data.Dataset:
"""Create a dataset that contains the generated confidence maps.
Args:
input_ds: A dataset with elements that contain the keys "instance_image",
"center_instance" and, if the attribute `all_instances` is True,
"all_instances".
Returns:
A `tf.data.Dataset` with the same keys as the input, as well as
"instance_confidence_maps" and, if the attribute `all_instances` is True,
"all_instance_confidence_maps" keys containing the generated confidence
maps.
Notes:
The output stride is relative to the current scale of the image. To map
points on the confidence maps to the raw image, first multiply them by the
output stride, and then scale the x- and y-coordinates by the "scale" key.
Importantly, the sigma will be proportional to the current image grid, not
the original grid prior to scaling operations.
"""
# Infer image dimensions to generate sampling grid.
test_example = next(iter(input_ds))
image_height = test_example["instance_image"].shape[0]
image_width = test_example["instance_image"].shape[1]
# Generate sampling grid vectors.
xv, yv = make_grid_vectors(
image_height=image_height,
image_width=image_width,
output_stride=self.output_stride,
)
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.all_instances:
example["all_instance_confidence_maps"] = make_multi_confmaps(
example["all_instances"],
xv=xv,
yv=yv,
sigma=self.sigma * self.output_stride,
)
return example
# Map transformation.
output_ds = input_ds.map(
generate_confmaps, num_parallel_calls=tf.data.experimental.AUTOTUNE
)
return output_ds
def test_make_grid_vector():
xv, yv = utils.make_grid_vectors(image_height=4,
image_width=3,
output_stride=1)
assert xv.dtype == tf.float32
assert xv.shape == (3, )
assert yv.dtype == tf.float32
assert yv.shape == (4, )
np.testing.assert_allclose(xv, [0, 1, 2])
np.testing.assert_allclose(yv, [0, 1, 2, 3])
xv, yv = utils.make_grid_vectors(image_height=4,
image_width=3,
output_stride=2)
np.testing.assert_allclose(xv, [0, 2])
np.testing.assert_allclose(yv, [0, 2])
def test_make_multi_pafs():
xv, yv = make_grid_vectors(image_height=3, image_width=3, output_stride=1)
edge_source = tf.cast(
[
[[1, 0.5], [0, 0]],
[[1, 0.5], [0, 0]],
],
tf.float32,
)
edge_destination = tf.cast(
[
[[1, 1.5], [2, 2]],
[[1, 1.5], [2, 2]],
],
tf.float32,
)
sigma = 1.0
pafs = edge_maps.make_multi_pafs(
xv=xv,
yv=yv,
edge_sources=edge_source,
edge_destinations=edge_destination,
sigma=sigma,
)
np.testing.assert_allclose(
pafs,
[
[
[[0.0, 0.916], [1.414, 1.414]],
[[0.0, 1.938], [1.248, 1.248]],
[[0.0, 0.916], [0.191, 0.191]],
],
[
[[0.0, 1.213], [1.248, 1.248]],
[[0.0, 2.0], [1.414, 1.414]],
[[0.0, 1.213], [1.248, 1.248]],
],
[
[[0.0, 0.916], [0.191, 0.191]],
[[0.0, 1.938], [1.248, 1.248]],
[[0.0, 0.916], [1.414, 1.414]],
],
],
atol=1e-3,
)
def test_make_multi_confmaps():
xv, yv = make_grid_vectors(image_height=4, image_width=5, output_stride=1)
instances = tf.cast([
[[0.5, 1.], [2., 2.]],
[[1.5, 1.], [2., 3.]],
[[np.nan, np.nan], [-1., 5.]],
], tf.float32)
cms = make_multi_confmaps(instances, xv=xv, yv=yv, sigma=1.)
assert cms.shape == (4, 5, 2)
assert cms.dtype == tf.float32
cm0 = make_confmaps(instances[0], xv=xv, yv=yv, sigma=1.)
cm1 = make_confmaps(instances[1], xv=xv, yv=yv, sigma=1.)
cm2 = make_confmaps(instances[2], xv=xv, yv=yv, sigma=1.)
np.testing.assert_array_equal(
cms, tf.reduce_max(tf.stack([cm0, cm1, cm2], axis=-1), axis=-1))
def test_make_multi_confmaps_with_offsets():
xv, yv = make_grid_vectors(image_height=4, image_width=5, output_stride=1)
instances = tf.cast(
[
[[0.5, 1.0], [2.0, 2.0]],
[[1.5, 1.0], [2.0, 3.0]],
[[np.nan, np.nan], [-1.0, 5.0]],
],
tf.float32,
)
cms, offsets = make_multi_confmaps_with_offsets(instances,
xv,
yv,
stride=1,
sigma=1.0,
offsets_threshold=0.2)
assert offsets.shape == (4, 5, 2, 2)
def test_edge_confidence_map():
xv, yv = make_grid_vectors(image_height=3, image_width=3, output_stride=1)
edge_source = tf.cast([[1, 0.5], [0, 0]], tf.float32)
edge_destination = tf.cast([[1, 1.5], [2, 2]], tf.float32)
sigma = 1.0
edge_confidence_map = edge_maps.make_edge_maps(
xv=xv,
yv=yv,
edge_source=edge_source,
edge_destination=edge_destination,
sigma=sigma,
)
np.testing.assert_allclose(
edge_confidence_map,
[
[[0.458, 1.000], [0.969, 0.882], [0.458, 0.135]],
[[0.607, 0.882], [1.000, 1.000], [0.607, 0.882]],
[[0.458, 0.135], [0.969, 0.882], [0.458, 1.000]],
],
atol=1e-3,
)
def transform_dataset(self, input_ds: tf.data.Dataset) -> tf.data.Dataset:
"""Create a dataset that contains the generated confidence maps.
Args:
input_ds: A dataset with elements that contain the keys "image",
"instances" and "skeleton_inds".
Returns:
A `tf.data.Dataset` with the same keys as the input, as well as
"part_affinity_fields".
The "part_affinity_fields" key will be a tensor of shape
(grid_height, grid_width, n_edges, 2) containing the combined part affinity
fields of all instances in the frame.
If the `flatten_channels` attribute is set to True, the last 2 axes of the
"part_affinity_fields" are flattened to produce a tensor of shape
(grid_height, grid_width, n_edges * 2). This is a convenient form when
training models as a rank-4 (batched) tensor will generally be expected.
Notes:
The output stride is relative to the current scale of the image. To map
points on the part affinity fields to the raw image, first multiply them by
the output stride, and then scale the x- and y-coordinates by the "scale"
key.
Importantly, the sigma will be proportional to the current image grid, not
the original grid prior to scaling operations.
"""
# Infer image dimensions to generate sampling grid.
test_example = next(iter(input_ds))
image_height = test_example["image"].shape[0]
image_width = test_example["image"].shape[1]
# Generate sampling grid vectors.
xv, yv = make_grid_vectors(
image_height=image_height,
image_width=image_width,
output_stride=self.output_stride,
)
grid_height = len(yv)
grid_width = len(xv)
# Pull out edge indices.
# TODO: Multi-skeleton support.
edge_inds = tf.cast(self.skeletons[0].edge_inds, dtype=tf.int32)
n_edges = len(edge_inds)
def generate_pafs(example):
"""Local processing function for dataset mapping."""
edge_sources, edge_destinations = get_edge_points(
example["instances"], edge_inds
)
edge_sources = tf.ensure_shape(edge_sources, (None, n_edges, 2))
edge_destinations = tf.ensure_shape(edge_destinations, (None, n_edges, 2))
pafs = make_multi_pafs(
xv=xv,
yv=yv,
edge_sources=edge_sources,
edge_destinations=edge_destinations,
sigma=self.sigma,
)
pafs = tf.ensure_shape(pafs, (grid_height, grid_width, n_edges, 2))
if self.flatten_channels:
pafs = tf.reshape(pafs, [grid_height, grid_width, n_edges * 2])
pafs = tf.ensure_shape(pafs, (grid_height, grid_width, n_edges * 2))
example["part_affinity_fields"] = pafs
return example
# Map transformation.
output_ds = input_ds.map(
generate_pafs, num_parallel_calls=tf.data.experimental.AUTOTUNE
)
return output_ds
def transform_dataset(self, input_ds: tf.data.Dataset) -> tf.data.Dataset:
"""Create a dataset that contains the generated confidence maps.
Args:
input_ds: A dataset with elements that contain the keys `"instances"` and
`"image"`.
Returns:
A `tf.data.Dataset` with the same keys as the input, as well as
`"confidence_maps"` containing the generated confidence maps.
Notes:
The output stride is relative to the current scale of the image. To map
points on the confidence maps to the raw image, first multiply them by the
output stride, and then scale the x- and y-coordinates by the "scale" key.
Importantly, the `sigma` will be proportional to the current image grid, not
the original grid prior to scaling operations.
"""
# Infer image dimensions to generate sampling grid.
test_example = next(iter(input_ds))
image_height = test_example["image"].shape[0]
image_width = test_example["image"].shape[1]
# Generate sampling grid vectors.
xv, yv = make_grid_vectors(
image_height=image_height,
image_width=image_width,
output_stride=self.output_stride,
)
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
# Map transformation.
output_ds = input_ds.map(
generate_confmaps,
num_parallel_calls=tf.data.experimental.AUTOTUNE)
return output_ds
def transform_dataset(self, input_ds: tf.data.Dataset) -> tf.data.Dataset:
"""Create a dataset that contains the generated confidence maps.
Args:
input_ds: A dataset with elements that contain the keys `"image"`, `"scale"`
and either "instances" or "centroids" depending on whether the
`centroids` attribute is set to `True`.
Returns:
A `tf.data.Dataset` with the same keys as the input, as well as a
`"confidence_maps"` or `"centroid_confidence_maps"` key containing the
generated confidence maps.
If the `with_offsets` attribute is `True`, example will contain a
`"offsets"` key.
Notes:
The output stride is relative to the current scale of the image. To map
points on the confidence maps to the raw image, first multiply them by the
output stride, and then scale the x- and y-coordinates by the `"scale"` key.
Importantly, the `sigma` will be proportional to the current image grid, not
the original grid prior to scaling operations.
"""
# Infer image dimensions to generate the full scale sampling grid.
test_example = next(iter(input_ds))
image_height = test_example["image"].shape[0]
image_width = test_example["image"].shape[1]
# Generate sampling grid vectors.
xv, yv = make_grid_vectors(
image_height=image_height,
image_width=image_width,
output_stride=self.output_stride,
)
def generate_multi_confmaps(example):
"""Local processing function for dataset mapping."""
if self.centroids:
points = tf.expand_dims(example["centroids"], axis=1)
cm_key = "centroid_confidence_maps"
else:
points = example["instances"]
cm_key = "confidence_maps"
if self.with_offsets:
cms, offsets = make_multi_confmaps_with_offsets(
points,
xv,
yv,
self.output_stride,
self.sigma * self.output_stride,
self.offsets_threshold,
)
example["offsets"] = offsets
if self.flatten_offsets:
shape = tf.shape(example["offsets"])
example["offsets"] = tf.reshape(example["offsets"],
[shape[0], shape[1], -1])
else:
cms = make_multi_confmaps(points, xv, yv,
self.sigma * self.output_stride)
example[cm_key] = cms
return example
# Map transformation.
output_ds = input_ds.map(
generate_multi_confmaps,
num_parallel_calls=tf.data.experimental.AUTOTUNE)
return output_ds
