Some small updates.

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

View File

@ -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(

View File

@ -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)

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,
)