You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
171 lines
6.2 KiB
171 lines
6.2 KiB
# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
|
|
#
|
|
# 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.
|
|
# ==============================================================================
|
|
|
|
"""Keypoint box coder.
|
|
|
|
The keypoint box coder follows the coding schema described below (this is
|
|
similar to the FasterRcnnBoxCoder, except that it encodes keypoints in addition
|
|
to box coordinates):
|
|
ty = (y - ya) / ha
|
|
tx = (x - xa) / wa
|
|
th = log(h / ha)
|
|
tw = log(w / wa)
|
|
tky0 = (ky0 - ya) / ha
|
|
tkx0 = (kx0 - xa) / wa
|
|
tky1 = (ky1 - ya) / ha
|
|
tkx1 = (kx1 - xa) / wa
|
|
...
|
|
where x, y, w, h denote the box's center coordinates, width and height
|
|
respectively. Similarly, xa, ya, wa, ha denote the anchor's center
|
|
coordinates, width and height. tx, ty, tw and th denote the anchor-encoded
|
|
center, width and height respectively. ky0, kx0, ky1, kx1, ... denote the
|
|
keypoints' coordinates, and tky0, tkx0, tky1, tkx1, ... denote the
|
|
anchor-encoded keypoint coordinates.
|
|
"""
|
|
|
|
import tensorflow as tf
|
|
|
|
from object_detection.core import box_coder
|
|
from object_detection.core import box_list
|
|
from object_detection.core import standard_fields as fields
|
|
|
|
EPSILON = 1e-8
|
|
|
|
|
|
class KeypointBoxCoder(box_coder.BoxCoder):
|
|
"""Keypoint box coder."""
|
|
|
|
def __init__(self, num_keypoints, scale_factors=None):
|
|
"""Constructor for KeypointBoxCoder.
|
|
|
|
Args:
|
|
num_keypoints: Number of keypoints to encode/decode.
|
|
scale_factors: List of 4 positive scalars to scale ty, tx, th and tw.
|
|
In addition to scaling ty and tx, the first 2 scalars are used to scale
|
|
the y and x coordinates of the keypoints as well. If set to None, does
|
|
not perform scaling.
|
|
"""
|
|
self._num_keypoints = num_keypoints
|
|
|
|
if scale_factors:
|
|
assert len(scale_factors) == 4
|
|
for scalar in scale_factors:
|
|
assert scalar > 0
|
|
self._scale_factors = scale_factors
|
|
self._keypoint_scale_factors = None
|
|
if scale_factors is not None:
|
|
self._keypoint_scale_factors = tf.expand_dims(tf.tile(
|
|
[tf.to_float(scale_factors[0]), tf.to_float(scale_factors[1])],
|
|
[num_keypoints]), 1)
|
|
|
|
@property
|
|
def code_size(self):
|
|
return 4 + self._num_keypoints * 2
|
|
|
|
def _encode(self, boxes, anchors):
|
|
"""Encode a box and keypoint collection with respect to anchor collection.
|
|
|
|
Args:
|
|
boxes: BoxList holding N boxes and keypoints to be encoded. Boxes are
|
|
tensors with the shape [N, 4], and keypoints are tensors with the shape
|
|
[N, num_keypoints, 2].
|
|
anchors: BoxList of anchors.
|
|
|
|
Returns:
|
|
a tensor representing N anchor-encoded boxes of the format
|
|
[ty, tx, th, tw, tky0, tkx0, tky1, tkx1, ...] where tky0 and tkx0
|
|
represent the y and x coordinates of the first keypoint, tky1 and tkx1
|
|
represent the y and x coordinates of the second keypoint, and so on.
|
|
"""
|
|
# Convert anchors to the center coordinate representation.
|
|
ycenter_a, xcenter_a, ha, wa = anchors.get_center_coordinates_and_sizes()
|
|
ycenter, xcenter, h, w = boxes.get_center_coordinates_and_sizes()
|
|
keypoints = boxes.get_field(fields.BoxListFields.keypoints)
|
|
keypoints = tf.transpose(tf.reshape(keypoints,
|
|
[-1, self._num_keypoints * 2]))
|
|
num_boxes = boxes.num_boxes()
|
|
|
|
# Avoid NaN in division and log below.
|
|
ha += EPSILON
|
|
wa += EPSILON
|
|
h += EPSILON
|
|
w += EPSILON
|
|
|
|
tx = (xcenter - xcenter_a) / wa
|
|
ty = (ycenter - ycenter_a) / ha
|
|
tw = tf.log(w / wa)
|
|
th = tf.log(h / ha)
|
|
|
|
tiled_anchor_centers = tf.tile(
|
|
tf.stack([ycenter_a, xcenter_a]), [self._num_keypoints, 1])
|
|
tiled_anchor_sizes = tf.tile(
|
|
tf.stack([ha, wa]), [self._num_keypoints, 1])
|
|
tkeypoints = (keypoints - tiled_anchor_centers) / tiled_anchor_sizes
|
|
|
|
# Scales location targets as used in paper for joint training.
|
|
if self._scale_factors:
|
|
ty *= self._scale_factors[0]
|
|
tx *= self._scale_factors[1]
|
|
th *= self._scale_factors[2]
|
|
tw *= self._scale_factors[3]
|
|
tkeypoints *= tf.tile(self._keypoint_scale_factors, [1, num_boxes])
|
|
|
|
tboxes = tf.stack([ty, tx, th, tw])
|
|
return tf.transpose(tf.concat([tboxes, tkeypoints], 0))
|
|
|
|
def _decode(self, rel_codes, anchors):
|
|
"""Decode relative codes to boxes and keypoints.
|
|
|
|
Args:
|
|
rel_codes: a tensor with shape [N, 4 + 2 * num_keypoints] representing N
|
|
anchor-encoded boxes and keypoints
|
|
anchors: BoxList of anchors.
|
|
|
|
Returns:
|
|
boxes: BoxList holding N bounding boxes and keypoints.
|
|
"""
|
|
ycenter_a, xcenter_a, ha, wa = anchors.get_center_coordinates_and_sizes()
|
|
|
|
num_codes = tf.shape(rel_codes)[0]
|
|
result = tf.unstack(tf.transpose(rel_codes))
|
|
ty, tx, th, tw = result[:4]
|
|
tkeypoints = result[4:]
|
|
if self._scale_factors:
|
|
ty /= self._scale_factors[0]
|
|
tx /= self._scale_factors[1]
|
|
th /= self._scale_factors[2]
|
|
tw /= self._scale_factors[3]
|
|
tkeypoints /= tf.tile(self._keypoint_scale_factors, [1, num_codes])
|
|
|
|
w = tf.exp(tw) * wa
|
|
h = tf.exp(th) * ha
|
|
ycenter = ty * ha + ycenter_a
|
|
xcenter = tx * wa + xcenter_a
|
|
ymin = ycenter - h / 2.
|
|
xmin = xcenter - w / 2.
|
|
ymax = ycenter + h / 2.
|
|
xmax = xcenter + w / 2.
|
|
decoded_boxes_keypoints = box_list.BoxList(
|
|
tf.transpose(tf.stack([ymin, xmin, ymax, xmax])))
|
|
|
|
tiled_anchor_centers = tf.tile(
|
|
tf.stack([ycenter_a, xcenter_a]), [self._num_keypoints, 1])
|
|
tiled_anchor_sizes = tf.tile(
|
|
tf.stack([ha, wa]), [self._num_keypoints, 1])
|
|
keypoints = tkeypoints * tiled_anchor_sizes + tiled_anchor_centers
|
|
keypoints = tf.reshape(tf.transpose(keypoints),
|
|
[-1, self._num_keypoints, 2])
|
|
decoded_boxes_keypoints.add_field(fields.BoxListFields.keypoints, keypoints)
|
|
return decoded_boxes_keypoints
|