crosstyan/RapidPoseTriangulation
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Some small updates.

Browse Source
This commit is contained in:
Daniel
2024-12-16 14:34:44 +01:00
parent 07f75c53ee
commit 17ff41010a
3 changed files with 25 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

31
scripts/utils_2d_pose_ort.py
View File

@ -302,6 +302,7 @@ class RTMDet(BaseModel):
self,
model_path: str,
conf_threshold: float,
min_area_fraction: float,
warmup: int = 30,
):
super(RTMDet, self).__init__(model_path, warmup)
@ -309,9 +310,8 @@ class RTMDet(BaseModel):
self.conf_threshold = conf_threshold
self.letterbox = LetterBox(self.target_size, fill_value=114)
min_area_scale = 0.025 * 0.025
img_area = self.target_size[0] * self.target_size[1]
self.min_area = img_area * min_area_scale
self.min_area = img_area * min_area_fraction
def preprocess(self, image: np.ndarray):
image = self.letterbox.resize_image(image)
@ -330,11 +330,18 @@ class RTMDet(BaseModel):
keep = boxes[:, 4] > self.conf_threshold
boxes = boxes[keep]
if len(boxes) == 0:
return np.array([])
# Drop boxes with too small area
boxes = boxes.astype(np.float32)
areas = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
keep = areas >= self.min_area
boxes = boxes[keep]
if len(boxes) == 0:
return np.array([])
paddings, scale, _ = self.letterbox.calc_params(image.shape)
boxes[:, 0] -= paddings[0]
@ -403,17 +410,20 @@ class RTMPose(BaseModel):
class TopDown:
def __init__(
self,
det_model_path,
pose_model_path,
box_conf_threshold=0.6,
warmup=30,
det_model_path: str,
pose_model_path: str,
box_conf_threshold: float,
box_min_area: float,
warmup: int = 30,
):
if (not det_model_path.endswith(".onnx")) or (
not pose_model_path.endswith(".onnx")
):
raise ValueError("Only ONNX models are supported.")
self.det_model = RTMDet(det_model_path, box_conf_threshold, warmup)
self.det_model = RTMDet(
det_model_path, box_conf_threshold, box_min_area, warmup
)
self.pose_model = RTMPose(pose_model_path, warmup)
def predict(self, image):
@ -428,15 +438,14 @@ class TopDown:
# ==================================================================================================
def load_model():
def load_model(min_bbox_score=0.3, min_bbox_area=0.1 * 0.1):
print("Loading onnx model ...")
model = TopDown(
# "/RapidPoseTriangulation/extras/mmdeploy/exports/rtmdet-nano_320x320_extra-steps.onnx",
"/RapidPoseTriangulation/extras/mmdeploy/exports/rtmdet-nano_320x320_fp16_extra-steps.onnx",
# "/RapidPoseTriangulation/extras/mmdeploy/exports/rtmpose-m_384x288_extra-steps.onnx",
"/RapidPoseTriangulation/extras/mmdeploy/exports/rtmpose-m_384x288_fp16_extra-steps.onnx",
box_conf_threshold=0.3,
box_conf_threshold=min_bbox_score,
box_min_area=min_bbox_area,
warmup=30,
)