karolmajek/Axial-DeepLab-SWideRNet
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 16"""This file contains code to get a sample from a dataset."""17 18import functools19 20import numpy as np21import tensorflow as tf22 23from deeplab2 import common24from deeplab2.data import dataset_utils25from deeplab2.data.preprocessing import input_preprocessing as preprocessing26 27 28def _compute_gaussian_from_std(sigma):29 """Computes the Gaussian and its size from a given standard deviation."""30 size = int(6 * sigma + 3)31 x = np.arange(size, dtype=np.float)32 y = x[:, np.newaxis]33 x0, y0 = 3 * sigma + 1, 3 * sigma + 134 return np.exp(-((x - x0)**2 + (y - y0)**2) / (2 * sigma**2)), size35 36 37class PanopticSampleGenerator:38 """This class generates samples from images and labels."""39 40 def __init__(self,41 dataset_info,42 is_training,43 crop_size,44 min_resize_value=None,45 max_resize_value=None,46 resize_factor=None,47 min_scale_factor=1.,48 max_scale_factor=1.,49 scale_factor_step_size=0,50 autoaugment_policy_name=None,51 only_semantic_annotations=False,52 thing_id_mask_annotations=False,53 max_thing_id=128,54 sigma=8,55 focus_small_instances=None):56 """Initializes the panoptic segmentation generator.57 58 Args:59 dataset_info: A dictionary with the following keys.60 - `name`: String, dataset name.61 - `ignore_label`: Integer, ignore label.62 - `class_has_instances_list`: A list of integers indicating which63 class has instance annotations.64 - `panoptic_label_divisor`: Integer, panoptic label divisor.65 - `num_classes`: Integer, number of classes.66 - `is_video_dataset`: Boolean, is video dataset or not.67 is_training: Boolean, is training mode or not.68 crop_size: Image crop size [height, width].69 min_resize_value: A 2-tuple of (height, width), desired minimum value70 after resize. If a single element is given, then height and width share71 the same value. None, empty or having 0 indicates no minimum value will72 be used.73 max_resize_value: A 2-tuple of (height, width), maximum allowed value74 after resize. If a single element is given, then height and width75 share the same value. None, empty or having 0 indicates no maximum76 value will be used.77 resize_factor: Resized dimensions are multiple of factor plus one.78 min_scale_factor: Minimum scale factor for random scale augmentation.79 max_scale_factor: Maximum scale factor for random scale augmentation.80 scale_factor_step_size: The step size from min scale factor to max scale81 factor. The input is randomly scaled based on the value of82 (min_scale_factor, max_scale_factor, scale_factor_step_size).83 autoaugment_policy_name: String, autoaugment policy name. See84 autoaugment_policy.py for available policies.85 only_semantic_annotations: An optional flag indicating whether the model86 needs only semantic annotations (default: False).87 thing_id_mask_annotations: An optional flag indicating whether the model88 needs thing_id_mask annotations. When `thing_id_mask_annotations` is89 True, we will additionally return mask annotation for each `thing`90 instance, encoded with a unique thing_id. This ground-truth annotation91 could be used to learn a better segmentation mask for each instance.92 `thing_id` indicates the number of unique thing-ID to each instance in93 an image, starting the counting from 0 (default: False).94 max_thing_id: The maximum number of possible thing instances per image. It95 is used together when thing_id_mask_annotations = True, representing the96 maximum thing ID encoded in the thing_id_mask. (default: 128).97 sigma: The standard deviation of the Gaussian used to encode the center98 keypoint (default: 8).99 focus_small_instances: An optional dict that defines how to deal with100 small instances (default: None):101 -`threshold`: An integer defining the threshold pixel number for an102 instance to be considered small.103 -`weight`: A number that defines the loss weight for small instances.104 """105 self._dataset_info = dataset_info106 self._ignore_label = self._dataset_info['ignore_label']107 self._only_semantic_annotations = only_semantic_annotations108 self._sigma = sigma109 self._instance_area_threshold = 0110 self._small_instance_weight = 1.0111 self._thing_id_mask_annotations = thing_id_mask_annotations112 self._max_thing_id = max_thing_id113 self._is_training = is_training114 self._preprocessing_fn = functools.partial(115 preprocessing.preprocess_image_and_label,116 crop_height=crop_size[0],117 crop_width=crop_size[1],118 min_resize_value=min_resize_value,119 max_resize_value=max_resize_value,120 resize_factor=resize_factor,121 min_scale_factor=min_scale_factor,122 max_scale_factor=max_scale_factor,123 scale_factor_step_size=scale_factor_step_size,124 autoaugment_policy_name=autoaugment_policy_name,125 ignore_label=self._ignore_label *126 self._dataset_info['panoptic_label_divisor'],127 is_training=self._is_training)128 129 if focus_small_instances is not None:130 self._instance_area_threshold = focus_small_instances['threshold']131 self._small_instance_weight = focus_small_instances['weight']132 133 self._gaussian, self._gaussian_size = _compute_gaussian_from_std(134 self._sigma)135 self._gaussian = tf.cast(tf.reshape(self._gaussian, [-1]), tf.float32)136 137 def __call__(self, sample_dict):138 """Gets a sample.139 140 Args:141 sample_dict: A dictionary with the following keys and values:142 - `image`: A tensor of shape [image_height, image_width, 3].143 - `image_name`: String, image name.144 - `label`: A tensor of shape [label_height, label_width, 1] or None.145 - `height`: An integer specifying the height of the image.146 - `width`: An integer specifying the width of the image.147 - `sequence`: An optional string specifying the sequence name.148 - `prev_image`: An optional tensor of the same shape as `image`.149 - `prev_label`: An optional tensor of the same shape as `label`.150 - `next_image`: An optional next-frame tensor of the shape of `image`.151 - `next_label`: An optional next-frame tensor of the shape of `label`.152 153 Returns:154 sample: A dictionary storing required data for panoptic segmentation.155 """156 return self.call(**sample_dict)157 158 def call(self,159 image,160 image_name,161 label,162 height,163 width,164 sequence='',165 prev_image=None,166 prev_label=None,167 next_image=None,168 next_label=None):169 """Gets a sample.170 171 Args:172 image: A tensor of shape [image_height, image_width, 3].173 image_name: String, image name.174 label: A tensor of shape [label_height, label_width, 1] or None.175 height: An integer specifying the height of the image.176 width: An integer specifying the width of the image.177 sequence: An optional string specifying the sequence name.178 prev_image: An optional tensor of shape [image_height, image_width, 3].179 prev_label: An optional tensor of shape [label_height, label_width, 1].180 next_image: An optional tensor of shape [image_height, image_width, 3].181 next_label: An optional tensor of shape [label_height, label_width, 1].182 183 Returns:184 sample: A dictionary storing required data for panoptic segmentation.185 186 Raises:187 ValueError: An error occurs when the label shape is invalid.188 NotImplementedError: An error occurs when thing_id_mask_annotations comes189 together with prev_image or prev_label, not currently implemented.190 """191 if label is not None:192 label.get_shape().assert_is_compatible_with(193 tf.TensorShape([None, None, 1]))194 original_label = tf.cast(label, dtype=tf.int32, name='original_label')195 if next_label is not None:196 original_next_label = tf.cast(197 next_label, dtype=tf.int32, name='original_next_label')198 # Reusing the preprocessing function for both next and prev samples.199 if next_image is not None:200 resized_image, image, label, next_image, next_label = (201 self._preprocessing_fn(202 image, label, prev_image=next_image, prev_label=next_label))203 else:204 resized_image, image, label, prev_image, prev_label = (205 self._preprocessing_fn(206 image, label, prev_image=prev_image, prev_label=prev_label))207 sample = {208 common.IMAGE: image209 }210 if prev_image is not None:211 sample[common.IMAGE] = tf.concat([image, prev_image], axis=2)212 if next_image is not None:213 sample[common.NEXT_IMAGE] = next_image214 sample[common.IMAGE] = tf.concat([image, next_image], axis=2)215 if label is not None:216 # Panoptic label for crowd regions will be ignore_label.217 semantic_label, panoptic_label, thing_mask, crowd_region = (218 dataset_utils.get_semantic_and_panoptic_label(219 self._dataset_info, label, self._ignore_label))220 sample[common.GT_SEMANTIC_KEY] = tf.squeeze(semantic_label, axis=2)221 semantic_weights = tf.ones_like(semantic_label, dtype=tf.float32)222 sample[common.SEMANTIC_LOSS_WEIGHT_KEY] = tf.squeeze(223 semantic_weights, axis=2)224 sample[common.GT_IS_CROWD] = tf.squeeze(crowd_region, axis=2)225 226 if not self._only_semantic_annotations:227 # The sample will have the original label including crowd regions.228 sample[common.GT_PANOPTIC_KEY] = tf.squeeze(label, axis=2)229 # Compute center loss for all non-crowd and non-ignore pixels.230 non_crowd_and_non_ignore_regions = tf.logical_and(231 tf.logical_not(crowd_region),232 tf.not_equal(semantic_label, self._ignore_label))233 sample[common.CENTER_LOSS_WEIGHT_KEY] = tf.squeeze(tf.cast(234 non_crowd_and_non_ignore_regions, tf.float32), axis=2)235 # Compute regression loss only for thing pixels that are not crowd.236 non_crowd_things = tf.logical_and(237 tf.logical_not(crowd_region), thing_mask)238 sample[common.REGRESSION_LOSS_WEIGHT_KEY] = tf.squeeze(tf.cast(239 non_crowd_things, tf.float32), axis=2)240 241 prev_panoptic_label = None242 next_panoptic_label = None243 if prev_label is not None:244 _, prev_panoptic_label, _, _ = (245 dataset_utils.get_semantic_and_panoptic_label(246 self._dataset_info, prev_label, self._ignore_label))247 if next_label is not None:248 _, next_panoptic_label, _, _ = (249 dataset_utils.get_semantic_and_panoptic_label(250 self._dataset_info, next_label, self._ignore_label))251 (sample[common.GT_INSTANCE_CENTER_KEY],252 sample[common.GT_INSTANCE_REGRESSION_KEY],253 sample[common.SEMANTIC_LOSS_WEIGHT_KEY],254 prev_center_map,255 frame_center_offsets,256 next_offset) = self._generate_gt_center_and_offset(257 panoptic_label, semantic_weights, prev_panoptic_label,258 next_panoptic_label)259 260 sample[common.GT_INSTANCE_REGRESSION_KEY] = tf.cast(261 sample[common.GT_INSTANCE_REGRESSION_KEY], tf.float32)262 263 if next_label is not None:264 sample[common.GT_NEXT_INSTANCE_REGRESSION_KEY] = tf.cast(265 next_offset, tf.float32)266 sample[common.NEXT_REGRESSION_LOSS_WEIGHT_KEY] = tf.cast(267 tf.greater(tf.reduce_sum(tf.abs(next_offset), axis=2), 0),268 tf.float32)269 270 # Only squeeze center map and semantic loss weights, as regression map271 # has two channels (x and y offsets).272 sample[common.GT_INSTANCE_CENTER_KEY] = tf.squeeze(273 sample[common.GT_INSTANCE_CENTER_KEY], axis=2)274 sample[common.SEMANTIC_LOSS_WEIGHT_KEY] = tf.squeeze(275 sample[common.SEMANTIC_LOSS_WEIGHT_KEY], axis=2)276 277 if prev_label is not None:278 sample[common.GT_FRAME_OFFSET_KEY] = frame_center_offsets279 sample[common.GT_FRAME_OFFSET_KEY] = tf.cast(280 sample[common.GT_FRAME_OFFSET_KEY], tf.float32)281 frame_offsets_present = tf.logical_or(282 tf.not_equal(frame_center_offsets[..., 0], 0),283 tf.not_equal(frame_center_offsets[..., 1], 0))284 sample[common.FRAME_REGRESSION_LOSS_WEIGHT_KEY] = tf.cast(285 frame_offsets_present, tf.float32)286 if self._is_training:287 sample[common.IMAGE] = tf.concat(288 [sample[common.IMAGE], prev_center_map], axis=2)289 290 if self._thing_id_mask_annotations:291 if any([prev_image is not None,292 prev_label is not None,293 next_image is not None,294 next_label is not None]):295 raise NotImplementedError(296 'Current implementation of Max-DeepLab does not support '297 + 'prev_image, prev_label, next_image, or next_label.')298 thing_id_mask, thing_id_class = (299 self._generate_thing_id_mask_and_class(300 panoptic_label, non_crowd_things))301 sample[common.GT_THING_ID_MASK_KEY] = tf.squeeze(302 thing_id_mask, axis=2)303 sample[common.GT_THING_ID_CLASS_KEY] = thing_id_class304 305 if not self._is_training:306 # Resized image is only used during visualization.307 sample[common.RESIZED_IMAGE] = resized_image308 sample[common.IMAGE_NAME] = image_name309 sample[common.GT_SIZE_RAW] = tf.stack([height, width], axis=0)310 if self._dataset_info['is_video_dataset']:311 sample[common.SEQUENCE_ID] = sequence312 # Keep original labels for evaluation.313 if label is not None:314 orig_semantic_label, _, _, orig_crowd_region = (315 dataset_utils.get_semantic_and_panoptic_label(316 self._dataset_info, original_label, self._ignore_label))317 sample[common.GT_SEMANTIC_RAW] = tf.squeeze(orig_semantic_label, axis=2)318 if not self._only_semantic_annotations:319 sample[common.GT_PANOPTIC_RAW] = tf.squeeze(original_label, axis=2)320 sample[common.GT_IS_CROWD_RAW] = tf.squeeze(orig_crowd_region)321 if next_label is not None:322 sample[common.GT_NEXT_PANOPTIC_RAW] = tf.squeeze(323 original_next_label, axis=2)324 return sample325 326 def _generate_thing_id_mask_and_class(self,327 panoptic_label,328 non_crowd_things):329 """Generates the ground-truth thing-ID masks and their class labels.330 331 It computes thing-ID mask and class with unique ID for each thing instance.332 `thing_id` indicates the number of unique thing-ID to each instance in an333 image, starting the counting from 0. Each pixel in thing_id_mask is labeled334 with the corresponding thing-ID.335 336 Args:337 panoptic_label: A tf.Tensor of shape [height, width, 1].338 non_crowd_things: A tf.Tensor of shape [height, width, 1], indicating339 non-crowd and thing-class regions.340 341 Returns:342 thing_id_mask: A tf.Tensor of shape [height, width, 1]. It assigns each343 non-crowd thing instance a unique mask-ID label, starting from 0.344 Unassigned pixels are set to -1.345 thing_id_class: A tf.Tensor of shape [max_thing_id]. It contains semantic346 ID of each instance assigned to thing_id_mask. The remaining347 (max_thing_id - num_things) elements are set to -1.348 349 Raises:350 ValueError: An error occurs when the thing-ID mask contains stuff or crowd351 region.352 ValueError: An error occurs when thing_count is greater or equal to353 self._max_thing_id.354 355 """356 unique_ids, _ = tf.unique(tf.reshape(panoptic_label, [-1]))357 thing_id_mask = -tf.ones_like(panoptic_label)358 thing_id_class = -tf.ones(self._max_thing_id)359 thing_count = 0360 for panoptic_id in unique_ids:361 semantic_id = panoptic_id // self._dataset_info['panoptic_label_divisor']362 # Filter out IDs that are not thing instances (i.e., IDs for ignore_label,363 # stuff classes or crowd). Stuff classes and crowd regions both have IDs364 # of the form panoptic_id = semantic_id * label_divisor (i.e., instance id365 # = 0)366 if (semantic_id == self._dataset_info['ignore_label'] or367 panoptic_id % self._dataset_info['panoptic_label_divisor'] == 0):368 continue369 370 assert_stuff_crowd = tf.debugging.Assert(371 tf.reduce_all(non_crowd_things[panoptic_label == panoptic_id]),372 ['thing-ID mask here must not contain stuff or crowd region.'])373 with tf.control_dependencies([assert_stuff_crowd]):374 panoptic_id = tf.identity(panoptic_id)375 376 thing_id_mask = tf.where(panoptic_label == panoptic_id,377 thing_count, thing_id_mask)378 379 assert_thing_count = tf.debugging.Assert(380 thing_count < self._max_thing_id,381 ['thing_count must be smaller than self._max_thing_id.'])382 with tf.control_dependencies([assert_thing_count]):383 thing_count = tf.identity(thing_count)384 385 thing_id_class = tf.tensor_scatter_nd_update(386 thing_id_class, [[thing_count]], [semantic_id])387 thing_count += 1388 return thing_id_mask, thing_id_class389 390 def _generate_prev_centers_with_noise(self,391 panoptic_label,392 offset_noise_factor=0.05,393 false_positive_rate=0.2,394 false_positive_noise_factor=0.05):395 """Generates noisy center predictions for the previous frame.396 397 Args:398 panoptic_label: A tf.Tensor of shape [height, width, 1].399 offset_noise_factor: An optional float defining the maximum fraction of400 the object size that is used to displace the previous center.401 false_positive_rate: An optional float indicating at which probability402 false positives should be added.403 false_positive_noise_factor: An optional float defining the maximum404 fraction of the object size that is used to displace the false positive405 center.406 407 Returns:408 A tuple of (center, ids_to_center) with both being tf.Tensor of shape409 [height, width, 1] and shape [N, 2] where N is the number of unique IDs.410 """411 height = tf.shape(panoptic_label)[0]412 width = tf.shape(panoptic_label)[1]413 414 # Pad center to make boundary handling easier.415 center_pad_begin = int(round(3 * self._sigma + 1))416 center_pad_end = int(round(3 * self._sigma + 2))417 center_pad = center_pad_begin + center_pad_end418 419 center = tf.zeros((height + center_pad, width + center_pad))420 unique_ids, _ = tf.unique(tf.reshape(panoptic_label, [-1]))421 ids_to_center_x = tf.zeros_like(unique_ids, dtype=tf.int32)422 ids_to_center_y = tf.zeros_like(unique_ids, dtype=tf.int32)423 424 for panoptic_id in unique_ids:425 semantic_id = panoptic_id // self._dataset_info['panoptic_label_divisor']426 # Filter out IDs that should be ignored, are stuff classes or crowd.427 # Stuff classes and crowd regions both have IDs of the form panoptic_id =428 # semantic_id * label_divisor429 if (semantic_id == self._dataset_info['ignore_label'] or430 panoptic_id % self._dataset_info['panoptic_label_divisor'] == 0):431 continue432 433 # Convert [[y0, x0, 0], ...] to [[y0, ...], [x0, ...], [0, ...]].434 mask_index = tf.cast(435 tf.transpose(tf.where(panoptic_label == panoptic_id)), tf.float32)436 centers = tf.reduce_mean(mask_index, axis=1)437 bbox_size = (438 tf.reduce_max(mask_index, axis=1) - tf.reduce_min(mask_index, axis=1))439 440 # Add noise.441 center_y = (442 centers[0] + tf.random.normal([], dtype=tf.float32) *443 offset_noise_factor * bbox_size[0])444 center_x = (445 centers[1] + tf.random.normal([], dtype=tf.float32) *446 offset_noise_factor * bbox_size[1])447 448 center_x = tf.minimum(449 tf.maximum(tf.cast(tf.round(center_x), tf.int32), 0), width - 1)450 center_y = tf.minimum(451 tf.maximum(tf.cast(tf.round(center_y), tf.int32), 0), height - 1)452 453 id_index = tf.where(tf.equal(panoptic_id, unique_ids))454 ids_to_center_x = tf.tensor_scatter_nd_update(455 ids_to_center_x, id_index, tf.expand_dims(center_x, axis=0))456 ids_to_center_y = tf.tensor_scatter_nd_update(457 ids_to_center_y, id_index, tf.expand_dims(center_y, axis=0))458 459 def add_center_gaussian(center_x_coord, center_y_coord, center):460 # Due to the padding with center_pad_begin in center, the computed461 # center becomes the upper left corner in the center tensor.462 upper_left = center_x_coord, center_y_coord463 bottom_right = (upper_left[0] + self._gaussian_size,464 upper_left[1] + self._gaussian_size)465 466 indices_x, indices_y = tf.meshgrid(467 tf.range(upper_left[0], bottom_right[0]),468 tf.range(upper_left[1], bottom_right[1]))469 indices = tf.transpose(470 tf.stack([tf.reshape(indices_y, [-1]),471 tf.reshape(indices_x, [-1])]))472 473 return tf.tensor_scatter_nd_max(474 center, indices, self._gaussian, name='center_scatter')475 476 center = add_center_gaussian(center_x, center_y, center)477 # Generate false positives.478 center_y = (479 tf.cast(center_y, dtype=tf.float32) +480 tf.random.normal([], dtype=tf.float32) * false_positive_noise_factor *481 bbox_size[0])482 center_x = (483 tf.cast(center_x, dtype=tf.float32) +484 tf.random.normal([], dtype=tf.float32) * false_positive_noise_factor *485 bbox_size[1])486 487 center_x = tf.minimum(488 tf.maximum(tf.cast(tf.round(center_x), tf.int32), 0), width - 1)489 center_y = tf.minimum(490 tf.maximum(tf.cast(tf.round(center_y), tf.int32), 0), height - 1)491 # Draw a sample to decide whether to add a false positive or not.492 center = center + tf.cast(493 tf.random.uniform([], dtype=tf.float32) < false_positive_rate,494 tf.float32) * (495 add_center_gaussian(center_x, center_y, center) - center)496 497 center = center[center_pad_begin:(center_pad_begin + height),498 center_pad_begin:(center_pad_begin + width)]499 center = tf.expand_dims(center, -1)500 return center, unique_ids, ids_to_center_x, ids_to_center_y501 502 def _generate_gt_center_and_offset(self,503 panoptic_label,504 semantic_weights,505 prev_panoptic_label=None,506 next_panoptic_label=None):507 """Generates the ground-truth center and offset from the panoptic labels.508 509 Additionally, the per-pixel weights for the semantic branch are increased510 for small instances. In case, prev_panoptic_label is passed, it also511 computes the previous center heatmap with random noise and the offsets512 between center maps.513 514 Args:515 panoptic_label: A tf.Tensor of shape [height, width, 1].516 semantic_weights: A tf.Tensor of shape [height, width, 1].517 prev_panoptic_label: An optional tf.Tensor of shape [height, width, 1].518 next_panoptic_label: An optional tf.Tensor of shape [height, width, 1].519 520 Returns:521 A tuple (center, offsets, weights, prev_center, frame_offset*,522 next_offset) with each being a tf.Tensor of shape [height, width, 1 (2*)].523 If prev_panoptic_label is None, prev_center and frame_offset are None.524 If next_panoptic_label is None, next_offset is None.525 """526 height = tf.shape(panoptic_label)[0]527 width = tf.shape(panoptic_label)[1]528 529 # Pad center to make boundary handling easier.530 center_pad_begin = int(round(3 * self._sigma + 1))531 center_pad_end = int(round(3 * self._sigma + 2))532 center_pad = center_pad_begin + center_pad_end533 534 center = tf.zeros((height + center_pad, width + center_pad))535 offset_x = tf.zeros((height, width, 1), dtype=tf.int32)536 offset_y = tf.zeros((height, width, 1), dtype=tf.int32)537 unique_ids, _ = tf.unique(tf.reshape(panoptic_label, [-1]))538 539 prev_center = None540 frame_offsets = None541 # Due to loop handling in tensorflow, these variables had to be defined for542 # all cases.543 frame_offset_x = tf.zeros((height, width, 1), dtype=tf.int32)544 frame_offset_y = tf.zeros((height, width, 1), dtype=tf.int32)545 546 # Next-frame instance offsets.547 next_offset = None548 next_offset_y = tf.zeros((height, width, 1), dtype=tf.int32)549 next_offset_x = tf.zeros((height, width, 1), dtype=tf.int32)550 551 if prev_panoptic_label is not None:552 (prev_center, prev_unique_ids, prev_centers_x, prev_centers_y553 ) = self._generate_prev_centers_with_noise(prev_panoptic_label)554 555 for panoptic_id in unique_ids:556 semantic_id = panoptic_id // self._dataset_info['panoptic_label_divisor']557 # Filter out IDs that should be ignored, are stuff classes or crowd.558 # Stuff classes and crowd regions both have IDs of the form panopti_id =559 # semantic_id * label_divisor560 if (semantic_id == self._dataset_info['ignore_label'] or561 panoptic_id % self._dataset_info['panoptic_label_divisor'] == 0):562 continue563 564 # Convert [[y0, x0, 0], ...] to [[y0, ...], [x0, ...], [0, ...]].565 mask_index = tf.transpose(tf.where(panoptic_label == panoptic_id))566 mask_y_index = mask_index[0]567 mask_x_index = mask_index[1]568 569 next_mask_index = None570 next_mask_y_index = None571 next_mask_x_index = None572 if next_panoptic_label is not None:573 next_mask_index = tf.transpose(574 tf.where(next_panoptic_label == panoptic_id))575 next_mask_y_index = next_mask_index[0]576 next_mask_x_index = next_mask_index[1]577 578 instance_area = tf.shape(mask_x_index)579 if instance_area < self._instance_area_threshold:580 semantic_weights = tf.where(panoptic_label == panoptic_id,581 self._small_instance_weight,582 semantic_weights)583 584 centers = tf.reduce_mean(tf.cast(mask_index, tf.float32), axis=1)585 586 center_x = tf.cast(tf.round(centers[1]), tf.int32)587 center_y = tf.cast(tf.round(centers[0]), tf.int32)588 589 # Due to the padding with center_pad_begin in center, the computed center590 # becomes the upper left corner in the center tensor.591 upper_left = center_x, center_y592 bottom_right = (upper_left[0] + self._gaussian_size,593 upper_left[1] + self._gaussian_size)594 595 indices_x, indices_y = tf.meshgrid(596 tf.range(upper_left[0], bottom_right[0]),597 tf.range(upper_left[1], bottom_right[1]))598 indices = tf.transpose(599 tf.stack([tf.reshape(indices_y, [-1]),600 tf.reshape(indices_x, [-1])]))601 602 center = tf.tensor_scatter_nd_max(603 center, indices, self._gaussian, name='center_scatter')604 offset_y = tf.tensor_scatter_nd_update(605 offset_y,606 tf.transpose(mask_index),607 center_y - tf.cast(mask_y_index, tf.int32),608 name='offset_y_scatter')609 offset_x = tf.tensor_scatter_nd_update(610 offset_x,611 tf.transpose(mask_index),612 center_x - tf.cast(mask_x_index, tf.int32),613 name='offset_x_scatter')614 if prev_panoptic_label is not None:615 mask = tf.equal(prev_unique_ids, panoptic_id)616 if tf.math.count_nonzero(mask) > 0:617 prev_center_x = prev_centers_x[mask]618 prev_center_y = prev_centers_y[mask]619 620 frame_offset_y = tf.tensor_scatter_nd_update(621 frame_offset_y,622 tf.transpose(mask_index),623 prev_center_y - tf.cast(mask_y_index, tf.int32),624 name='frame_offset_y_scatter')625 frame_offset_x = tf.tensor_scatter_nd_update(626 frame_offset_x,627 tf.transpose(mask_index),628 prev_center_x - tf.cast(mask_x_index, tf.int32),629 name='frame_offset_x_scatter')630 if next_panoptic_label is not None:631 next_offset_y = tf.tensor_scatter_nd_update(632 next_offset_y,633 tf.transpose(next_mask_index),634 center_y - tf.cast(next_mask_y_index, tf.int32),635 name='next_offset_y_scatter')636 next_offset_x = tf.tensor_scatter_nd_update(637 next_offset_x,638 tf.transpose(next_mask_index),639 center_x - tf.cast(next_mask_x_index, tf.int32),640 name='next_offset_x_scatter')641 642 offset = tf.concat([offset_y, offset_x], axis=2)643 center = center[center_pad_begin:(center_pad_begin + height),644 center_pad_begin:(center_pad_begin + width)]645 center = tf.expand_dims(center, -1)646 if prev_panoptic_label is not None:647 frame_offsets = tf.concat([frame_offset_y, frame_offset_x], axis=2)648 if next_panoptic_label is not None:649 next_offset = tf.concat([next_offset_y, next_offset_x], axis=2)650 return (center, offset, semantic_weights, prev_center, frame_offsets,651 next_offset)652