86 lines
3.2 KiB
Python

import numpy as np
import kneron_preprocessing
"""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].
output_size (array): output size of network input
Returns:
np.ndarray: The transform matrix.
"""
def get_affine_transform(center, scale, output_size):
if not isinstance(scale, np.ndarray) and not isinstance(scale, list):
scale = np.array([scale, scale])
src_w = scale[0]
dst_w = output_size[1]
dst_h = output_size[0]
src_dir =[0.0, src_w * -0.5]
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
src[1, :] = center + src_dir
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, :])
trans = kneron_preprocessing.similarity_transform(np.float32(dst).flatten().tolist(),np.float32(src).flatten().tolist(),type='float')
return trans
def get_3rd_point(a, b):
direct = a - b
return b + np.array([-direct[1], direct[0]], dtype=np.float32)
def _bbox_to_center_and_scale( bbox):
x, y, w, h = bbox
center = np.zeros(2, dtype=np.float32)
center[0] = x + w / 2.0
center[1] = y + h / 2.0
scale = np.array([w , h],
dtype=np.float32)
return center, scale
'''
return preprocess affine tranfor,
:param input_img: image path
:param bbox: bounding box of size 4+ [x,y,w,h,.....]
:param image_size: determines size of output
:return: preprocessed input of size (1, image_size 0, image_size 1, 3)
'''
def postprocess_(input_img,bbox,image_size=[384,288],scale_ext=[1.09,1.135],**kwargs):
aspect_ratio = image_size[1]/image_size[0]
if isinstance(input_img, str):
try:
data_numpy = kneron_preprocessing.API.load_image(input_img)
except:
try:
data_numpy = kneron_preprocessing.API.load_bin(input_img, **kwargs)
except:
print('input format error')
assert 0
else:
assert isinstance(image, np.ndarray)
if data_numpy is None:
raise ValueError('=> fail to read {}'.format(image_file))
center,scale = _bbox_to_center_and_scale(bbox)
scale[0] *= scale_ext[0]
scale[1] *= scale_ext[1]
# fit the ratio
if scale[0] > aspect_ratio * scale[1]:
scale[1] = scale[0] * 1.0 / aspect_ratio
else:
scale[0] = scale[1] * 1.0 * aspect_ratio
trans = get_affine_transform(center, scale, image_size)
height = data_numpy.shape[0]
width = data_numpy.shape[1]
pre_info = {'centers':[center.tolist()],'scales':[scale.tolist()],'height':height,'width':width}
img = kneron_preprocessing.warpAffine(data_numpy,Matrix=trans,warp_size= (int(image_size[1]), int(image_size[0]))) #reversed in affine for new one
img = kneron_preprocessing.API.norm(img)
img = img[np.newaxis,:,:,:]
return img,pre_info