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 16r"""Creates STEP panoptic map from semantic and instance maps.17 18This script implements the process of merging semantic maps (from our extra19annotations[1]) and instance maps (collected from the MOTS[2]) to obtain the20STEP panoptic map.21 22[1] Mark Weber, etc. STEP: Segmenting and Tracking Every Pixel, arXiv:2102.1185923[2] Paul Voigtlaender, etc. Multi-object tracking and segmentation. CVPR, 201924 25To run this script, you need to install opencv-python (>=4.4.0).26e.g. In Linux, run27$pip install opencv-python28 29The input directory structure should be as follows:30 31+ INPUT_SEMANTIC_MAP_ROOT_DIR32 + train33 + sequence_id34 - *.png35 ...36 + val37 38+ INPUT_INSTANCE_MAP_ROOT_DIR39 + train40 + sequence_id41 - *.png42 ...43 + val44 45+ OUTPUT_PANOPTIC_MAP_ROOT_DIR (generated)46 + train47 + sequence_id48 - *.png49 ...50 + val51 52The ground-truth panoptic map is generated and encoded as the following in PNG53format:54 R: semantic_id55 G: instance_id // 25656 B: instance % 25657 58The generated panoptic maps will be used by ../build_step_data.py to create59tfrecords for training and evaluation.60 61Example to run the scipt:62 63```bash64 python deeplab2/data/utils/create_step_panoptic_maps.py \65 --input_semantic_map_root_dir=...66 ...67```68"""69 70import os71from typing import Any, Sequence, Union72 73from absl import app74from absl import flags75from absl import logging76import cv277import numpy as np78from PIL import Image79import tensorflow as tf80 81FLAGS = flags.FLAGS82flags.DEFINE_string('input_semantic_map_root_dir', None,83 'Path to a directory containing the semantic map.')84flags.DEFINE_string('input_instance_root_dir', None,85 'Path to a directory containing the instance map.')86flags.DEFINE_string('output_panoptic_map_root_dir', None,87 'Path to a directory where we write the panoptic map.')88flags.DEFINE_integer(89 'kernel_size', 15, 'Kernel size to extend instance object boundary when '90 'merging it with semantic map.')91flags.DEFINE_enum('dataset_name', 'kitti-step',92 ['kitti-step', 'motchallenge-step'], 'Name of the dataset')93 94# The label definition below follows Cityscapes label definition in95# https://www.cityscapes-dataset.com/.96MOTCHALLENGE_MERGED_CLASSES = (0, 3, 4, 5, 6, 7, 9, 13, 14, 15, 16, 17)97NUM_VALID_CLASSES = 1998SEMANTIC_CAR = 1399SEMANTIC_PERSON = 11100SEMANTIC_VOID = 255101INSTANCE_CAR = 1102INSTANCE_PERSON = 2103INSTANCE_LABEL_DIVISOR = 1000104 105 106def encode_panoptic_map(panoptic_map: np.ndarray) -> np.ndarray:107 """Encodes the panoptic map in three channel image format."""108 # Encoding format: R: semantic | G: instance // 256 | B: instance % 256109 semantic_id = panoptic_map // INSTANCE_LABEL_DIVISOR110 instance_id = panoptic_map % INSTANCE_LABEL_DIVISOR111 return np.dstack(112 (semantic_id, instance_id // 256, instance_id % 256)).astype(np.uint8)113 114 115def load_image(image_path: str) -> np.ndarray:116 """Loads an image as numpy array."""117 with tf.io.gfile.GFile(image_path, 'rb') as f:118 return np.array(Image.open(f))119 120 121def _update_motchallege_label_map(semantic_map: np.ndarray) -> np.ndarray:122 """Updates semantic map by merging some classes."""123 # For MOTChallenge dataset, we merge some classes since they are less124 # representative:125 #--------------------------------------------------------------126 # Original index | Updated index| Note127 #----------------+--------------+------------------------------128 # 0 | 1 | map road to sidewalk129 # 1 | 1 | keep sidewalk130 # 2 | 2 | keep building131 # 3 | 255 | not present anyway132 # 4 | 255 | remove fence133 # 5 | 255 | remove pole134 # 6 | 255 | remove traffic light135 # 7 | 255 | not present anyway136 # 8 | 8 | keep vegetation137 # 9 | 8 | map terrain to vegetation138 # 10 | 10 | keep sky139 # 11 | 11 | keep pedestrain140 # 12 | 12 | keep rider141 # 13 | 255 | remove car142 # 14 | 255 | not present anyway143 # 15 | 255 | not present anyway144 # 16 | 255 | not present anyway145 # 17 | 255 | remove motorcycle146 # 18 | 18 | keep bicycle147 # 255 | 255 | keep void148 #--------------------------------------------------------------149 for label in MOTCHALLENGE_MERGED_CLASSES:150 if label == 0:151 semantic_map[semantic_map == label] = 1152 elif label == 9:153 semantic_map[semantic_map == label] = 8154 else:155 semantic_map[semantic_map == label] = 255156 return semantic_map157 158 159def _compute_panoptic_id(semantic_id: Union[int, np.ndarray],160 instance_id: Union[int, np.ndarray]) -> Any:161 """Gets the panoptic id by combining semantic and instance id."""162 return semantic_id * INSTANCE_LABEL_DIVISOR + instance_id163 164 165def _remap_motchallege_semantic_indices(panoptic_id: np.ndarray) -> np.ndarray:166 """Updates MOTChallenge semantic map by re-mapping label indices."""167 semantic_id = panoptic_id // INSTANCE_LABEL_DIVISOR168 instance_id = panoptic_id % INSTANCE_LABEL_DIVISOR169 # Re-mapping index170 # 1 -> 0: sidewalk171 # 2 -> 1: building172 # 8 -> 2: vegetation173 # 10 -> 3: sky174 # 11 -> 4: pedestrain175 # 12 -> 5: rider176 # 18 -> 6: bicycle177 # 255 -> 255: void178 all_labels = set(range(NUM_VALID_CLASSES))179 for i, label in enumerate(180 sorted(all_labels - set(MOTCHALLENGE_MERGED_CLASSES))):181 semantic_id[semantic_id == label] = i182 return _compute_panoptic_id(semantic_id, instance_id)183 184 185def _get_semantic_maps(semantic_map_root: str, dataset_split: str,186 sequence_id: str) -> Sequence[str]:187 """Gets files for the specified data type and dataset split."""188 search_files = os.path.join(semantic_map_root, dataset_split, sequence_id,189 '*')190 filenames = tf.io.gfile.glob(search_files)191 return sorted(filenames)192 193 194class StepPanopticMapGenerator(object):195 """Class to generate and write panoptic map from semantic and instance map."""196 197 def __init__(self, kernel_size: int, dataset_name: str):198 self.kernel_size = kernel_size199 self.is_mots_challenge = (dataset_name == 'motchallenge-step')200 201 def _update_semantic_label_map(self, instance_map: np.ndarray,202 semantic_map: np.ndarray) -> np.ndarray:203 """Updates semantic map by leveraging semantic map and instance map."""204 kernel = np.ones((self.kernel_size, self.kernel_size), np.uint8)205 updated_semantic_map = semantic_map.astype(np.int32)206 if self.is_mots_challenge:207 updated_semantic_map = _update_motchallege_label_map(updated_semantic_map)208 for label in (SEMANTIC_CAR, SEMANTIC_PERSON):209 semantic_mask = (semantic_map == label)210 if label == SEMANTIC_PERSON:211 # The instance ids are encoded according to212 # https://www.vision.rwth-aachen.de/page/mots213 instance_mask = (214 instance_map // INSTANCE_LABEL_DIVISOR == INSTANCE_PERSON)215 elif label == SEMANTIC_CAR:216 instance_mask = instance_map // INSTANCE_LABEL_DIVISOR == INSTANCE_CAR217 # Run dilation on the instance map to merge it with semantic map.218 instance_mask = instance_mask.astype(np.uint8)219 dilated_instance_mask = cv2.dilate(instance_mask, kernel)220 void_boundary = np.logical_and(dilated_instance_mask - instance_mask,221 semantic_mask)222 updated_semantic_map[void_boundary] = SEMANTIC_VOID223 return updated_semantic_map224 225 def merge_panoptic_map(self, semantic_map: np.ndarray,226 instance_map: np.ndarray) -> np.ndarray:227 """Merges semantic labels with given instance map."""228 # Use semantic_map as the base map.229 updated_semantic_map = self._update_semantic_label_map(230 instance_map, semantic_map)231 panoptic_map = _compute_panoptic_id(updated_semantic_map, 0)232 # Merge instance.233 mask_car = instance_map // INSTANCE_LABEL_DIVISOR == INSTANCE_CAR234 # The instance map has index from 0 but the panoptic map's instance index235 # will start from 1.236 instance_id = (instance_map[mask_car] % INSTANCE_LABEL_DIVISOR) + 1237 panoptic_map[mask_car] = _compute_panoptic_id(SEMANTIC_CAR,238 instance_id.astype(np.int32))239 mask_person = instance_map // INSTANCE_LABEL_DIVISOR == INSTANCE_PERSON240 instance_id = (instance_map[mask_person] % INSTANCE_LABEL_DIVISOR) + 1241 panoptic_map[mask_person] = _compute_panoptic_id(242 SEMANTIC_PERSON, instance_id.astype(np.int32))243 244 # Remap label indices.245 if self.is_mots_challenge:246 panoptic_map = _remap_motchallege_semantic_indices(panoptic_map)247 return panoptic_map248 249 def build_panoptic_maps(self, semantic_map_root: str, instance_map_root: str,250 dataset_split: str, sequence_id: str,251 panoptic_map_root: str):252 """Creates panoptic maps and save them as PNG format.253 254 Args:255 semantic_map_root: Semantic map root folder.256 instance_map_root: Instance map root folder.257 dataset_split: Train/Val/Test split of the data.258 sequence_id: Sequence id of the data.259 panoptic_map_root: Panoptic map root folder where the encoded panoptic260 maps will be saved.261 """262 semantic_maps = _get_semantic_maps(semantic_map_root, dataset_split,263 sequence_id)264 for semantic_map_path in semantic_maps:265 image_name = os.path.basename(semantic_map_path)266 instance_map_path = os.path.join(instance_map_root, dataset_split,267 sequence_id, image_name)268 if not tf.io.gfile.exists(instance_map_path):269 logging.warn('Could not find instance map for %s', semantic_map_path)270 continue271 semantic_map = load_image(semantic_map_path)272 instance_map = load_image(instance_map_path)273 panoptic_map = self.merge_panoptic_map(semantic_map, instance_map)274 encoded_panoptic_map = Image.fromarray(275 encode_panoptic_map(panoptic_map)).convert('RGB')276 panoptic_map_path = os.path.join(panoptic_map_root, dataset_split,277 sequence_id, image_name)278 with tf.io.gfile.GFile(panoptic_map_path, 'wb') as f:279 encoded_panoptic_map.save(f, format='PNG')280 281 282def main(argv: Sequence[str]) -> None:283 if len(argv) > 1:284 raise app.UsageError('Too many command-line arguments.')285 286 panoptic_map_generator = StepPanopticMapGenerator(FLAGS.kernel_size,287 FLAGS.dataset_name)288 for dataset_split in ('train', 'val', 'test'):289 sem_dir = os.path.join(FLAGS.input_semantic_map_root_dir, dataset_split)290 if not tf.io.gfile.exists(sem_dir):291 logging.info('Split %s not found.', dataset_split)292 continue293 for set_dir in tf.io.gfile.listdir(sem_dir):294 tf.io.gfile.makedirs(295 os.path.join(FLAGS.output_panoptic_map_root_dir, dataset_split,296 set_dir))297 logging.info('Start to create panoptic map for split %s, sequence %s.',298 dataset_split, set_dir)299 panoptic_map_generator.build_panoptic_maps(300 FLAGS.input_semantic_map_root_dir, FLAGS.input_instance_root_dir,301 dataset_split, set_dir, FLAGS.output_panoptic_map_root_dir)302 303 304if __name__ == '__main__':305 app.run(main)306