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 16"""This file contains functions to post-process ViP-DeepLab results."""17 18import numpy as np19 20 21def stitch_video_panoptic_prediction(22    concat_panoptic: np.ndarray,23    next_panoptic: np.ndarray,24    label_divisor: int,25    overlap_offset: int = 128,26    combine_offset: int = 2 ** 32) -> np.ndarray:27  """The stitching algorithm in ViP-DeepLab.28 29  This function stitches a pair of image panoptic predictions to form video30  panoptic predictions by propagating instance IDs from concat_panoptic to31  next_panoptic based on IoU matching.32 33  Siyuan Qiao, Yukun Zhu, Hartwig Adam, Alan Yuille, and Liang-Chieh Chen.34  "ViP-DeepLab: Learning Visual Perception with Depth-aware Video Panoptic35  Segmentation." CVPR, 2021.36 37  Args:38    concat_panoptic: Panoptic prediction of the next frame by concatenating39      it with the current frame.40    next_panoptic: Panoptic prediction of the next frame.41    label_divisor: An integer specifying the label divisor of the dataset.42    overlap_offset: An integer offset to avoid overlap between the IDs in43      next_panoptic and the propagated IDs from concat_panoptic.44    combine_offset: An integer offset to combine concat and next panoptic.45 46  Returns:47    Panoptic prediction of the next frame with the instance IDs propragated48      from the concatenated panoptic prediction.49  """50  def _ids_to_counts(id_array: np.ndarray):51    """Given a numpy array, a mapping from each entry to its count."""52    ids, counts = np.unique(id_array, return_counts=True)53    return dict(zip(ids, counts))54  new_panoptic = next_panoptic.copy()55  # Increase the panoptic instance ID to avoid overlap.56  new_category = new_panoptic // label_divisor57  new_instance = new_panoptic % label_divisor58  # We skip 0 which is reserved for crowd.59  instance_mask = new_instance > 060  new_instance[instance_mask] = new_instance[instance_mask] + overlap_offset61  new_panoptic = new_category * label_divisor + new_instance62  # Pre-compute areas for all the segments.63  concat_segment_areas = _ids_to_counts(concat_panoptic)64  next_segment_areas = _ids_to_counts(next_panoptic)65  # Combine concat_panoptic and next_panoptic.66  intersection_id_array = (concat_panoptic.astype(np.int64) *67                           combine_offset + next_panoptic.astype(np.int64))68  intersection_areas = _ids_to_counts(intersection_id_array)69  # Compute IoU and sort them.70  intersection_ious = []71  for intersection_id, intersection_area in intersection_areas.items():72    concat_panoptic_label = int(intersection_id // combine_offset)73    next_panoptic_label = int(intersection_id % combine_offset)74    concat_category_label = concat_panoptic_label // label_divisor75    next_category_label = next_panoptic_label // label_divisor76    if concat_category_label != next_category_label:77      continue78    concat_instance_label = concat_panoptic_label % label_divisor79    next_instance_label = next_panoptic_label % label_divisor80    # We skip 0 which is reserved for crowd.81    if concat_instance_label == 0 or next_instance_label == 0:82      continue83    union = (84        concat_segment_areas[concat_panoptic_label] +85        next_segment_areas[next_panoptic_label] -86        intersection_area)87    iou = intersection_area / union88    intersection_ious.append([89        concat_panoptic_label, next_panoptic_label, iou])90  intersection_ious = sorted(91      intersection_ious, key=lambda e: e[2])92  # Build mapping and inverse mapping. Two-way mapping guarantees 1-to-193  # matching.94  map_concat_to_next = {}95  map_next_to_concat = {}96  for (concat_panoptic_label, next_panoptic_label,97         iou) in intersection_ious:98    map_concat_to_next[concat_panoptic_label] = next_panoptic_label99    map_next_to_concat[next_panoptic_label] = concat_panoptic_label100  # Match and propagate.101  for (concat_panoptic_label,102         next_panoptic_label) in map_concat_to_next.items():103    if map_next_to_concat[next_panoptic_label] == concat_panoptic_label:104      propagate_mask = next_panoptic == next_panoptic_label105      new_panoptic[propagate_mask] = concat_panoptic_label106  return new_panoptic107