karolmajek/maxdeeplab

1// Copyright 2021 The Deeplab2 Authors.2//3// Licensed under the Apache License, Version 2.0 (the "License");4// you may not use this file except in compliance with the License.5// You may obtain a copy of the License at6//7//     http://www.apache.org/licenses/LICENSE-2.08//9// Unless required by applicable law or agreed to in writing, software10// distributed under the License is distributed on an "AS IS" BASIS,11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12// See the License for the specific language governing permissions and13// limitations under the License.14 15#include <cstdint>16#define EIGEN_USE_THREADS17 18#define _USE_MATH_DEFINES19 20#include <algorithm>21#include <iterator>22#include <set>23#include <unordered_map>24#include <unordered_set>25#include <vector>26 27#include /*third_party*/"tensorflow/core/framework/op_kernel.h"28#include /*third_party*/"tensorflow/core/framework/register_types.h"29#include /*third_party*/"tensorflow/core/framework/tensor.h"30#include /*third_party*/"tensorflow/core/framework/tensor_shape.h"31#include /*third_party*/"tensorflow/core/framework/types.h"32#include /*third_party*/"tensorflow/core/lib/core/errors.h"33#include /*third_party*/"tensorflow/core/lib/core/status.h"34#include /*third_party*/"tensorflow/core/platform/logging.h"35#include /*third_party*/"merge_semantic_and_instance_maps_op_kernel.h" // local headers36 37namespace tensorflow_models {38namespace deeplab {39namespace deeplab2 {40 41namespace {42 43using tensorflow::Tensor;44using tensorflow::TensorShape;45using tensorflow::TTypes;46using tensorflow::errors::InvalidArgument;47 48}  // namespace49 50namespace functor {51 52// This function merges the semantic segmentation and class-agnostic53// instance segmentation to form the panoptic segmentation. In particular,54// the class label of each instance mask is inferred from the majority55// votes from the corresponding pixels in the semantic segmentation. This56// operation is first poposed in the DeeperLab paper and adopted by the57// Panoptic-DeepLab.58// - DeeperLab: Single-Shot Image Parser, T-J Yang, et al. arXiv:1902.05093.59// - Panoptic-DeepLab, B. Cheng, et al. In CVPR, 2020.60// Specialization of MergeSemanticAndInstanceMaps< for CPU.61template <>62void MergeSemanticAndInstanceMaps<Eigen::ThreadPoolDevice>::operator()(63    const Eigen::ThreadPoolDevice& d,64    typename TTypes<int32_t, 3>::ConstTensor semantic_maps,65    typename TTypes<int32_t, 3>::ConstTensor instance_maps,66    const std::unordered_set<int32_t>& thing_ids_set, int label_divisor,67    int stuff_area_limit, int void_label,68    typename TTypes<int32_t, 3>::Tensor parsing_maps) {69  const int num_batches = semantic_maps.dimension(0);70  const int height = semantic_maps.dimension(1);71  const int width = semantic_maps.dimension(2);72 73  for (int b = 0; b < num_batches; ++b) {74    // A vector to keep track of which pixels are predicted as `thing` or75    // `stuff` class.76    std::vector<bool> is_thing(height * width, true);77 78    // For each instance, find its corresponding histogram of semantic labels.79    // Suppose car label = 2 and road label = 5, and predicted instance 3 has80    // 5 pixels predicted as car and 20 pixels predicted as road. Then,81    // instance_id_to_semantic_histogram[3][2] = 5 and82    // instance_id_to_semantic_histogram[3][5] = 20.83    using InstanceIdType = int32_t;84    using SemanticLabelType = int32_t;85    using CountsType = int32_t;86    std::unordered_map<InstanceIdType,87                       std::unordered_map<SemanticLabelType, CountsType>>88        instance_id_to_semantic_histogram;89    // A map from stuff label to area.90    std::unordered_map<SemanticLabelType, CountsType> stuff_label_to_area;91    for (int h = 0; h < height; ++h) {92      for (int w = 0; w < width; ++w) {93        const int semantic_val = semantic_maps(b, h, w);94        if (thing_ids_set.find(semantic_val) == thing_ids_set.end()) {95          // Skip if it is `stuff`.96          is_thing[w + width * h] = false;97          ++stuff_label_to_area[semantic_val];98          continue;99        }100        const int instance_val = instance_maps(b, h, w);101        ++instance_id_to_semantic_histogram[instance_val][semantic_val];102      }103    }104    // Keep track of how many instances for each semantic_label.105    std::unordered_map<SemanticLabelType, CountsType>106        semantic_label_to_instance_counts;107    // Find the new semantic label and instance id for each instance. We use108    // majority vote to find the new semantic label while reorder the instance109    // id in the following way. In the original instance map, every instance110    // has a different instance id. In the new instance map, every instance111    // `in the same semantic class` should have a different id, but instances112    // `in different semantic classes` can have the same instance id. This113    // reduces the maximum instance label value and avoids the problem of114    // combining the two maps with the label_divisor.115    std::unordered_map<InstanceIdType,116                       std::pair<SemanticLabelType, InstanceIdType>>117        instance_id_to_new_semantic_label_and_instance_id;118    for (const auto& instance_to_histogram :119         instance_id_to_semantic_histogram) {120      const int instance_val = instance_to_histogram.first;121      const std::unordered_map<SemanticLabelType, CountsType>122          semantic_histogram = instance_to_histogram.second;123      int semantic_label = -1;124      int max_count = 0;125      // Find the majority semantic label.126      for (const auto& semantic_to_count : semantic_histogram) {127        // Break ties deterministically by select the smaller semantic label.128        if (semantic_to_count.second > max_count ||129            (semantic_to_count.second == max_count &&130             semantic_to_count.first < semantic_label)) {131          max_count = semantic_to_count.second;132          semantic_label = semantic_to_count.first;133        }134      }135      ++semantic_label_to_instance_counts[semantic_label];136      // For `thing` class, we set instance id starting from 1, while for137      // `stuff` class, we use instance id 0.138      instance_id_to_new_semantic_label_and_instance_id[instance_val] = {139          semantic_label, semantic_label_to_instance_counts[semantic_label]};140    }141    // Create a new semantic map by assigning the majority semantic label for142    // each instance.143    std::vector<SemanticLabelType> semantic_map(height * width);144    // Create a new instance map by assigning ordered instance id's.145    std::vector<InstanceIdType> instance_map(height * width);146    for (int h = 0; h < height; ++h) {147      for (int w = 0; w < width; ++w) {148        const int pixel = w + width * h;149        if (is_thing[pixel]) {150          const int instance_val = instance_maps(b, h, w);151          // Assign the majority semantic vote in the new semantic map, and152          // reorder the instance id in the new instance map.153          std::tie(semantic_map[pixel], instance_map[pixel]) =154              instance_id_to_new_semantic_label_and_instance_id[instance_val];155        } else {156          // If current pixel belongs to `stuff` class, keep the same semantic157          // label in the new semantic map. We also check if its area is158          // smaller than the stuff_area_limit_ or not. If true, we re-assign159          // the segment with void_label_.160          const int semantic_val = semantic_maps(b, h, w);161          if (stuff_area_limit > 0 &&162              stuff_label_to_area[semantic_val] <= stuff_area_limit) {163            semantic_map[pixel] = void_label;164          } else {165            semantic_map[pixel] = semantic_val;166          }167          // If current pixel belongs to `stuff` class, assign 0 in the new168          // instance map.169          instance_map[pixel] = 0;170        }171      }172    }173    // Merge those semantic map and instance map.174    for (int h = 0; h < height; ++h) {175      for (int w = 0; w < width; ++w) {176        const int pixel = w + width * h;177        parsing_maps(b, h, w) =178            semantic_map[pixel] * label_divisor + instance_map[pixel];179      }180    }181  }182}183 184template <>185std::unordered_set<int32_t> Convert1DInt32TensorToSet(186    const Eigen::ThreadPoolDevice& d, const Tensor& tensor) {187  std::unordered_set<int32_t> target_set;188  const int n_vals = tensor.dim_size(0);189  typename TTypes<int32_t, 1>::ConstTensor tensor_data =190      tensor.tensor<int32_t, 1>();191  for (int i = 0; i < n_vals; i++) {192    target_set.insert(tensor_data(i));193  }194 195  return target_set;196}197 198}  // namespace functor199 200template <typename Device>201class MergeSemanticAndInstanceMapsOp : public tensorflow::OpKernel {202 public:203  explicit MergeSemanticAndInstanceMapsOp(204      tensorflow::OpKernelConstruction* context)205      : OpKernel(context) {206    OP_REQUIRES_OK(context, context->GetAttr("label_divisor", &label_divisor_));207    OP_REQUIRES(context, label_divisor_ > 0,208                InvalidArgument("Label divisor must be positive."));209    OP_REQUIRES_OK(context,210                   context->GetAttr("stuff_area_limit", &stuff_area_limit_));211    OP_REQUIRES(context, stuff_area_limit_ >= 0,212                InvalidArgument("Stuff area limit must be non-negative."));213    OP_REQUIRES_OK(context, context->GetAttr("void_label", &void_label_));214    OP_REQUIRES(context, void_label_ >= 0,215                InvalidArgument("Void label must be non-negative."));216  }217 218  void Compute(tensorflow::OpKernelContext* context) override {219    // Extract the inputs.220    const Tensor& semantic_maps = context->input(0);221    const Tensor& instance_maps = context->input(1);222    const Tensor& thing_ids_tensor = context->input(2);223 224    // Convert thing_ids_tensor into a set.225    std::unordered_set<int32_t> thing_ids_set =226        functor::Convert1DInt32TensorToSet(context->eigen_device<Device>(),227                                           thing_ids_tensor);228 229    // Extract the constants.230    const int batch = semantic_maps.dim_size(0);231    const int height = semantic_maps.dim_size(1);232    const int width = semantic_maps.dim_size(2);233 234    // Check input shapes.235    OP_REQUIRES(context,236                instance_maps.dim_size(0) == batch &&237                    instance_maps.dim_size(1) == height &&238                    instance_maps.dim_size(2) == width,239                InvalidArgument(240                    "Expect semantic and instance maps have the same shape.",241                    instance_maps.shape().DebugString()));242 243    Tensor* parsing_maps = nullptr;244    OP_REQUIRES_OK(context,245                   context->allocate_output(246                       0, TensorShape({batch, height, width}), &parsing_maps));247 248    functor::MergeSemanticAndInstanceMaps<Device>()(249        context->eigen_device<Device>(), semantic_maps.tensor<int32_t, 3>(),250        instance_maps.tensor<int32_t, 3>(), thing_ids_set, label_divisor_,251        stuff_area_limit_, void_label_, parsing_maps->tensor<int32_t, 3>());252  }253 254 private:255  // Label divisor, the value used to combine the semantic and instance map to256  // generate the parsing map.257  int label_divisor_;258 259  // Stuff area limit is used to remove predicted stuff segments whose area are260  // smaller than it.261  int stuff_area_limit_;262 263  // Removed predicted stuff segments are re-assigned with void label.264  int void_label_;265};266 267REGISTER_KERNEL_BUILDER(268    Name("MergeSemanticAndInstanceMaps").Device(tensorflow::DEVICE_CPU),269    MergeSemanticAndInstanceMapsOp<Eigen::ThreadPoolDevice>);270 271#ifdef GOOGLE_CUDA272REGISTER_KERNEL_BUILDER(273    Name("MergeSemanticAndInstanceMaps").Device(tensorflow::DEVICE_GPU),274    MergeSemanticAndInstanceMapsOp<Eigen::GpuDevice>)275#endif  // GOOGLE_CUDA276 277}  // namespace deeplab2278}  // namespace deeplab279}  // namespace tensorflow_models280