keypoint_preprocess.py

# Copyright (c) 2021 PaddlePaddle 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.
"""
this code is based on https://github.com/open-mmlab/mmpose/mmpose/core/post_processing/post_transforms.py
"""
import cv2
import numpy as np


class EvalAffine(object):
 def __init__(self, size, stride=64):
 super(EvalAffine, self).__init__()
 self.size = size
 self.stride = stride

 def __call__(self, image, im_info):
 s = self.size
 h, w, _ = image.shape
 trans, size_resized = get_affine_mat_kernel(h, w, s, inv=False)
 image_resized = cv2.warpAffine(image, trans, size_resized)
 return image_resized, im_info


def get_affine_mat_kernel(h, w, s, inv=False):
 if w < h:
 w_ = s
 h_ = int(np.ceil((s / w * h) / 64.) * 64)
 scale_w = w
 scale_h = h_ / w_ * w

 else:
 h_ = s
 w_ = int(np.ceil((s / h * w) / 64.) * 64)
 scale_h = h
 scale_w = w_ / h_ * h

 center = np.array([np.round(w / 2.), np.round(h / 2.)])

 size_resized = (w_, h_)
 trans = get_affine_transform(
 center, np.array([scale_w, scale_h]), 0, size_resized, inv=inv)

 return trans, size_resized


def get_affine_transform(center,
 input_size,
 rot,
 output_size,
 shift=(0., 0.),
 inv=False):
 """Get the affine transform matrix, given the center/scale/rot/output_size.

 Args:
 center (np.ndarray[2, ]): Center of the bounding box (x, y).
 scale (np.ndarray[2, ]): Scale of the bounding box
 wrt [width, height].
 rot (float): Rotation angle (degree).
 output_size (np.ndarray[2, ]): Size of the destination heatmaps.
 shift (0-100%): Shift translation ratio wrt the width/height.
 Default (0., 0.).
 inv (bool): Option to inverse the affine transform direction.
 (inv=False: src->dst or inv=True: dst->src)

 Returns:
 np.ndarray: The transform matrix.
 """
 assert len(center) == 2
 assert len(output_size) == 2
 assert len(shift) == 2
 if not isinstance(input_size, (np.ndarray, list)):
 input_size = np.array([input_size, input_size], dtype=np.float32)
 scale_tmp = input_size

 shift = np.array(shift)
 src_w = scale_tmp[0]
 dst_w = output_size[0]
 dst_h = output_size[1]

 rot_rad = np.pi * rot / 180
 src_dir = rotate_point([0., src_w * -0.5], rot_rad)
 dst_dir = np.array([0., dst_w * -0.5])

 src = np.zeros((3, 2), dtype=np.float32)
 src[0, :] = center + scale_tmp * shift
 src[1, :] = center + src_dir + scale_tmp * shift
 src[2, :] = _get_3rd_point(src[0, :], src[1, :])

 dst = np.zeros((3, 2), dtype=np.float32)
 dst[0, :] = [dst_w * 0.5, dst_h * 0.5]
 dst[1, :] = np.array([dst_w * 0.5, dst_h * 0.5]) + dst_dir
 dst[2, :] = _get_3rd_point(dst[0, :], dst[1, :])

 if inv:
 trans = cv2.getAffineTransform(np.float32(dst), np.float32(src))
 else:
 trans = cv2.getAffineTransform(np.float32(src), np.float32(dst))

 return trans


def get_warp_matrix(theta, size_input, size_dst, size_target):
 """This code is based on 
 https://github.com/open-mmlab/mmpose/blob/master/mmpose/core/post_processing/post_transforms.py

 Calculate the transformation matrix under the constraint of unbiased.
 Paper ref: Huang et al. The Devil is in the Details: Delving into Unbiased
 Data Processing for Human Pose Estimation (CVPR 2020).

 Args:
 theta (float): Rotation angle in degrees.
 size_input (np.ndarray): Size of input image [w, h].
 size_dst (np.ndarray): Size of output image [w, h].
 size_target (np.ndarray): Size of ROI in input plane [w, h].

 Returns:
 matrix (np.ndarray): A matrix for transformation.
 """
 theta = np.deg2rad(theta)
 matrix = np.zeros((2, 3), dtype=np.float32)
 scale_x = size_dst[0] / size_target[0]
 scale_y = size_dst[1] / size_target[1]
 matrix[0, 0] = np.cos(theta) * scale_x
 matrix[0, 1] = -np.sin(theta) * scale_x
 matrix[0, 2] = scale_x * (
 -0.5 * size_input[0] * np.cos(theta) + 0.5 * size_input[1] *
 np.sin(theta) + 0.5 * size_target[0])
 matrix[1, 0] = np.sin(theta) * scale_y
 matrix[1, 1] = np.cos(theta) * scale_y
 matrix[1, 2] = scale_y * (
 -0.5 * size_input[0] * np.sin(theta) - 0.5 * size_input[1] *
 np.cos(theta) + 0.5 * size_target[1])
 return matrix


