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"""Script to generate test data for cityscapes."""17 18import collections19import json20import os21 22from absl import app23from absl import flags24from absl import logging25import numpy as np26from PIL import Image27import tensorflow as tf28 29# resources dependency30 31from deeplab2.data import data_utils32from deeplab2.data import dataset33 34flags.DEFINE_string(35    'panoptic_annotation_path',36    'deeplab2/data/testdata/'37    'dummy_prediction.png',38    'Path to annotated test image with cityscapes encoding.')39flags.DEFINE_string(40    'panoptic_gt_output_path',41    'deeplab2/data/testdata/'42    'dummy_gt_for_vps.png',43    'Path to annotated test image with Video Panoptic Segmentation encoding.')44flags.DEFINE_string(45    'output_cityscapes_root',46    'deeplab2/data/testdata/',47    'Path to output root directory.')48 49FLAGS = flags.FLAGS50 51# Cityscapes label, using `TrainId`.52_CITYSCAPES_IGNORE = 25553# Each valid (not ignored) label below is a tuple of (TrainId, EvalId)54_CITYSCAPES_CAR = (13, 26)55_CITYSCAPES_TREE = (8, 21)56_CITYSCAPES_SKY = (10, 23)57_CITYSCAPES_BUILDING = (2, 11)58_CITYSCAPES_ROAD = (0, 7)59 60_IS_CROWD = 'is_crowd'61_NOT_CROWD = 'not_crowd'62 63_CLASS_HAS_INSTANCES_LIST = dataset.CITYSCAPES_PANOPTIC_INFORMATION.class_has_instances_list64_PANOPTIC_LABEL_DIVISOR = dataset.CITYSCAPES_PANOPTIC_INFORMATION.panoptic_label_divisor65_FILENAME_PREFIX = 'dummy_000000_000000'66 67 68def create_test_data(annotation_path):69  """Creates cityscapes panoptic annotation, vps annotation and segment info.70 71  Our Video Panoptic Segmentation (VPS) encoding uses ID == semantic trainID *72  1000 + instance ID (starting at 1) with instance ID == 0 marking73  crowd regions.74 75  Args:76    annotation_path: The path to the annotation to be loaded.77 78  Returns:79    A tuple of cityscape annotation, vps annotation and segment infos.80  """81  # Convert panoptic labels to cityscapes label format.82 83  # Dictionary mapping converted panoptic annotation to its corresponding84  # Cityscapes label. Here the key is encoded by converting each RGB pixel85  # value to 1 * R + 256 * G + 256 * 256 * B.86  panoptic_label_to_cityscapes_label = {87      0: (_CITYSCAPES_IGNORE, _NOT_CROWD),88      31110: (_CITYSCAPES_CAR, _NOT_CROWD),89      31354: (_CITYSCAPES_CAR, _IS_CROWD),90      35173: (_CITYSCAPES_CAR, _NOT_CROWD),91      488314: (_CITYSCAPES_CAR, _IS_CROWD),92      549788: (_CITYSCAPES_CAR, _IS_CROWD),93      1079689: (_CITYSCAPES_CAR, _IS_CROWD),94      1341301: (_CITYSCAPES_CAR, _NOT_CROWD),95      1544590: (_CITYSCAPES_CAR, _NOT_CROWD),96      1926498: (_CITYSCAPES_CAR, _NOT_CROWD),97      4218944: (_CITYSCAPES_TREE, _NOT_CROWD),98      4251840: (_CITYSCAPES_SKY, _NOT_CROWD),99      6959003: (_CITYSCAPES_BUILDING, _NOT_CROWD),100      # To be merged with the building segment above.101      8396960: (_CITYSCAPES_BUILDING, _NOT_CROWD),102      8413312: (_CITYSCAPES_ROAD, _NOT_CROWD),103  }104  with tf.io.gfile.GFile(annotation_path, 'rb') as f:105    panoptic = data_utils.read_image(f.read())106 107  # Input panoptic annotation is RGB color coded, here we convert each pixel108  # to a unique number to avoid comparing 3-tuples.109  panoptic = np.dot(panoptic, [1, 256, 256 * 256])110  # Creates cityscapes panoptic map. Cityscapes use ID == semantic EvalId for111  # `stuff` segments and `thing` segments with `iscrowd` label, and112  # ID == semantic EvalId * 1000 + instance ID (starting from 0) for other113  # `thing` segments.114  cityscapes_panoptic = np.zeros_like(panoptic, dtype=np.int32)115  # Creates Video Panoptic Segmentation (VPS) map. We use ID == semantic116  # trainID * 1000 + instance ID (starting at 1) with instance ID == 0 marking117  # crowd regions.118  vps_panoptic = np.zeros_like(panoptic, dtype=np.int32)119  num_instances_per_class = collections.defaultdict(int)120  unique_labels = np.unique(panoptic)121 122  # Dictionary that maps segment id to segment info.123  segments_info = {}124  for label in unique_labels:125    cityscapes_label, is_crowd = panoptic_label_to_cityscapes_label[label]126    selected_pixels = panoptic == label127 128    if cityscapes_label == _CITYSCAPES_IGNORE:129      vps_panoptic[selected_pixels] = (130          _CITYSCAPES_IGNORE * _PANOPTIC_LABEL_DIVISOR)131      continue132 133    train_id, eval_id = tuple(cityscapes_label)134    cityscapes_id = eval_id135    vps_id = train_id * _PANOPTIC_LABEL_DIVISOR136    if train_id in _CLASS_HAS_INSTANCES_LIST:137      # `thing` class.138      if is_crowd != _IS_CROWD:139        cityscapes_id = (140            eval_id * _PANOPTIC_LABEL_DIVISOR +141            num_instances_per_class[train_id])142        # First instance should have ID 1.143        vps_id += num_instances_per_class[train_id] + 1144        num_instances_per_class[train_id] += 1145 146    cityscapes_panoptic[selected_pixels] = cityscapes_id147    vps_panoptic[selected_pixels] = vps_id148    pixel_area = int(np.sum(selected_pixels))149    if cityscapes_id in segments_info:150      logging.info('Merging segments with label %d into segment %d', label,151                   cityscapes_id)152      segments_info[cityscapes_id]['area'] += pixel_area153    else:154      segments_info[cityscapes_id] = {155          'area': pixel_area,156          'category_id': train_id,157          'id': cityscapes_id,158          'iscrowd': 1 if is_crowd == _IS_CROWD else 0,159      }160 161  cityscapes_panoptic = np.dstack([162      cityscapes_panoptic % 256, cityscapes_panoptic // 256,163      cityscapes_panoptic // 256 // 256164  ])165  vps_panoptic = np.dstack(166      [vps_panoptic % 256, vps_panoptic // 256, vps_panoptic // 256 // 256])167  return (cityscapes_panoptic.astype(np.uint8), vps_panoptic.astype(np.uint8),168          list(segments_info.values()))169 170 171def main(argv):172  if len(argv) > 1:173    raise app.UsageError('Too many command-line arguments.')174 175  data_path = FLAGS.panoptic_annotation_path  # OSS: removed internal filename loading.176  panoptic_map, vps_map, segments_info = create_test_data(data_path)177  panoptic_map_filename = _FILENAME_PREFIX + '_gtFine_panoptic.png'178  panoptic_map_path = os.path.join(FLAGS.output_cityscapes_root, 'gtFine',179                                   'cityscapes_panoptic_dummy_trainId',180                                   panoptic_map_filename)181 182  gt_output_path = FLAGS.panoptic_gt_output_path  # OSS: removed internal filename loading.183  with tf.io.gfile.GFile(gt_output_path, 'wb') as f:184    Image.fromarray(vps_map).save(f, format='png')185 186  panoptic_map_path = panoptic_map_path  # OSS: removed internal filename loading.187  with tf.io.gfile.GFile(panoptic_map_path, 'wb') as f:188    Image.fromarray(panoptic_map).save(f, format='png')189 190  json_annotation = {191      'annotations': [{192          'file_name': _FILENAME_PREFIX + '_gtFine_panoptic.png',193          'image_id': _FILENAME_PREFIX,194          'segments_info': segments_info195      }]196  }197  json_annotation_path = os.path.join(FLAGS.output_cityscapes_root, 'gtFine',198                                      'cityscapes_panoptic_dummy_trainId.json')199  json_annotation_path = json_annotation_path  # OSS: removed internal filename loading.200  with tf.io.gfile.GFile(json_annotation_path, 'w') as f:201    json.dump(json_annotation, f, indent=2)202 203 204if __name__ == '__main__':205  app.run(main)206