mirror of
https://github.com/deepinsight/insightface.git
synced 2025-12-30 08:02:27 +00:00
123 lines
3.6 KiB
Python
123 lines
3.6 KiB
Python
#!/usr/bin/env python
|
|
# -*- encoding: utf-8 -*-
|
|
"""
|
|
@Author : Qingping Zheng
|
|
@Contact : qingpingzheng2014@gmail.com
|
|
@File : transforms.py
|
|
@Time : 10/01/21 00:00 PM
|
|
@Desc :
|
|
@License : Licensed under the Apache License, Version 2.0 (the "License");
|
|
@Copyright : Copyright 2022 The Authors. All Rights Reserved.
|
|
"""
|
|
from __future__ import absolute_import
|
|
from __future__ import division
|
|
from __future__ import print_function
|
|
|
|
|
|
import numpy as np
|
|
import cv2
|
|
|
|
|
|
def flip_back(output_flipped, matched_parts):
|
|
'''
|
|
ouput_flipped: numpy.ndarray(batch_size, num_joints, height, width)
|
|
'''
|
|
assert output_flipped.ndim == 4,\
|
|
'output_flipped should be [batch_size, num_joints, height, width]'
|
|
|
|
output_flipped = output_flipped[:, :, :, ::-1]
|
|
|
|
for pair in matched_parts:
|
|
tmp = output_flipped[:, pair[0], :, :].copy()
|
|
output_flipped[:, pair[0], :, :] = output_flipped[:, pair[1], :, :]
|
|
output_flipped[:, pair[1], :, :] = tmp
|
|
|
|
return output_flipped
|
|
|
|
|
|
def transform_parsing(pred, center, scale, width, height, input_size):
|
|
|
|
if center is not None:
|
|
trans = get_affine_transform(center, scale, 0, input_size, inv=1)
|
|
target_pred = cv2.warpAffine(
|
|
pred,
|
|
trans,
|
|
(int(width), int(height)), #(int(width), int(height)),
|
|
flags=cv2.INTER_NEAREST,
|
|
borderMode=cv2.BORDER_CONSTANT,
|
|
borderValue=(0))
|
|
else:
|
|
target_pred = cv2.resize(pred, (int(width), int(height)), interpolation=cv2.INTER_NEAREST)
|
|
|
|
return target_pred
|
|
|
|
|
|
def get_affine_transform(center,
|
|
scale,
|
|
rot,
|
|
output_size,
|
|
shift=np.array([0, 0], dtype=np.float32),
|
|
inv=0):
|
|
if not isinstance(scale, np.ndarray) and not isinstance(scale, list):
|
|
print(scale)
|
|
scale = np.array([scale, scale])
|
|
|
|
scale_tmp = scale
|
|
|
|
src_w = scale_tmp[0]
|
|
dst_w = output_size[1]
|
|
dst_h = output_size[0]
|
|
|
|
rot_rad = np.pi * rot / 180
|
|
src_dir = get_dir([0, src_w * -0.5], rot_rad)
|
|
dst_dir = np.array([0, dst_w * -0.5], np.float32)
|
|
|
|
src = np.zeros((3, 2), dtype=np.float32)
|
|
dst = np.zeros((3, 2), dtype=np.float32)
|
|
src[0, :] = center + scale_tmp * shift
|
|
src[1, :] = center + src_dir + scale_tmp * shift
|
|
dst[0, :] = [dst_w * 0.5, dst_h * 0.5]
|
|
dst[1, :] = np.array([dst_w * 0.5, dst_h * 0.5]) + dst_dir
|
|
|
|
src[2:, :] = get_3rd_point(src[0, :], src[1, :])
|
|
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 affine_transform(pt, t):
|
|
new_pt = np.array([pt[0], pt[1], 1.]).T
|
|
new_pt = np.dot(t, new_pt)
|
|
return new_pt[:2]
|
|
|
|
|
|
def get_3rd_point(a, b):
|
|
direct = a - b
|
|
return b + np.array([-direct[1], direct[0]], dtype=np.float32)
|
|
|
|
|
|
def get_dir(src_point, rot_rad):
|
|
sn, cs = np.sin(rot_rad), np.cos(rot_rad)
|
|
|
|
src_result = [0, 0]
|
|
src_result[0] = src_point[0] * cs - src_point[1] * sn
|
|
src_result[1] = src_point[0] * sn + src_point[1] * cs
|
|
|
|
return src_result
|
|
|
|
|
|
def crop(img, center, scale, output_size, rot=0):
|
|
trans = get_affine_transform(center, scale, rot, output_size)
|
|
|
|
dst_img = cv2.warpAffine(img,
|
|
trans,
|
|
(int(output_size[1]), int(output_size[0])),
|
|
flags=cv2.INTER_LINEAR)
|
|
|
|
return dst_img
|