def process( self, class_predictions_with_background: tf.Tensor, num_object_detections: tf.Tensor, detection_object_boxes: tf.Tensor, groundtruth_object_boxes: tf.Tensor, groundtruth_object_classes: tf.Tensor, groundtruth_object_weights: Optional[tf.Tensor] = None ) -> Dict[str, tf.Tensor]: # pylint: disable=arguments-differ,too-many-locals # parent save method has more generic signature detection_object_boxes = tf.stop_gradient(detection_object_boxes) num_classes_with_bg = ( class_predictions_with_background.shape.as_list()[-1]) groundtruth_classes_with_background = tf.one_hot( indices=groundtruth_object_classes, depth=num_classes_with_bg, dtype=tf.float32) (class_targets, class_targets_weights, _, _) = _batch_assign_targets_dynamic( self._detector_target_assigner, anchors=detection_object_boxes, groundtruth_object_boxes=groundtruth_object_boxes, groundtruth_classes_with_background= groundtruth_classes_with_background, num_classes=num_classes_with_bg - 1, groundtruth_object_weights=groundtruth_object_weights) max_num_proposals = tf.shape(detection_object_boxes)[1] normalizer = _get_normalizer( tf.shape(detection_object_boxes)[0], num_object_detections, max_num_proposals) proposals_batch_indicator = _get_proposals_batch_indicator( num_object_detections, max_num_proposals) second_stage_cls_losses = self._loss_fn( class_predictions_with_background, class_targets, weights=class_targets_weights, losses_mask=None) normalizer = broadcast_with_expand_to(normalizer, second_stage_cls_losses) proposals_batch_indicator = broadcast_with_expand_to( proposals_batch_indicator, second_stage_cls_losses) second_stage_cls_losses = second_stage_cls_losses / normalizer second_stage_cls_loss = tf.reduce_sum(second_stage_cls_losses * proposals_batch_indicator, name="loss_classification_match") return {"loss_classification_match": second_stage_cls_loss}
def _cutout_keypoints(object_keypoints, cutout_box): (cutout_box_ymin, cutout_box_xmin, cutout_box_ymax, cutout_box_xmax) = tf.split(cutout_box, 4, -1) object_keypoints_y = object_keypoints[..., 0] object_keypoints_x = object_keypoints[..., 1] valid_keypoints_mask_first = tf.reduce_all(tf.logical_and( tf.greater(object_keypoints, 0), tf.less_equal(object_keypoints, 1.0)), -1, keepdims=True) valid_keypoints_mask_cutout = tf.reduce_any(tf.stack([ tf.greater_equal(object_keypoints_y, cutout_box_ymax), tf.less_equal(object_keypoints_y, cutout_box_ymin), tf.greater_equal(object_keypoints_x, cutout_box_xmax), tf.less_equal(object_keypoints_x, cutout_box_xmin), ], -1), -1, keepdims=True) valid_keypoints_mask = tf.logical_and(valid_keypoints_mask_first, valid_keypoints_mask_cutout) keypoints_cutout = tf.where( broadcast_with_expand_to(valid_keypoints_mask, object_keypoints), object_keypoints, tf.zeros_like(object_keypoints)) return keypoints_cutout
def cutout(self, image: tf.Tensor, replace_value: Union[float, tf.Tensor] = 0.0) -> tf.Tensor: """ Apply cutout to the given image Parameters ---------- image input image replace_value the value to which the rectangle is set Returns ------- img_with_cutoff the input images after applying cutout """ image_size = tf.shape(image)[:-1] image_size_float = tf.cast(image_size, tf.float32) cut_lengths_absolute = tf.cast( tf.floor(self.random_variables["cut_lengths"] * image_size_float), tf.int32) cut_offset_absolute = tf.cast( tf.floor(self.random_variables["cut_offset"] * image_size_float), tf.int32) mask = self._create_cutout_mask(cut_offset_absolute, cut_lengths_absolute, image_size) mask = broadcast_with_expand_to(mask, image) mask_with_replaced_value = tf.cast(mask, image.dtype) * replace_value image_with_cutoff = tf.where(mask, mask_with_replaced_value, image) return image_with_cutoff
def _mask_single_batch_input_to_classes( item: tf.Tensor, classes_mask: tf.Tensor, rearranged_indices: tf.Tensor) -> tf.Tensor: mask_broadcasted = broadcast_with_expand_to(classes_mask, item) item_masked = tf.where(mask_broadcasted, item, tf.zeros_like(item)) item_masked_rearranged = tf.batch_gather(item_masked, rearranged_indices) return item_masked_rearranged
def _flip_keypoints_up_down(keypoints): keypoints_y, keypoints_x = tf.split(keypoints, 2, -1) keypoints_y_flipped = 1 - keypoints_y keypoints_flipped = tf.concat([keypoints_y_flipped, keypoints_x], -1) valid_keypoints = tf.reduce_any(tf.greater(keypoints, 0), -1) keypoints_flipped = tf.where( broadcast_with_expand_to(valid_keypoints, keypoints), keypoints_flipped, tf.zeros_like(keypoints_flipped)) return keypoints_flipped
def _flip_boxes_up_down(boxes): ymin, xmin, ymax, xmax = tf.split(boxes, 4, -1) height = ymax - ymin ymin = 1.0 - height - ymin ymax = ymin + height boxes_flipped = tf.concat([ymin, xmin, ymax, xmax], -1) valid_boxes = tf.greater(height, 0) boxes_flipped = tf.where(broadcast_with_expand_to(valid_boxes, boxes), boxes_flipped, tf.zeros_like(boxes_flipped)) return boxes_flipped
def _flip_boxes_left_right(boxes): ymin, xmin, ymax, xmax = tf.split(boxes, 4, -1) width = xmax - xmin xmin = 1.0 - width - xmin xmax = xmin + width boxes_flipped = tf.concat([ymin, xmin, ymax, xmax], -1) valid_boxes = tf.greater(width, 0) boxes_flipped = tf.where(broadcast_with_expand_to(valid_boxes, boxes), boxes_flipped, tf.zeros_like(boxes_flipped)) return boxes_flipped
def _rotate_keypoints(self, object_keypoints: tf.Tensor) -> tf.Tensor: rotation_angle = self.random_variables["rotation_angle"] keypoints_rotated = _rotate_points(object_keypoints, rotation_angle) valid_keypoints_mask_first = tf.reduce_all(tf.logical_and( tf.greater(object_keypoints, 0), tf.less_equal(object_keypoints, 1.0)), -1, keepdims=True) valid_keypoints_mask_after_rotation = tf.reduce_all(tf.logical_and( tf.greater(keypoints_rotated, 0), tf.less_equal(keypoints_rotated, 1.0)), -1, keepdims=True) valid_keypoints_mask = tf.logical_and( valid_keypoints_mask_first, valid_keypoints_mask_after_rotation) keypoints_rotated = tf.where( broadcast_with_expand_to(valid_keypoints_mask, object_keypoints), keypoints_rotated, tf.zeros_like(object_keypoints)) return keypoints_rotated
def _crop_keypoints(self, object_keypoints: tf.Tensor) -> tf.Tensor: crop_offset = self.random_variables["crop_offset"] crop_scale = self.random_variables["crop_scale"] keypoints_with_offset = object_keypoints - crop_offset keypoints_with_offset_scaled = keypoints_with_offset / crop_scale valid_keypoints_mask_first = tf.reduce_all(tf.logical_and( tf.greater(object_keypoints, 0), tf.less_equal(object_keypoints, 1.0)), -1, keepdims=True) valid_keypoints_mask_after_crop = tf.reduce_all(tf.logical_and( tf.greater(keypoints_with_offset_scaled, 0), tf.less_equal(keypoints_with_offset_scaled, 1.0)), -1, keepdims=True) valid_keypoints_mask = tf.logical_and(valid_keypoints_mask_first, valid_keypoints_mask_after_crop) keypoints_cropped = tf.where( broadcast_with_expand_to(valid_keypoints_mask, object_keypoints), keypoints_with_offset_scaled, tf.zeros_like(object_keypoints)) return keypoints_cropped
def process( self, *, detection_object_boxes: tf.Tensor, num_object_detections: tf.Tensor, detection_object_keypoints_heatmaps: tf.Tensor, groundtruth_object_boxes: tf.Tensor, groundtruth_object_keypoints: tf.Tensor, groundtruth_object_weights: Optional[tf.Tensor] = None ) -> Dict[str, tf.Tensor]: # pylint: disable=arguments-differ,too-many-locals # parent save method has more generic signature detection_object_boxes = tf.stop_gradient(detection_object_boxes) heatmaps_shape = tf.shape(detection_object_keypoints_heatmaps)[2:4] groundtruth_object_keypoints = tf.maximum(groundtruth_object_keypoints, 0) (target_keypoints_encoded, keypoints_masks, target_weights) = self._match_groundtruth_keypoints( detection_object_boxes, groundtruth_object_boxes, groundtruth_object_keypoints, groundtruth_object_weights) target_keypoints_heatmaps = self._get_target_keypoints_heatmaps( target_keypoints_encoded, keypoints_masks, heatmaps_shape) max_num_detections = tf.shape(detection_object_boxes)[1] batch_indicator = _get_proposals_batch_indicator( num_object_detections, max_num_detections) normalizer = tf.maximum( tf.reduce_sum(batch_indicator * target_weights * tf.reduce_prod(tf.cast(heatmaps_shape, tf.float32))), 1.0) detection_object_keypoints_heatmaps_r = tf.reshape( detection_object_keypoints_heatmaps, [ tf.shape(detection_object_keypoints_heatmaps)[0], -1, tf.shape(detection_object_keypoints_heatmaps)[-1] ]) target_keypoints_heatmaps_r = tf.reshape(target_keypoints_heatmaps, [ tf.shape(target_keypoints_heatmaps)[0], -1, tf.shape(target_keypoints_heatmaps)[-1] ]) weights = tf.reshape( broadcast_with_expand_to(target_weights, target_keypoints_heatmaps[..., 0]), [tf.shape(target_weights)[0], -1]) batch_indicator_r = tf.reshape( broadcast_with_expand_to(batch_indicator, target_keypoints_heatmaps[..., 0]), [tf.shape(target_weights)[0], -1]) if "Classification" in self.heatmaps_loss_name: weights = weights[..., tf.newaxis] batch_indicator_r = batch_indicator_r[..., tf.newaxis] heatmaps_loss_samples = self._loss_fn( detection_object_keypoints_heatmaps_r, target_keypoints_heatmaps_r, weights=weights, ) / normalizer loss = tf.reduce_sum(heatmaps_loss_samples * batch_indicator_r, name="loss_keypoints_heatmaps") return {"loss_keypoints_heatmaps": loss}