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 16r"""Converts STEP (KITTI-STEP or MOTChallenge-STEP) data to sharded TFRecord file format with tf.train.Example protos.17 18The expected directory structure of the STEP dataset should be as follows:19 20  + {KITTI | MOTChallenge}-STEP21    + images22       + train23         + sequence_id24           - *.{png|jpg}25           ...26       + val27       + test28    + panoptic_maps29       + train30         + sequence_id31           - *.png32           ...33       + val34 35The ground-truth panoptic map is encoded as the following in PNG format:36 37  R: semantic_id38  G: instance_id // 25639  B: instance % 25640 41See ./utils/create_step_panoptic_maps.py for more details of how we create the42panoptic map by merging semantic and instance maps.43 44The output Example proto contains the following fields:45 46  image/encoded: encoded image content.47  image/filename: image filename.48  image/format: image file format.49  image/height: image height.50  image/width: image width.51  image/channels: image channels.52  image/segmentation/class/encoded: encoded panoptic segmentation content.53  image/segmentation/class/format: segmentation encoding format.54  video/sequence_id: sequence ID of the frame.55  video/frame_id: ID of the frame of the video sequence.56 57The output panoptic segmentation map stored in the Example will be the raw bytes58of an int32 panoptic map, where each pixel is assigned to a panoptic ID:59 60  panoptic ID = semantic ID * label divisor (1000) + instance ID61 62where semantic ID will be the same with `category_id` (use TrainId) for63each segment, and ignore label for pixels not belong to any segment.64 65The instance ID will be 0 for pixels belonging to66  1) `stuff` class67  2) `thing` class with `iscrowd` label68  3) pixels with ignore label69and [1, label divisor) otherwise.70 71Example to run the scipt:72 73   python deeplab2/data/build_step_data.py \74     --step_root=${STEP_ROOT} \75     --output_dir=${OUTPUT_DIR}76"""77 78import math79import os80 81from typing import Iterator, Sequence, Tuple, Optional82 83from absl import app84from absl import flags85from absl import logging86import numpy as np87 88from PIL import Image89 90import tensorflow as tf91 92from deeplab2.data import data_utils93 94FLAGS = flags.FLAGS95 96flags.DEFINE_string('step_root', None, 'STEP dataset root folder.')97 98flags.DEFINE_string('output_dir', None,99                    'Path to save converted TFRecord of TensorFlow examples.')100flags.DEFINE_bool(101    'use_two_frames', False, 'Flag to separate between 1 frame '102    'per TFExample or 2 consecutive frames per TFExample.')103 104_PANOPTIC_LABEL_FORMAT = 'raw'105_NUM_SHARDS = 10106_IMAGE_FOLDER_NAME = 'images'107_PANOPTIC_MAP_FOLDER_NAME = 'panoptic_maps'108_LABEL_MAP_FORMAT = 'png'109_INSTANCE_LABEL_DIVISOR = 1000110_ENCODED_INSTANCE_LABEL_DIVISOR = 256111_TF_RECORD_PATTERN = '%s-%05d-of-%05d.tfrecord'112_FRAME_ID_PATTERN = '%06d'113 114 115def _get_image_info_from_path(image_path: str) -> Tuple[str, str]:116  """Gets image info including sequence id and image id.117 118  Image path is in the format of '.../split/sequence_id/image_id.png',119  where `sequence_id` refers to the id of the video sequence, and `image_id` is120  the id of the image in the video sequence.121 122  Args:123    image_path: Absolute path of the image.124 125  Returns:126    sequence_id, and image_id as strings.127  """128  sequence_id = image_path.split('/')[-2]129  image_id = os.path.splitext(os.path.basename(image_path))[0]130  return sequence_id, image_id131 132 133def _get_images_per_shard(step_root: str, dataset_split: str,134                          sharded_by_sequence: bool) -> Iterator[Sequence[str]]:135  """Gets files for the specified data type and dataset split.136 137  Args:138    step_root: String, Path to STEP dataset root folder.139    dataset_split: String, dataset split ('train', 'val', 'test')140    sharded_by_sequence: Whether the images should be sharded by sequence or141      even split.142 143  Yields:144    A list of sorted file lists. Each inner list corresponds to one shard and is145    a list of files for this shard.146  """147  search_files = os.path.join(step_root, _IMAGE_FOLDER_NAME, dataset_split, '*',148                              '*')149  filenames = sorted(tf.io.gfile.glob(search_files))150  num_per_even_shard = int(math.ceil(len(filenames) / _NUM_SHARDS))151 152  sequence_ids = [os.path.basename(os.path.dirname(name)) for name in filenames]153  images_per_shard = []154  for i, name in enumerate(filenames):155    images_per_shard.append(name)156    shard_data = (i == len(filenames) - 1)157    # Sharded by sequence id.158    shard_data = shard_data or (sharded_by_sequence and159                                sequence_ids[i + 1] != sequence_ids[i])160    # Sharded evenly.161    shard_data = shard_data or (not sharded_by_sequence and162                                len(images_per_shard) == num_per_even_shard)163    if shard_data:164      yield images_per_shard165      images_per_shard = []166 167 168def _decode_panoptic_map(panoptic_map_path: str) -> Optional[str]:169  """Decodes the panoptic map from encoded image file.170 171  Args:172    panoptic_map_path: Path to the panoptic map image file.173 174  Returns:175    Panoptic map as an encoded int32 numpy array bytes or None if not existing.176  """177  if not tf.io.gfile.exists(panoptic_map_path):178    return None179  with tf.io.gfile.GFile(panoptic_map_path, 'rb') as f:180    panoptic_map = np.array(Image.open(f)).astype(np.int32)181  semantic_map = panoptic_map[:, :, 0]182  instance_map = (183      panoptic_map[:, :, 1] * _ENCODED_INSTANCE_LABEL_DIVISOR +184      panoptic_map[:, :, 2])185  panoptic_map = semantic_map * _INSTANCE_LABEL_DIVISOR + instance_map186  return panoptic_map.tobytes()187 188 189def _get_previous_frame_path(image_path: str) -> str:190  """Gets previous frame path. If not exists, duplicate it with image_path."""191  frame_id, frame_ext = os.path.splitext(os.path.basename(image_path))192  folder_dir = os.path.dirname(image_path)193  prev_frame_id = _FRAME_ID_PATTERN % (int(frame_id) - 1)194  prev_image_path = os.path.join(folder_dir, prev_frame_id + frame_ext)195  # If first frame, duplicates it.196  if not tf.io.gfile.exists(prev_image_path):197    tf.compat.v1.logging.warn(198        'Could not find previous frame %s of frame %d, duplicate the previous '199        'frame with the current frame.', prev_image_path, int(frame_id))200    prev_image_path = image_path201  return prev_image_path202 203 204def _create_panoptic_tfexample(image_path: str,205                               panoptic_map_path: str,206                               use_two_frames: bool,207                               is_testing: bool = False) -> tf.train.Example:208  """Creates a TF example for each image.209 210  Args:211    image_path: Path to the image.212    panoptic_map_path: Path to the panoptic map (as an image file).213    use_two_frames: Whether to encode consecutive two frames in the Example.214    is_testing: Whether it is testing data. If so, skip adding label data.215 216  Returns:217    TF example proto.218  """219  with tf.io.gfile.GFile(image_path, 'rb') as f:220    image_data = f.read()221  label_data = None222  if not is_testing:223    label_data = _decode_panoptic_map(panoptic_map_path)224  image_name = os.path.basename(image_path)225  image_format = image_name.split('.')[1].lower()226  sequence_id, frame_id = _get_image_info_from_path(image_path)227  prev_image_data = None228  prev_label_data = None229  if use_two_frames:230    # Previous image.231    prev_image_path = _get_previous_frame_path(image_path)232    with tf.io.gfile.GFile(prev_image_path, 'rb') as f:233      prev_image_data = f.read()234    # Previous panoptic map.235    if not is_testing:236      prev_panoptic_map_path = _get_previous_frame_path(panoptic_map_path)237      prev_label_data = _decode_panoptic_map(prev_panoptic_map_path)238  return data_utils.create_video_tfexample(239      image_data,240      image_format,241      image_name,242      label_format=_PANOPTIC_LABEL_FORMAT,243      sequence_id=sequence_id,244      image_id=frame_id,245      label_data=label_data,246      prev_image_data=prev_image_data,247      prev_label_data=prev_label_data)248 249 250def _convert_dataset(step_root: str,251                     dataset_split: str,252                     output_dir: str,253                     use_two_frames: bool = False):254  """Converts the specified dataset split to TFRecord format.255 256  Args:257    step_root: String, Path to STEP dataset root folder.258    dataset_split: String, the dataset split (e.g., train, val).259    output_dir: String, directory to write output TFRecords to.260    use_two_frames: Whether to encode consecutive two frames in the Example.261  """262  # For val and test set, if we run with use_two_frames, we should create a263  # sorted tfrecord per sequence.264  create_tfrecord_per_sequence = ('train'265                                  not in dataset_split) and use_two_frames266  is_testing = 'test' in dataset_split267 268  image_files_per_shard = list(269      _get_images_per_shard(step_root, dataset_split,270                            sharded_by_sequence=create_tfrecord_per_sequence))271  num_shards = len(image_files_per_shard)272 273  for shard_id, image_list in enumerate(image_files_per_shard):274    shard_filename = _TF_RECORD_PATTERN % (dataset_split, shard_id, num_shards)275    output_filename = os.path.join(output_dir, shard_filename)276    with tf.io.TFRecordWriter(output_filename) as tfrecord_writer:277      for image_path in image_list:278        sequence_id, image_id = _get_image_info_from_path(image_path)279        panoptic_map_path = os.path.join(280            step_root, _PANOPTIC_MAP_FOLDER_NAME, dataset_split, sequence_id,281            '%s.%s' % (image_id, _LABEL_MAP_FORMAT))282        example = _create_panoptic_tfexample(image_path, panoptic_map_path,283                                             use_two_frames, is_testing)284        tfrecord_writer.write(example.SerializeToString())285 286 287def main(argv: Sequence[str]) -> None:288  if len(argv) > 1:289    raise app.UsageError('Too many command-line arguments.')290  tf.io.gfile.makedirs(FLAGS.output_dir)291  for dataset_split in ('train', 'val', 'test'):292    logging.info('Starts to processing STEP dataset split %s.', dataset_split)293    _convert_dataset(FLAGS.step_root, dataset_split, FLAGS.output_dir,294                     FLAGS.use_two_frames)295 296 297if __name__ == '__main__':298  app.run(main)299