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 16r"""Converts Depth-aware Video Panoptic Segmentation (DVPS) data to sharded TFRecord file format with tf.train.Example protos.17 18The expected directory structure of the DVPS dataset should be as follows:19 20  + DVPS_ROOT21    + train | val22      - ground-truth depth maps (*_depth.png)23      - ground-truth panoptic maps (*_gtFine_instanceTrainIds.png)24      - images (*_leftImg8bit.png)25    + test26      - images (*_leftImg8bit.png)27 28The ground-truth panoptic map is encoded as the following in PNG format:29 30  panoptic ID = semantic ID * panoptic divisor (1000) + instance ID31 32 33The output Example proto contains the following fields:34 35  image/encoded: encoded image content.36  image/filename: image filename.37  image/format: image file format.38  image/height: image height.39  image/width: image width.40  image/channels: image channels.41  image/segmentation/class/encoded: encoded panoptic segmentation content.42  image/segmentation/class/format: segmentation encoding format.43  image/depth/encoded: encoded depth content.44  image/depth/format: depth encoding format.45  video/sequence_id: sequence ID of the frame.46  video/frame_id: ID of the frame of the video sequence.47  next_image/encoded: encoded next-frame image content.48  next_image/segmentation/class/encoded: encoded panoptic segmentation content49    of the next frame.50 51The output panoptic segmentation map stored in the Example will be the raw bytes52of an int32 panoptic map, where each pixel is assigned to a panoptic ID:53 54  panoptic ID = semantic ID * panoptic divisor (1000) + instance ID55 56where semantic ID will be the same with `category_id` for each segment, and57ignore label for pixels not belong to any segment.58 59The depth map will be the raw bytes of an int32 depth map, where each pixel is:60 61  depth map = depth ground truth * 25662 63Example to run the scipt:64 65   python deeplab2/data/build_dvps_data.py \66     --dvps_root=${DVPS_ROOT} \67     --output_dir=${OUTPUT_DIR}68"""69 70import math71import os72 73from typing import Sequence, Tuple, Optional74 75from absl import app76from absl import flags77from absl import logging78import numpy as np79 80from PIL import Image81 82import tensorflow as tf83 84from deeplab2.data import data_utils85 86FLAGS = flags.FLAGS87 88flags.DEFINE_string('dvps_root', None, 'DVPS dataset root folder.')89 90flags.DEFINE_string('output_dir', None,91                    'Path to save converted TFRecord of TensorFlow examples.')92 93_PANOPTIC_DEPTH_FORMAT = 'raw'94_NUM_SHARDS = 100095_TF_RECORD_PATTERN = '%s-%05d-of-%05d.tfrecord'96_IMAGE_SUFFIX = '_leftImg8bit.png'97_LABEL_SUFFIX = '_gtFine_instanceTrainIds.png'98_DEPTH_SUFFIX = '_depth.png'99 100 101def _get_image_info_from_path(image_path: str) -> Tuple[str, str]:102  """Gets image info including sequence id and image id.103 104  Image path is in the format of '{sequence_id}_{image_id}_*.png',105  where `sequence_id` refers to the id of the video sequence, and `image_id` is106  the id of the image in the video sequence.107 108  Args:109    image_path: Absolute path of the image.110 111  Returns:112    sequence_id, and image_id as strings.113  """114  image_path = os.path.basename(image_path)115  return tuple(image_path.split('_')[:2])116 117 118def _get_images(dvps_root: str, dataset_split: str) -> Sequence[str]:119  """Gets files for the specified data type and dataset split.120 121  Args:122    dvps_root: String, path to DVPS dataset root folder.123    dataset_split: String, dataset split ('train', 'val', 'test').124 125  Returns:126    A list of sorted file names under dvps_root and dataset_split.127  """128  search_files = os.path.join(dvps_root, dataset_split, '*' + _IMAGE_SUFFIX)129  filenames = tf.io.gfile.glob(search_files)130  return sorted(filenames)131 132 133def _decode_panoptic_or_depth_map(map_path: str) -> Optional[str]:134  """Decodes the panoptic or depth map from encoded image file.135 136  Args:137    map_path: Path to the panoptic or depth map image file.138 139  Returns:140    Panoptic or depth map as an encoded int32 numpy array bytes or None if not141      existing.142  """143  if not tf.io.gfile.exists(map_path):144    return None145  with tf.io.gfile.GFile(map_path, 'rb') as f:146    decoded_map = np.array(Image.open(f)).astype(np.int32)147  return decoded_map.tobytes()148 149 150def _get_next_frame_path(image_path: str) -> Optional[str]:151  """Gets next frame path.152 153  If not exists, return None.154 155  The files are named {sequence_id}_{frame_id}*. To get the path of the next156  frame, this function keeps sequence_id and increase the frame_id by 1. It157  finds all the files matching this pattern, and returns the corresponding158  file path matching the input type.159 160  Args:161    image_path: String, path to the image.162 163  Returns:164    A string for the path of the next frame of the given image path or None if165      the given image path is the last frame of the sequence.166  """167  sequence_id, image_id = _get_image_info_from_path(image_path)168  next_image_id = '{:06d}'.format(int(image_id) + 1)169  next_image_name = sequence_id + '_' + next_image_id170  next_image_path = None171  for suffix in (_IMAGE_SUFFIX, _LABEL_SUFFIX):172    if image_path.endswith(suffix):173      next_image_path = os.path.join(174          os.path.dirname(image_path), next_image_name + suffix)175      if not tf.io.gfile.exists(next_image_path):176        return None177  return next_image_path178 179 180def _create_tfexample(image_path: str, panoptic_map_path: str,181                      depth_map_path: str) -> Optional[tf.train.Example]:182  """Creates a TF example for each image.183 184  Args:185    image_path: Path to the image.186    panoptic_map_path: Path to the panoptic map (as an image file).187    depth_map_path: Path to the depth map (as an image file).188 189  Returns:190    TF example proto.191  """192  with tf.io.gfile.GFile(image_path, 'rb') as f:193    image_data = f.read()194  label_data = _decode_panoptic_or_depth_map(panoptic_map_path)195  depth_data = _decode_panoptic_or_depth_map(depth_map_path)196  image_name = os.path.basename(image_path)197  image_format = image_name.split('.')[1].lower()198  sequence_id, frame_id = _get_image_info_from_path(image_path)199  next_image_data = None200  next_label_data = None201  # Next image.202  next_image_path = _get_next_frame_path(image_path)203  # If there is no next image, no examples will be created.204  if next_image_path is None:205    return None206  with tf.io.gfile.GFile(next_image_path, 'rb') as f:207    next_image_data = f.read()208  # Next panoptic map.209  next_panoptic_map_path = _get_next_frame_path(panoptic_map_path)210  next_label_data = _decode_panoptic_or_depth_map(next_panoptic_map_path)211  return data_utils.create_video_and_depth_tfexample(212      image_data,213      image_format,214      image_name,215      label_format=_PANOPTIC_DEPTH_FORMAT,216      sequence_id=sequence_id,217      image_id=frame_id,218      label_data=label_data,219      next_image_data=next_image_data,220      next_label_data=next_label_data,221      depth_data=depth_data,222      depth_format=_PANOPTIC_DEPTH_FORMAT)223 224 225def _convert_dataset(dvps_root: str, dataset_split: str, output_dir: str):226  """Converts the specified dataset split to TFRecord format.227 228  Args:229    dvps_root: String, path to DVPS dataset root folder.230    dataset_split: String, the dataset split (e.g., train, val, test).231    output_dir: String, directory to write output TFRecords to.232  """233  image_files = _get_images(dvps_root, dataset_split)234  num_images = len(image_files)235 236  num_per_shard = int(math.ceil(len(image_files) / _NUM_SHARDS))237 238  for shard_id in range(_NUM_SHARDS):239    shard_filename = _TF_RECORD_PATTERN % (dataset_split, shard_id, _NUM_SHARDS)240    output_filename = os.path.join(output_dir, shard_filename)241    with tf.io.TFRecordWriter(output_filename) as tfrecord_writer:242      start_idx = shard_id * num_per_shard243      end_idx = min((shard_id + 1) * num_per_shard, num_images)244      for i in range(start_idx, end_idx):245        image_path = image_files[i]246        panoptic_map_path = image_path.replace(_IMAGE_SUFFIX, _LABEL_SUFFIX)247        depth_map_path = image_path.replace(_IMAGE_SUFFIX, _DEPTH_SUFFIX)248        example = _create_tfexample(image_path, panoptic_map_path,249                                    depth_map_path)250        if example is not None:251          tfrecord_writer.write(example.SerializeToString())252 253 254def main(argv: Sequence[str]) -> None:255  if len(argv) > 1:256    raise app.UsageError('Too many command-line arguments.')257  tf.io.gfile.makedirs(FLAGS.output_dir)258  for dataset_split in ('train', 'val', 'test'):259    logging.info('Starts to processing DVPS dataset split %s.', dataset_split)260    _convert_dataset(FLAGS.dvps_root, dataset_split, FLAGS.output_dir)261 262 263if __name__ == '__main__':264  app.run(main)265