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