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 build_step_data."""17 18import os19 20from absl import flags21import numpy as np22from PIL import Image23import tensorflow as tf24 25from deeplab2.data import build_step_data26 27FLAGS = flags.FLAGS28 29 30class BuildStepDataTest(tf.test.TestCase):31 32  def setUp(self):33    super().setUp()34    self.data_dir = FLAGS.test_tmpdir35    self.height = 10036    self.width = 10037    self.sequence_id = '010'38 39  def _create_images(self, split):40    image_path = os.path.join(self.data_dir, build_step_data._IMAGE_FOLDER_NAME,41                              split, self.sequence_id)42    panoptic_map_path = os.path.join(self.data_dir,43                                     build_step_data._PANOPTIC_MAP_FOLDER_NAME,44                                     split, self.sequence_id)45 46    tf.io.gfile.makedirs(image_path)47    tf.io.gfile.makedirs(panoptic_map_path)48    self.panoptic_maps = {}49    for image_id in [101, 100]:50      self.panoptic_maps[image_id] = self._create_image_and_panoptic_map(51          image_path, panoptic_map_path, image_id)52 53  def _create_image_and_panoptic_map(self, image_path, panoptic_path, image_id):54    """Creates dummy images and panoptic maps."""55    # Dummy image.56    image = np.random.randint(57        0, 255, (self.height, self.width, 3), dtype=np.uint8)58    with tf.io.gfile.GFile(59        os.path.join(image_path, '%06d.png' % image_id), 'wb') as f:60      Image.fromarray(image).save(f, format='PNG')61 62    # Dummy panoptic map.63    semantic = np.random.randint(64        0, 20, (self.height, self.width), dtype=np.int32)65    instance = np.random.randint(66        0, 1000, (self.height, self.width), dtype=np.int32)67    encoded_panoptic_map = np.dstack(68        (semantic, instance // 256, instance % 256)).astype(np.uint8)69    with tf.io.gfile.GFile(70        os.path.join(panoptic_path, '%06d.png' % image_id), 'wb') as f:71      Image.fromarray(encoded_panoptic_map).save(f, format='PNG')72    decoded_panoptic_map = semantic * 1000 + instance73    return decoded_panoptic_map74 75  def test_build_step_dataset_correct(self):76    split = 'train'77    self._create_images(split)78    build_step_data._convert_dataset(79        step_root=self.data_dir,80        dataset_split=split,81        output_dir=FLAGS.test_tmpdir)82    # We will have 2 shards with each shard containing 1 image.83    num_shards = 284    output_record = os.path.join(85        FLAGS.test_tmpdir, build_step_data._TF_RECORD_PATTERN %86        (split, 0, num_shards))87    self.assertTrue(tf.io.gfile.exists(output_record))88 89    # Parses tf record.90    image_ids = sorted(self.panoptic_maps)91    for i, raw_record in enumerate(92        tf.data.TFRecordDataset([output_record]).take(5)):93      image_id = image_ids[i]94      example = tf.train.Example.FromString(raw_record.numpy())95      panoptic_map = np.fromstring(96          example.features.feature['image/segmentation/class/encoded']97          .bytes_list.value[0],98          dtype=np.int32).reshape((self.height, self.width))99      np.testing.assert_array_equal(panoptic_map, self.panoptic_maps[image_id])100      self.assertEqual(101          example.features.feature['video/sequence_id'].bytes_list.value[0],102          b'010')103      self.assertEqual(104          example.features.feature['video/frame_id'].bytes_list.value[0],105          b'%06d' % image_id)106 107  def test_build_step_dataset_correct_with_two_frames(self):108    split = 'train'109    self._create_images(split)110    build_step_data._convert_dataset(111        step_root=self.data_dir,112        dataset_split=split,113        output_dir=FLAGS.test_tmpdir, use_two_frames=True)114    num_shards = 2115    output_record = os.path.join(116        FLAGS.test_tmpdir, build_step_data._TF_RECORD_PATTERN %117        (split, 0, num_shards))118    self.assertTrue(tf.io.gfile.exists(output_record))119 120    # Parses tf record.121    image_ids = sorted(self.panoptic_maps)122    for i, raw_record in enumerate(123        tf.data.TFRecordDataset([output_record]).take(5)):124      image_id = image_ids[i]125      example = tf.train.Example.FromString(raw_record.numpy())126      panoptic_map = np.fromstring(127          example.features.feature['image/segmentation/class/encoded']128          .bytes_list.value[0],129          dtype=np.int32).reshape((self.height, self.width))130      np.testing.assert_array_equal(panoptic_map, self.panoptic_maps[image_id])131      prev_panoptic_map = np.fromstring(132          example.features.feature['prev_image/segmentation/class/encoded']133          .bytes_list.value[0],134          dtype=np.int32).reshape((self.height, self.width))135      if i == 0:136        # First frame.137        np.testing.assert_array_equal(panoptic_map, prev_panoptic_map)138      else:139        # Not a first frame.140        np.testing.assert_array_equal(prev_panoptic_map, self.panoptic_maps[0])141      self.assertEqual(142          example.features.feature['video/sequence_id'].bytes_list.value[0],143          b'010')144      self.assertEqual(145          example.features.feature['video/frame_id'].bytes_list.value[0],146          b'%06d' % image_id)147 148  def test_build_step_dataset_with_two_frames_shared_by_sequence(self):149    split = 'val'150    self._create_images(split)151    build_step_data._convert_dataset(152        step_root=self.data_dir,153        dataset_split=split,154        output_dir=FLAGS.test_tmpdir, use_two_frames=True)155    # Only one shard since there is only one sequence for the val set.156    num_shards = 1157    output_record = os.path.join(158        FLAGS.test_tmpdir, build_step_data._TF_RECORD_PATTERN %159        (split, 0, num_shards))160    self.assertTrue(tf.io.gfile.exists(output_record))161 162 163if __name__ == '__main__':164  tf.test.main()165