From 17ff41010a1170ee02dfe9c36d9895950920f53d Mon Sep 17 00:00:00 2001 From: Daniel Date: Mon, 16 Dec 2024 14:34:44 +0100 Subject: [PATCH] Some small updates. --- scripts/test_triangulate.py | 4 ++-- scripts/utils_2d_pose.py | 6 +++--- scripts/utils_2d_pose_ort.py | 31 ++++++++++++++++++++----------- 3 files changed, 25 insertions(+), 16 deletions(-) diff --git a/scripts/test_triangulate.py b/scripts/test_triangulate.py index 364b992..831a7e6 100644 --- a/scripts/test_triangulate.py +++ b/scripts/test_triangulate.py @@ -283,7 +283,7 @@ def main(): if any((whole_body[k] for k in whole_body)): kpt_model = utils_2d_pose.load_wb_model() else: - kpt_model = utils_2d_pose.load_model() + kpt_model = utils_2d_pose.load_model(min_bbox_score=0.3) # Manually set matplotlib backend matplotlib.use("TkAgg") @@ -340,7 +340,7 @@ def main(): else: cameras = rpt.convert_cameras(camparams) roomp = [roomparams["room_size"], roomparams["room_center"]] - triangulator = rpt.Triangulator(min_match_score=0.95) + triangulator = rpt.Triangulator(min_match_score=0.94) stime = time.time() poses_3d = triangulator.triangulate_poses( diff --git a/scripts/utils_2d_pose.py b/scripts/utils_2d_pose.py index 6f7edd4..b94e416 100644 --- a/scripts/utils_2d_pose.py +++ b/scripts/utils_2d_pose.py @@ -46,7 +46,7 @@ def load_wb_model(): def get_2d_pose(model, imgs, num_joints=17): """See: https://mmpose.readthedocs.io/en/latest/user_guides/inference.html#basic-usage""" - result_generator = model(imgs, show=False) + result_generator = model(imgs, show=False, bbox_thr=0.3, nms_thr=0.3) new_poses = [] for _ in range(len(imgs)): result = next(result_generator) @@ -56,8 +56,8 @@ def get_2d_pose(model, imgs, num_joints=17): kpts = result["predictions"][0][i]["keypoints"] scores = result["predictions"][0][i]["keypoint_scores"] - kpts = np.array(kpts) - scores = np.array(scores).reshape(-1, 1) + kpts = np.asarray(kpts) + scores = np.asarray(scores).reshape(-1, 1) scores = np.clip(scores, 0, 1) pose = np.concatenate((kpts, scores), axis=-1) poses.append(pose) diff --git a/scripts/utils_2d_pose_ort.py b/scripts/utils_2d_pose_ort.py index 0cbc875..8c336f0 100644 --- a/scripts/utils_2d_pose_ort.py +++ b/scripts/utils_2d_pose_ort.py @@ -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, )