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"""Tests for build_coco_data."""17 18import json19import os20 21from absl import flags22import numpy as np23from PIL import Image24import tensorflow as tf25 26from deeplab2.data import build_coco_data27from deeplab2.data import coco_constants28 29FLAGS = flags.FLAGS30_TEST_FILE_NAME = '000000123456.png'31 32 33class BuildCOCODataTest(tf.test.TestCase):34 35  def setUp(self):36    super().setUp()37    self.data_dir = FLAGS.test_tmpdir38    self.height = 10039    self.width = 10040    self.split = 'train'41    image_path = os.path.join(self.data_dir,42                              build_coco_data._FOLDERS_MAP[self.split]['image'])43    panoptic_map_path = os.path.join(self.data_dir,44                                     build_coco_data._FOLDERS_MAP45                                     [self.split]['label'])46    tf.io.gfile.makedirs(panoptic_map_path)47    panoptic_map_path = os.path.join(panoptic_map_path,48                                     'panoptic_%s2017' % self.split)49 50    tf.io.gfile.makedirs(image_path)51    tf.io.gfile.makedirs(panoptic_map_path)52    self.panoptic_maps = {}53    image_id = int(_TEST_FILE_NAME[:-4])54    self.panoptic_maps[image_id] = self._create_image_and_panoptic_map(55        image_path, panoptic_map_path, image_id)56 57  def _create_image_and_panoptic_map(self, image_path, panoptic_path, image_id):58    def id2rgb(id_map):59      id_map_copy = id_map.copy()60      rgb_shape = tuple(list(id_map.shape) + [3])61      rgb_map = np.zeros(rgb_shape, dtype=np.uint8)62      for i in range(3):63        rgb_map[..., i] = id_map_copy % 25664        id_map_copy //= 25665      return rgb_map66 67    # Creates dummy images and panoptic maps.68    # Dummy image.69    image = np.random.randint(70        0, 255, (self.height, self.width, 3), dtype=np.uint8)71    with tf.io.gfile.GFile(72        os.path.join(image_path, '%012d.jpg' % image_id), 'wb') as f:73      Image.fromarray(image).save(f, format='JPEG')74 75    # Dummy panoptic map.76    semantic = np.random.randint(77        0, 201, (self.height, self.width), dtype=np.int32)78    instance_ = np.random.randint(79        0, 100, (self.height, self.width), dtype=np.int32)80    id_mapping = coco_constants.get_id_mapping()81    valid_semantic = id_mapping.keys()82    for i in range(201):83      if i not in valid_semantic:84        mask = (semantic == i)85        semantic[mask] = 086        instance_[mask] = 087 88    instance = instance_.copy()89    segments_info = []90    for sem in np.unique(semantic):91      ins_id = 192      if sem == 0:93        continue94      if id_mapping[sem] in build_coco_data._CLASS_HAS_INSTANCE_LIST:95        for ins in np.unique(instance_[semantic == sem]):96          instance[np.logical_and(semantic == sem, instance_ == ins)] = ins_id97          area = np.logical_and(semantic == sem, instance_ == ins).sum()98          idx = sem * 256 + ins_id99          iscrowd = 0100          segments_info.append({101              'id': idx.tolist(),102              'category_id': sem.tolist(),103              'area': area.tolist(),104              'iscrowd': iscrowd,105          })106          ins_id += 1107      else:108        instance[semantic == sem] = 0109        area = (semantic == sem).sum()110        idx = sem * 256111        iscrowd = 0112        segments_info.append({113            'id': idx.tolist(),114            'category_id': sem.tolist(),115            'area': area.tolist(),116            'iscrowd': iscrowd,117        })118 119    encoded_panoptic_map = semantic * 256 + instance120    encoded_panoptic_map = id2rgb(encoded_panoptic_map)121    with tf.io.gfile.GFile(122        os.path.join(panoptic_path, '%012d.png' % image_id), 'wb') as f:123      Image.fromarray(encoded_panoptic_map).save(f, format='PNG')124 125    for i in range(201):126      if i in valid_semantic:127        mask = (semantic == i)128        semantic[mask] = id_mapping[i]129 130    decoded_panoptic_map = semantic * 256 + instance131 132    # Write json file133    json_annotation = {134        'annotations': [135            {136                'file_name': _TEST_FILE_NAME,137                'image_id': int(_TEST_FILE_NAME[:-4]),138                'segments_info': segments_info139            }140        ]141    }142    json_annotation_path = os.path.join(self.data_dir,143                                        build_coco_data._FOLDERS_MAP144                                        [self.split]['label'],145                                        'panoptic_%s2017.json' % self.split)146    with tf.io.gfile.GFile(json_annotation_path, 'w') as f:147      json.dump(json_annotation, f, indent=2)148 149    return decoded_panoptic_map150 151  def test_build_coco_dataset_correct(self):152    build_coco_data._convert_dataset(153        coco_root=self.data_dir,154        dataset_split=self.split,155        output_dir=FLAGS.test_tmpdir)156    output_record = os.path.join(157        FLAGS.test_tmpdir, '%s-%05d-of-%05d.tfrecord' %158        (self.split, 0, build_coco_data._NUM_SHARDS))159    self.assertTrue(tf.io.gfile.exists(output_record))160 161    # Parses tf record.162    image_ids = sorted(self.panoptic_maps)163    for i, raw_record in enumerate(164        tf.data.TFRecordDataset([output_record]).take(5)):165      image_id = image_ids[i]166      example = tf.train.Example.FromString(raw_record.numpy())167      panoptic_map = np.fromstring(168          example.features.feature['image/segmentation/class/encoded']169          .bytes_list.value[0],170          dtype=np.int32).reshape((self.height, self.width))171      np.testing.assert_array_equal(panoptic_map, self.panoptic_maps[image_id])172 173if __name__ == '__main__':174  tf.test.main()175