karolmajek/maxdeeplab
0
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"""Converts COCO data to sharded TFRecord file format with Example protos.17 18Please check19 ../g3doc/setup/coco.md20for instructions.21"""22 23import collections24import json25import math26import os27 28from typing import Sequence, Tuple, Any29 30from absl import app31from absl import flags32from absl import logging33import numpy as np34import tensorflow as tf35 36from deeplab2.data import coco_constants37from deeplab2.data import data_utils38from deeplab2.data import dataset39 40FLAGS = flags.FLAGS41 42flags.DEFINE_string('coco_root', None, 'coco dataset root folder.')43 44flags.DEFINE_string('output_dir', None,45 'Path to save converted TFRecord of TensorFlow examples.')46 47flags.DEFINE_boolean('treat_crowd_as_ignore', True,48 'Whether to apply ignore labels to crowd pixels in '49 'panoptic label.')50 51_NUM_SHARDS = 100052 53 54_SPLITS_TO_SIZES = dataset.COCO_PANOPTIC_INFORMATION.splits_to_sizes55_IGNORE_LABEL = dataset.COCO_PANOPTIC_INFORMATION.ignore_label56_CLASS_HAS_INSTANCE_LIST = dataset.COCO_PANOPTIC_INFORMATION.class_has_instances_list57_PANOPTIC_LABEL_DIVISOR = dataset.COCO_PANOPTIC_INFORMATION.panoptic_label_divisor58_CLASS_MAPPING = coco_constants.get_id_mapping()59 60# A map from data type to folder name that saves the data.61_FOLDERS_MAP = {62 'train': {63 'image': 'train2017',64 'label': 'annotations',65 },66 'val': {67 'image': 'val2017',68 'label': 'annotations',69 },70 'test': {71 'image': 'test2017',72 'label': '',73 }74}75 76# A map from data type to data format.77_DATA_FORMAT_MAP = {78 'image': 'jpg',79 'label': 'png',80}81_PANOPTIC_LABEL_FORMAT = 'raw'82 83 84def _get_images(coco_root: str, dataset_split: str) -> Sequence[str]:85 """Gets files for the specified data type and dataset split.86 87 Args:88 coco_root: String, path to coco dataset root folder.89 dataset_split: String, dataset split ('train', 'val', 'test').90 91 Returns:92 A list of sorted file names.93 """94 pattern = '*.%s' % _DATA_FORMAT_MAP['image']95 search_files = os.path.join(96 coco_root, _FOLDERS_MAP[dataset_split]['image'], pattern)97 filenames = tf.io.gfile.glob(search_files)98 return sorted(filenames)99 100 101def _get_panoptic_annotation(coco_root: str, dataset_split: str,102 annotation_file_name: str) -> str:103 panoptic_folder = 'panoptic_%s2017' % dataset_split104 return os.path.join(coco_root, _FOLDERS_MAP[dataset_split]['label'],105 panoptic_folder, annotation_file_name)106 107 108def _read_segments(coco_root: str, dataset_split: str):109 """Reads segments information from json file.110 111 Args:112 coco_root: String, path to coco dataset root folder.113 dataset_split: String, dataset split.114 115 Returns:116 segments_dict: A dictionary that maps file prefix of annotation_file_name to117 a tuple of (panoptic annotation file name, segments). Please refer to118 _generate_panoptic_label() method on the detail structure of `segments`.119 120 Raises:121 ValueError: If found duplicated image id in annotations.122 """123 json_filename = os.path.join(124 coco_root, _FOLDERS_MAP[dataset_split]['label'],125 'panoptic_%s2017.json' % dataset_split)126 with tf.io.gfile.GFile(json_filename) as f:127 panoptic_dataset = json.load(f)128 129 segments_dict = {}130 for annotation in panoptic_dataset['annotations']:131 image_id = annotation['image_id']132 if image_id in segments_dict:133 raise ValueError('Image ID %s already exists' % image_id)134 annotation_file_name = annotation['file_name']135 segments = annotation['segments_info']136 137 segments_dict[os.path.splitext(annotation_file_name)[-2]] = (138 annotation_file_name, segments)139 140 return segments_dict141 142 143def _generate_panoptic_label(panoptic_annotation_file: str, segments:144 Any) -> np.ndarray:145 """Creates panoptic label map from annotations.146 147 Args:148 panoptic_annotation_file: String, path to panoptic annotation.149 segments: A list of dictionaries containing information of every segment.150 Read from panoptic_${DATASET_SPLIT}2017.json. This method consumes151 the following fields in each dictionary:152 - id: panoptic id153 - category_id: semantic class id154 - area: pixel area of this segment155 - iscrowd: if this segment is crowd region156 157 Returns:158 A 2D numpy int32 array with the same height / width with panoptic159 annotation. Each pixel value represents its panoptic ID. Please refer to160 g3doc/setup/coco.md for more details about how panoptic ID is assigned.161 """162 with tf.io.gfile.GFile(panoptic_annotation_file, 'rb') as f:163 panoptic_label = data_utils.read_image(f.read())164 165 if panoptic_label.mode != 'RGB':166 raise ValueError('Expect RGB image for panoptic label, gets %s' %167 panoptic_label.mode)168 169 panoptic_label = np.array(panoptic_label, dtype=np.int32)170 # COCO panoptic map is created by:171 # color = [segmentId % 256, segmentId // 256, segmentId // 256 // 256]172 panoptic_label = np.dot(panoptic_label, [1, 256, 256 * 256])173 174 semantic_label = np.ones_like(panoptic_label) * _IGNORE_LABEL175 instance_label = np.zeros_like(panoptic_label)176 # Running count of instances per semantic category.177 instance_count = collections.defaultdict(int)178 179 for segment in segments:180 selected_pixels = panoptic_label == segment['id']181 pixel_area = np.sum(selected_pixels)182 if pixel_area != segment['area']:183 raise ValueError('Expect %d pixels for segment %s, gets %d.' %184 (segment['area'], segment, pixel_area))185 186 category_id = segment['category_id']187 188 # Map the category_id to contiguous ids189 category_id = _CLASS_MAPPING[category_id]190 191 semantic_label[selected_pixels] = category_id192 193 if category_id in _CLASS_HAS_INSTANCE_LIST:194 if segment['iscrowd']:195 # COCO crowd pixels will have instance ID of 0.196 if FLAGS.treat_crowd_as_ignore:197 semantic_label[selected_pixels] = _IGNORE_LABEL198 continue199 # Non-crowd pixels will have instance ID starting from 1.200 instance_count[category_id] += 1201 if instance_count[category_id] >= _PANOPTIC_LABEL_DIVISOR:202 raise ValueError('Too many instances for category %d in this image.' %203 category_id)204 instance_label[selected_pixels] = instance_count[category_id]205 elif segment['iscrowd']:206 raise ValueError('Stuff class should not have `iscrowd` label.')207 208 panoptic_label = semantic_label * _PANOPTIC_LABEL_DIVISOR + instance_label209 return panoptic_label.astype(np.int32)210 211 212def _create_panoptic_label(coco_root: str, dataset_split: str, image_path: str,213 segments_dict: Any214 ) -> Tuple[str, str]:215 """Creates labels for panoptic segmentation.216 217 Args:218 coco_root: String, path to coco dataset root folder.219 dataset_split: String, dataset split ('train', 'val', 'test').220 image_path: String, path to the image file.221 segments_dict:222 Read from panoptic_${DATASET_SPLIT}2017.json. This method consumes223 the following fields in each dictionary:224 - id: panoptic id225 - category_id: semantic class id226 - area: pixel area of this segment227 - iscrowd: if this segment is crowd region228 229 Returns:230 A panoptic label where each pixel value represents its panoptic ID.231 Please refer to g3doc/setup/coco.md for more details about howpanoptic ID232 is assigned.233 A string indicating label format in TFRecord.234 """235 236 image_path = os.path.normpath(image_path)237 path_list = image_path.split(os.sep)238 file_name = path_list[-1]239 240 annotation_file_name, segments = segments_dict[241 os.path.splitext(file_name)[-2]]242 panoptic_annotation_file = _get_panoptic_annotation(coco_root,243 dataset_split,244 annotation_file_name)245 246 panoptic_label = _generate_panoptic_label(panoptic_annotation_file, segments)247 return panoptic_label.tostring(), _PANOPTIC_LABEL_FORMAT248 249 250def _convert_dataset(coco_root: str, dataset_split: str,251 output_dir: str) -> None:252 """Converts the specified dataset split to TFRecord format.253 254 Args:255 coco_root: String, path to coco dataset root folder.256 dataset_split: String, the dataset split (one of `train`, `val` and `test`).257 output_dir: String, directory to write output TFRecords to.258 """259 image_files = _get_images(coco_root, dataset_split)260 261 num_images = len(image_files)262 263 if dataset_split != 'test':264 segments_dict = _read_segments(coco_root, dataset_split)265 266 num_per_shard = int(math.ceil(len(image_files) / _NUM_SHARDS))267 268 for shard_id in range(_NUM_SHARDS):269 shard_filename = '%s-%05d-of-%05d.tfrecord' % (270 dataset_split, shard_id, _NUM_SHARDS)271 output_filename = os.path.join(output_dir, shard_filename)272 with tf.io.TFRecordWriter(output_filename) as tfrecord_writer:273 start_idx = shard_id * num_per_shard274 end_idx = min((shard_id + 1) * num_per_shard, num_images)275 for i in range(start_idx, end_idx):276 # Read the image.277 with tf.io.gfile.GFile(image_files[i], 'rb') as f:278 image_data = f.read()279 280 if dataset_split == 'test':281 label_data, label_format = None, None282 else:283 label_data, label_format = _create_panoptic_label(284 coco_root, dataset_split, image_files[i], segments_dict)285 286 # Convert to tf example.287 image_path = os.path.normpath(image_files[i])288 path_list = image_path.split(os.sep)289 file_name = path_list[-1]290 file_prefix = file_name.replace(_DATA_FORMAT_MAP['image'], '')291 example = data_utils.create_tfexample(image_data,292 'jpeg',293 file_prefix, label_data,294 label_format)295 296 tfrecord_writer.write(example.SerializeToString())297 298 299def main(unused_argv: Sequence[str]) -> None:300 tf.io.gfile.makedirs(FLAGS.output_dir)301 302 for dataset_split in ('train', 'val', 'test'):303 logging.info('Starts processing dataset split %s.', dataset_split)304 _convert_dataset(FLAGS.coco_root, dataset_split, FLAGS.output_dir)305 306 307if __name__ == '__main__':308 flags.mark_flags_as_required(['coco_root', 'output_dir'])309 app.run(main)310