Spaces:

akhaliq
/

deeplab2

Runtime error

deeplab2 / data /build_coco_data.py

akhaliq3

spaces demo

506da10 almost 3 years ago

No virus

11 kB

	# coding=utf-8
	# Copyright 2021 The Deeplab2 Authors.
	#
	# Licensed under the Apache License, Version 2.0 (the "License");
	# you may not use this file except in compliance with the License.
	# You may obtain a copy of the License at
	#
	# http://www.apache.org/licenses/LICENSE-2.0
	#
	# Unless required by applicable law or agreed to in writing, software
	# distributed under the License is distributed on an "AS IS" BASIS,
	# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	# See the License for the specific language governing permissions and
	# limitations under the License.

	"""Converts COCO data to sharded TFRecord file format with Example protos.

	Please check
	../g3doc/setup/coco.md
	for instructions.
	"""

	import collections
	import json
	import math
	import os

	from typing import Sequence, Tuple, Any

	from absl import app
	from absl import flags
	from absl import logging
	import numpy as np
	import tensorflow as tf

	from deeplab2.data import coco_constants
	from deeplab2.data import data_utils
	from deeplab2.data import dataset

	FLAGS = flags.FLAGS

	flags.DEFINE_string('coco_root', None, 'coco dataset root folder.')

	flags.DEFINE_string('output_dir', None,
	'Path to save converted TFRecord of TensorFlow examples.')

	flags.DEFINE_boolean('treat_crowd_as_ignore', True,
	'Whether to apply ignore labels to crowd pixels in '
	'panoptic label.')

	_NUM_SHARDS = 1000


	_SPLITS_TO_SIZES = dataset.COCO_PANOPTIC_INFORMATION.splits_to_sizes
	_IGNORE_LABEL = dataset.COCO_PANOPTIC_INFORMATION.ignore_label
	_CLASS_HAS_INSTANCE_LIST = dataset.COCO_PANOPTIC_INFORMATION.class_has_instances_list
	_PANOPTIC_LABEL_DIVISOR = dataset.COCO_PANOPTIC_INFORMATION.panoptic_label_divisor
	_CLASS_MAPPING = coco_constants.get_id_mapping()

	# A map from data type to folder name that saves the data.
	_FOLDERS_MAP = {
	'train': {
	'image': 'train2017',
	'label': 'annotations',
	},
	'val': {
	'image': 'val2017',
	'label': 'annotations',
	},
	'test': {
	'image': 'test2017',
	'label': '',
	}
	}

	# A map from data type to data format.
	_DATA_FORMAT_MAP = {
	'image': 'jpg',
	'label': 'png',
	}
	_PANOPTIC_LABEL_FORMAT = 'raw'


	def _get_images(coco_root: str, dataset_split: str) -> Sequence[str]:
	"""Gets files for the specified data type and dataset split.

	Args:
	coco_root: String, path to coco dataset root folder.
	dataset_split: String, dataset split ('train', 'val', 'test').

	Returns:
	A list of sorted file names.
	"""
	pattern = '*.%s' % _DATA_FORMAT_MAP['image']
	search_files = os.path.join(
	coco_root, _FOLDERS_MAP[dataset_split]['image'], pattern)
	filenames = tf.io.gfile.glob(search_files)
	return sorted(filenames)


	def _get_panoptic_annotation(coco_root: str, dataset_split: str,
	annotation_file_name: str) -> str:
	panoptic_folder = 'panoptic_%s2017' % dataset_split
	return os.path.join(coco_root, _FOLDERS_MAP[dataset_split]['label'],
	panoptic_folder, annotation_file_name)


	def _read_segments(coco_root: str, dataset_split: str):
	"""Reads segments information from json file.

	Args:
	coco_root: String, path to coco dataset root folder.
	dataset_split: String, dataset split.

	Returns:
	segments_dict: A dictionary that maps file prefix of annotation_file_name to
	a tuple of (panoptic annotation file name, segments). Please refer to
	_generate_panoptic_label() method on the detail structure of `segments`.

	Raises:
	ValueError: If found duplicated image id in annotations.
	"""
	json_filename = os.path.join(
	coco_root, _FOLDERS_MAP[dataset_split]['label'],
	'panoptic_%s2017.json' % dataset_split)
	with tf.io.gfile.GFile(json_filename) as f:
	panoptic_dataset = json.load(f)

	segments_dict = {}
	for annotation in panoptic_dataset['annotations']:
	image_id = annotation['image_id']
	if image_id in segments_dict:
	raise ValueError('Image ID %s already exists' % image_id)
	annotation_file_name = annotation['file_name']
	segments = annotation['segments_info']

	segments_dict[os.path.splitext(annotation_file_name)[-2]] = (
	annotation_file_name, segments)

	return segments_dict


	def _generate_panoptic_label(panoptic_annotation_file: str, segments:
	Any) -> np.ndarray:
	"""Creates panoptic label map from annotations.

	Args:
	panoptic_annotation_file: String, path to panoptic annotation.
	segments: A list of dictionaries containing information of every segment.
	Read from panoptic_${DATASET_SPLIT}2017.json. This method consumes
	the following fields in each dictionary:
	- id: panoptic id
	- category_id: semantic class id
	- area: pixel area of this segment
	- iscrowd: if this segment is crowd region

	Returns:
	A 2D numpy int32 array with the same height / width with panoptic
	annotation. Each pixel value represents its panoptic ID. Please refer to
	g3doc/setup/coco.md for more details about how panoptic ID is assigned.
	"""
	with tf.io.gfile.GFile(panoptic_annotation_file, 'rb') as f:
	panoptic_label = data_utils.read_image(f.read())

	if panoptic_label.mode != 'RGB':
	raise ValueError('Expect RGB image for panoptic label, gets %s' %
	panoptic_label.mode)

	panoptic_label = np.array(panoptic_label, dtype=np.int32)
	# COCO panoptic map is created by:
	# color = [segmentId % 256, segmentId // 256, segmentId // 256 // 256]
	panoptic_label = np.dot(panoptic_label, [1, 256, 256 * 256])

	semantic_label = np.ones_like(panoptic_label) * _IGNORE_LABEL
	instance_label = np.zeros_like(panoptic_label)
	# Running count of instances per semantic category.
	instance_count = collections.defaultdict(int)

	for segment in segments:
	selected_pixels = panoptic_label == segment['id']
	pixel_area = np.sum(selected_pixels)
	if pixel_area != segment['area']:
	raise ValueError('Expect %d pixels for segment %s, gets %d.' %
	(segment['area'], segment, pixel_area))

	category_id = segment['category_id']

	# Map the category_id to contiguous ids
	category_id = _CLASS_MAPPING[category_id]

	semantic_label[selected_pixels] = category_id

	if category_id in _CLASS_HAS_INSTANCE_LIST:
	if segment['iscrowd']:
	# COCO crowd pixels will have instance ID of 0.
	if FLAGS.treat_crowd_as_ignore:
	semantic_label[selected_pixels] = _IGNORE_LABEL
	continue
	# Non-crowd pixels will have instance ID starting from 1.
	instance_count[category_id] += 1
	if instance_count[category_id] >= _PANOPTIC_LABEL_DIVISOR:
	raise ValueError('Too many instances for category %d in this image.' %
	category_id)
	instance_label[selected_pixels] = instance_count[category_id]
	elif segment['iscrowd']:
	raise ValueError('Stuff class should not have `iscrowd` label.')

	panoptic_label = semantic_label * _PANOPTIC_LABEL_DIVISOR + instance_label
	return panoptic_label.astype(np.int32)


	def _create_panoptic_label(coco_root: str, dataset_split: str, image_path: str,
	segments_dict: Any
	) -> Tuple[str, str]:
	"""Creates labels for panoptic segmentation.

	Args:
	coco_root: String, path to coco dataset root folder.
	dataset_split: String, dataset split ('train', 'val', 'test').
	image_path: String, path to the image file.
	segments_dict:
	Read from panoptic_${DATASET_SPLIT}2017.json. This method consumes
	the following fields in each dictionary:
	- id: panoptic id
	- category_id: semantic class id
	- area: pixel area of this segment
	- iscrowd: if this segment is crowd region

	Returns:
	A panoptic label where each pixel value represents its panoptic ID.
	Please refer to g3doc/setup/coco.md for more details about howpanoptic ID
	is assigned.
	A string indicating label format in TFRecord.
	"""

	image_path = os.path.normpath(image_path)
	path_list = image_path.split(os.sep)
	file_name = path_list[-1]

	annotation_file_name, segments = segments_dict[
	os.path.splitext(file_name)[-2]]
	panoptic_annotation_file = _get_panoptic_annotation(coco_root,
	dataset_split,
	annotation_file_name)

	panoptic_label = _generate_panoptic_label(panoptic_annotation_file, segments)
	return panoptic_label.tostring(), _PANOPTIC_LABEL_FORMAT


	def _convert_dataset(coco_root: str, dataset_split: str,
	output_dir: str) -> None:
	"""Converts the specified dataset split to TFRecord format.

	Args:
	coco_root: String, path to coco dataset root folder.
	dataset_split: String, the dataset split (one of `train`, `val` and `test`).
	output_dir: String, directory to write output TFRecords to.
	"""
	image_files = _get_images(coco_root, dataset_split)

	num_images = len(image_files)

	if dataset_split != 'test':
	segments_dict = _read_segments(coco_root, dataset_split)

	num_per_shard = int(math.ceil(len(image_files) / _NUM_SHARDS))

	for shard_id in range(_NUM_SHARDS):
	shard_filename = '%s-%05d-of-%05d.tfrecord' % (
	dataset_split, shard_id, _NUM_SHARDS)
	output_filename = os.path.join(output_dir, shard_filename)
	with tf.io.TFRecordWriter(output_filename) as tfrecord_writer:
	start_idx = shard_id * num_per_shard
	end_idx = min((shard_id + 1) * num_per_shard, num_images)
	for i in range(start_idx, end_idx):
	# Read the image.
	with tf.io.gfile.GFile(image_files[i], 'rb') as f:
	image_data = f.read()

	if dataset_split == 'test':
	label_data, label_format = None, None
	else:
	label_data, label_format = _create_panoptic_label(
	coco_root, dataset_split, image_files[i], segments_dict)

	# Convert to tf example.
	image_path = os.path.normpath(image_files[i])
	path_list = image_path.split(os.sep)
	file_name = path_list[-1]
	file_prefix = file_name.replace(_DATA_FORMAT_MAP['image'], '')
	example = data_utils.create_tfexample(image_data,
	'jpeg',
	file_prefix, label_data,
	label_format)

	tfrecord_writer.write(example.SerializeToString())


	def main(unused_argv: Sequence[str]) -> None:
	tf.io.gfile.makedirs(FLAGS.output_dir)

	for dataset_split in ('train', 'val', 'test'):
	logging.info('Starts processing dataset split %s.', dataset_split)
	_convert_dataset(FLAGS.coco_root, dataset_split, FLAGS.output_dir)


	if __name__ == '__main__':
	flags.mark_flags_as_required(['coco_root', 'output_dir'])
	app.run(main)