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