def rotate_point(pt, angle_rad):
 """Rotate a point by an angle.

 Args:
 pt (list[float]): 2 dimensional point to be rotated
 angle_rad (float): rotation angle by radian

 Returns:
 list[float]: Rotated point.
 """
 assert len(pt) == 2
 sn, cs = np.sin(angle_rad), np.cos(angle_rad)
 new_x = pt[0] * cs - pt[1] * sn
 new_y = pt[0] * sn + pt[1] * cs
 rotated_pt = [new_x, new_y]

 return rotated_pt


def _get_3rd_point(a, b):
 """To calculate the affine matrix, three pairs of points are required. This
 function is used to get the 3rd point, given 2D points a & b.

 The 3rd point is defined by rotating vector `a - b` by 90 degrees
 anticlockwise, using b as the rotation center.

 Args:
 a (np.ndarray): point(x,y)
 b (np.ndarray): point(x,y)

 Returns:
 np.ndarray: The 3rd point.
 """
 assert len(a) == 2
 assert len(b) == 2
 direction = a - b
 third_pt = b + np.array([-direction[1], direction[0]], dtype=np.float32)

 return third_pt


class TopDownEvalAffine(object):
 """apply affine transform to image and coords

 Args:
 trainsize (list): [w, h], the standard size used to train
 use_udp (bool): whether to use Unbiased Data Processing.
 records(dict): the dict contained the image and coords

 Returns:
 records (dict): contain the image and coords after tranformed

 """

 def __init__(self, trainsize, use_udp=False):
 self.trainsize = trainsize
 self.use_udp = use_udp

 def __call__(self, image, im_info):
 rot = 0
 imshape = im_info['im_shape'][::-1]
 center = im_info['center'] if 'center' in im_info else imshape / 2.
 scale = im_info['scale'] if 'scale' in im_info else imshape
 if self.use_udp:
 trans = get_warp_matrix(
 rot, center * 2.0,
 [self.trainsize[0] - 1.0, self.trainsize[1] - 1.0], scale)
 image = cv2.warpAffine(
 image,
 trans, (int(self.trainsize[0]), int(self.trainsize[1])),
 flags=cv2.INTER_LINEAR)
 else:
 trans = get_affine_transform(center, scale, rot, self.trainsize)
 image = cv2.warpAffine(
 image,
 trans, (int(self.trainsize[0]), int(self.trainsize[1])),
 flags=cv2.INTER_LINEAR)

 return image, im_info


def expand_crop(images, rect, expand_ratio=0.3):
 imgh, imgw, c = images.shape
 label, conf, xmin, ymin, xmax, ymax = [int(x) for x in rect.tolist()]
 if label != 0:
 return None, None, None
 org_rect = [xmin, ymin, xmax, ymax]
 h_half = (ymax - ymin) * (1 + expand_ratio) / 2.
 w_half = (xmax - xmin) * (1 + expand_ratio) / 2.
 if h_half > w_half * 4 / 3:
 w_half = h_half * 0.75
 center = [(ymin + ymax) / 2., (xmin + xmax) / 2.]
 ymin = max(0, int(center[0] - h_half))
 ymax = min(imgh - 1, int(center[0] + h_half))
 xmin = max(0, int(center[1] - w_half))
 xmax = min(imgw - 1, int(center[1] + w_half))
 return images[ymin:ymax, xmin:xmax, :], [xmin, ymin, xmax, ymax], org_rect