karolmajek/maxdeeplab

1# coding=utf-82# Copyright 2021 The Deeplab2 Authors.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8#     http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15 16"""Utility functions to set up unit tests on Panoptic Segmentation code."""17 18import os19from typing import Mapping, Optional, Tuple20 21from absl import flags22import numpy as np23from PIL import Image24 25import tensorflow as tf26 27FLAGS = flags.FLAGS28 29_TEST_DATA_DIR = ('deeplab2/'30                  'evaluation/testdata')31 32 33def read_test_image(testdata_path: str,34                    image_format: Optional[str] = None) -> np.ndarray:35  """Loads a test image.36 37  Args:38    testdata_path: Image path relative to panoptic_segmentation/testdata as a39      string.40    image_format: Format of the image. Can be one of 'RGBA', 'RGB', or 'L'.41 42  Returns:43    The image, as a numpy array.44  """45  image_path = os.path.join(_TEST_DATA_DIR, testdata_path)46  with tf.io.gfile.GFile(image_path, 'rb') as f:47    image = Image.open(f)48    if image_format is not None:49      image = image.convert(image_format)50    return np.array(image)51 52 53def read_segmentation_with_rgb_color_map(54    image_testdata_path: str,55    rgb_to_semantic_label: Mapping[Tuple[int, int, int], int],56    output_dtype: Optional[np.dtype] = None) -> np.ndarray:57  """Reads a test segmentation as an image and a map from colors to labels.58 59  Args:60    image_testdata_path: Image path relative to panoptic_segmentation/testdata61      as a string.62    rgb_to_semantic_label: Mapping from RGB colors to integer labels as a63      dictionary.64    output_dtype: Type of the output labels. If None, defaults to the type of65      the provided color map.66 67  Returns:68    A 2D numpy array of labels.69 70  Raises:71    ValueError: On an incomplete `rgb_to_semantic_label`.72  """73  rgb_image = read_test_image(image_testdata_path, image_format='RGB')74  if len(rgb_image.shape) != 3 or rgb_image.shape[2] != 3:75    raise AssertionError('Expected RGB image, actual shape is %s' %76                         (rgb_image.shape,))77 78  num_pixels = rgb_image.shape[0] * rgb_image.shape[1]79  unique_colors = np.unique(np.reshape(rgb_image, [num_pixels, 3]), axis=0)80  if not set(map(tuple, unique_colors)).issubset(rgb_to_semantic_label.keys()):81    raise ValueError('RGB image has colors not in color map.')82 83  output_dtype = output_dtype or type(84      next(iter(rgb_to_semantic_label.values())))85  output_labels = np.empty(rgb_image.shape[:2], dtype=output_dtype)86  for rgb_color, int_label in rgb_to_semantic_label.items():87    color_array = np.array(rgb_color, ndmin=3)88    output_labels[np.all(rgb_image == color_array, axis=2)] = int_label89  return output_labels90 91 92def panoptic_segmentation_with_class_map(93    instance_testdata_path: str, instance_label_to_semantic_label: Mapping[int,94                                                                           int]95) -> Tuple[np.ndarray, np.ndarray]:96  """Reads in a panoptic segmentation with an instance map and a map to classes.97 98  Args:99    instance_testdata_path: Path to a grayscale instance map, given as a string100      and relative to panoptic_segmentation/testdata.101    instance_label_to_semantic_label: A map from instance labels to class102      labels.103 104  Returns:105    A tuple `(instance_labels, class_labels)` of numpy arrays.106 107  Raises:108    ValueError: On a mismatched set of instances in109    the110      `instance_label_to_semantic_label`.111  """112  instance_labels = read_test_image(instance_testdata_path, image_format='L')113  if set(np.unique(instance_labels)) != set(114      instance_label_to_semantic_label.keys()):115    raise ValueError('Provided class map does not match present instance ids.')116 117  class_labels = np.empty_like(instance_labels)118  for instance_id, class_id in instance_label_to_semantic_label.items():119    class_labels[instance_labels == instance_id] = class_id120 121  return instance_labels, class_labels122