karolmajek/Axial-DeepLab-SWideRNet

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"""Tests for data_utils."""17 18import io19import numpy as np20from PIL import Image21import tensorflow as tf22 23from deeplab2.data import data_utils24 25 26def _encode_png_image(image):27  """Helper method to encode input image in PNG format."""28  buffer = io.BytesIO()29  Image.fromarray(image).save(buffer, format='png')30  return buffer.getvalue()31 32 33class DataUtilsTest(tf.test.TestCase):34 35  def _create_test_image(self, height, width):36    rng = np.random.RandomState(319281498)37    return rng.randint(0, 255, size=(height, width, 3), dtype=np.uint8)38 39  def test_encode_and_decode(self):40    """Checks decode created tf.Example for semantic segmentation."""41    test_image_height = 2042    test_image_width = 1543    filename = 'dummy'44 45    image = self._create_test_image(test_image_height, test_image_width)46    # Take the last channel as dummy label.47    label = image[..., 0]48 49    example = data_utils.create_tfexample(50        image_data=_encode_png_image(image),51        image_format='png', filename=filename,52        label_data=_encode_png_image(label), label_format='png')53 54    # Parse created example, expect getting identical results.55    parser = data_utils.SegmentationDecoder(is_panoptic_dataset=False)56    parsed_tensors = parser(example.SerializeToString())57 58    self.assertIn('image', parsed_tensors)59    self.assertIn('image_name', parsed_tensors)60    self.assertIn('label', parsed_tensors)61    self.assertEqual(filename, parsed_tensors['image_name'])62    np.testing.assert_array_equal(image, parsed_tensors['image'].numpy())63    # Decoded label is a 3-D array with last dimension of 1.64    decoded_label = parsed_tensors['label'].numpy()65    np.testing.assert_array_equal(label, decoded_label[..., 0])66 67  def test_encode_and_decode_panoptic(self):68    test_image_height = 3169    test_image_width = 1770    filename = 'dummy'71 72    image = self._create_test_image(test_image_height, test_image_width)73    # Create dummy panoptic label in np.int32 dtype.74    label = np.dot(image.astype(np.int32), [1, 256, 256 * 256]).astype(np.int32)75    example = data_utils.create_tfexample(76        image_data=_encode_png_image(image),77        image_format='png', filename=filename,78        label_data=label.tostring(), label_format='raw')79 80    parser = data_utils.SegmentationDecoder(is_panoptic_dataset=True)81    parsed_tensors = parser(example.SerializeToString())82 83    self.assertIn('image', parsed_tensors)84    self.assertIn('image_name', parsed_tensors)85    self.assertIn('label', parsed_tensors)86    self.assertEqual(filename, parsed_tensors['image_name'])87    np.testing.assert_array_equal(image, parsed_tensors['image'].numpy())88    # Decoded label is a 3-D array with last dimension of 1.89    decoded_label = parsed_tensors['label'].numpy()90    np.testing.assert_array_equal(label, decoded_label[..., 0])91 92 93if __name__ == '__main__':94  tf.test.main()95