72 lines
3.0 KiB
Python

import sys, os
from base.postprocess import postprocess_hw
from .rsn_affine_postprocess import postprocess_
import numpy as np
class RsnAffineRunner:
def __init__(self,image_size=[384,288], scale_ext = [1.09,1.135], **kwargs):
"""
:param lib_path:
a folder that include all dependencies scripts of building block classes. (optional)
:param model_def_path:
a script that instantiates `target_model` based on the imported building blocks.
if `lib_path` is not defined, one should difine all building block classes in this script
before instantiating `target_model`
"""
self.init_config = locals()
self.init_config.update(kwargs)
self.image_size = image_size
self.scale_ext = scale_ext
def _exec_py_script(self, model_def_path):
"""
a script that import `target_model` into global() so that the runner can access
"""
with open(model_def_path, 'r') as f:
lines = f.readlines()
script_content = "".join(lines)
if 'target_model' in script_content:
self._target_model_instantiate(model_def_path)
else:
# declare everything from the script into global variable, dangerous!
# this part will be deprecated in near future
exec(script_content, globals())
'''
return affine transformation for pose estimation models and convert bounding boxes to kpt model input
:param pre_output: list of bounding box of size 4+ [[x,y,w,h,.....],[x2,y2,w2,h2,.....],.....]
:param p2: point 2
:return: distance between input pts
:param pre_output: list of bounding boxes with 4 + elements and format x, y, w, h [[x1,y1,w1,h1,....], [x2,y2,w2,h2,....],....]
:return: output of target image size. bounding box will be converted to output of dim 1 x image_size[0] x image_size[1] x 3
'''
def run(self,img,pre_results,**kwargs):
if len(pre_results)==0:
return [],[]
bboxes = []
infer_config = {
}
infer_config.update(self.init_config)
for bbox in pre_results:
bbox = bbox[:4]
bboxes.append([int(float(bbox[0])),int(float(bbox[1])),int(float(bbox[2])),int(float(bbox[3]))])
if len(bboxes)==0:
return [],[]
results=[]
results_info = []
center_bbox = 0
pre_config = {
}
pre_config.update(self.init_config)
pre_config.update(kwargs)
for i,bbox in enumerate(bboxes):
#no pre_config or preproc technically
post_config = {
"bbox":bbox,
"image_size":self.image_size,
"scale_ext":self.scale_ext
}
post_config.update(self.init_config)
post_config.update(kwargs)
postproc,post_info= postprocess_hw(img, postprocess_, **post_config)
results.append(postproc)
results_info.append(post_info)
return results,results_info