|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
"""Utility functions to set up unit tests on Panoptic Segmentation code.""" |
|
|
|
import os |
|
from typing import Mapping, Optional, Tuple |
|
|
|
from absl import flags |
|
import numpy as np |
|
from PIL import Image |
|
|
|
import tensorflow as tf |
|
|
|
FLAGS = flags.FLAGS |
|
|
|
_TEST_DATA_DIR = ('deeplab2/' |
|
'evaluation/testdata') |
|
|
|
|
|
def read_test_image(testdata_path: str, |
|
image_format: Optional[str] = None) -> np.ndarray: |
|
"""Loads a test image. |
|
|
|
Args: |
|
testdata_path: Image path relative to panoptic_segmentation/testdata as a |
|
string. |
|
image_format: Format of the image. Can be one of 'RGBA', 'RGB', or 'L'. |
|
|
|
Returns: |
|
The image, as a numpy array. |
|
""" |
|
image_path = os.path.join(_TEST_DATA_DIR, testdata_path) |
|
with tf.io.gfile.GFile(image_path, 'rb') as f: |
|
image = Image.open(f) |
|
if image_format is not None: |
|
image = image.convert(image_format) |
|
return np.array(image) |
|
|
|
|
|
def read_segmentation_with_rgb_color_map( |
|
image_testdata_path: str, |
|
rgb_to_semantic_label: Mapping[Tuple[int, int, int], int], |
|
output_dtype: Optional[np.dtype] = None) -> np.ndarray: |
|
"""Reads a test segmentation as an image and a map from colors to labels. |
|
|
|
Args: |
|
image_testdata_path: Image path relative to panoptic_segmentation/testdata |
|
as a string. |
|
rgb_to_semantic_label: Mapping from RGB colors to integer labels as a |
|
dictionary. |
|
output_dtype: Type of the output labels. If None, defaults to the type of |
|
the provided color map. |
|
|
|
Returns: |
|
A 2D numpy array of labels. |
|
|
|
Raises: |
|
ValueError: On an incomplete `rgb_to_semantic_label`. |
|
""" |
|
rgb_image = read_test_image(image_testdata_path, image_format='RGB') |
|
if len(rgb_image.shape) != 3 or rgb_image.shape[2] != 3: |
|
raise AssertionError('Expected RGB image, actual shape is %s' % |
|
(rgb_image.shape,)) |
|
|
|
num_pixels = rgb_image.shape[0] * rgb_image.shape[1] |
|
unique_colors = np.unique(np.reshape(rgb_image, [num_pixels, 3]), axis=0) |
|
if not set(map(tuple, unique_colors)).issubset(rgb_to_semantic_label.keys()): |
|
raise ValueError('RGB image has colors not in color map.') |
|
|
|
output_dtype = output_dtype or type( |
|
next(iter(rgb_to_semantic_label.values()))) |
|
output_labels = np.empty(rgb_image.shape[:2], dtype=output_dtype) |
|
for rgb_color, int_label in rgb_to_semantic_label.items(): |
|
color_array = np.array(rgb_color, ndmin=3) |
|
output_labels[np.all(rgb_image == color_array, axis=2)] = int_label |
|
return output_labels |
|
|
|
|
|
def panoptic_segmentation_with_class_map( |
|
instance_testdata_path: str, instance_label_to_semantic_label: Mapping[int, |
|
int] |
|
) -> Tuple[np.ndarray, np.ndarray]: |
|
"""Reads in a panoptic segmentation with an instance map and a map to classes. |
|
|
|
Args: |
|
instance_testdata_path: Path to a grayscale instance map, given as a string |
|
and relative to panoptic_segmentation/testdata. |
|
instance_label_to_semantic_label: A map from instance labels to class |
|
labels. |
|
|
|
Returns: |
|
A tuple `(instance_labels, class_labels)` of numpy arrays. |
|
|
|
Raises: |
|
ValueError: On a mismatched set of instances in |
|
the |
|
`instance_label_to_semantic_label`. |
|
""" |
|
instance_labels = read_test_image(instance_testdata_path, image_format='L') |
|
if set(np.unique(instance_labels)) != set( |
|
instance_label_to_semantic_label.keys()): |
|
raise ValueError('Provided class map does not match present instance ids.') |
|
|
|
class_labels = np.empty_like(instance_labels) |
|
for instance_id, class_id in instance_label_to_semantic_label.items(): |
|
class_labels[instance_labels == instance_id] = class_id |
|
|
|
return instance_labels, class_labels |
|
|
