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

Optional batched pose processing.

Browse Source
This commit is contained in:
Daniel
2024-12-18 16:22:08 +01:00
parent 7b8d209601
commit 07426fac2f
8 changed files with 151 additions and 75 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
scripts/test_skelda_dataset.py
View File

@ -53,6 +53,9 @@ default_min_match_score = 0.94
# If the number of cameras is high, and the views are not occluded, use a higher value
default_min_group_size = 1
# Batch poses per image for faster processing
# If most of the time only one person is in a image, disable it, because it is slightly slower then
default_batch_poses = True
datasets = {
"human36m": {
@ -62,6 +65,7 @@ datasets = {
"min_group_size": 1,
"min_bbox_score": 0.4,
"min_bbox_area": 0.1 * 0.1,
"batch_poses": False,
},
"panoptic": {
"path": "/datasets/panoptic/skelda/test.json",
@ -310,13 +314,14 @@ def main():
min_group_size = datasets[dataset_use].get("min_group_size", default_min_group_size)
min_bbox_score = datasets[dataset_use].get("min_bbox_score", default_min_bbox_score)
min_bbox_area = datasets[dataset_use].get("min_bbox_area", default_min_bbox_area)
batch_poses = datasets[dataset_use].get("batch_poses", default_batch_poses)
# Load 2D pose model
whole_body = test_triangulate.whole_body
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(min_bbox_score, min_bbox_area)
kpt_model = utils_2d_pose.load_model(min_bbox_score, min_bbox_area, batch_poses)
# Manually set matplotlib backend
try:

105
scripts/utils_2d_pose_ort.py
View File

@ -96,29 +96,12 @@ class BaseModel(ABC):
if "image" in iname:
ishape = list(self.input_shapes[i])
if "batch_size" in ishape:
if "TensorrtExecutionProvider" in self.providers:
# Using different images sizes for TensorRT warmup takes too long
ishape = [1, 1000, 1000, 3]
else:
ishape = [
1,
np.random.randint(300, 1000),
np.random.randint(300, 1000),
3,
]
max_batch_size = 10
ishape[0] = np.random.choice(
list(range(1, max_batch_size + 1))
)
tensor = np.random.random(ishape)
tensor = tensor * 255
elif "bbox" in iname:
tensor = np.array(
[
[
np.random.randint(30, 100),
np.random.randint(30, 100),
np.random.randint(200, 300),
np.random.randint(200, 300),
]
]
)
else:
raise ValueError("Undefined input type:", iname)
@ -401,35 +384,48 @@ class RTMPose(BaseModel):
self.target_size = (384, 288)
self.boxcrop = BoxCrop(self.target_size, padding_scale=1.25, fill_value=0)
def preprocess(self, image: np.ndarray, bbox: np.ndarray):
bbox = np.asarray(bbox)[0:4]
bbox += np.array([-0.5, -0.5, 0.5 - 1e-8, 0.5 - 1e-8])
bbox = bbox.round().astype(np.int32)
region = self.boxcrop.crop_resize_box(image, bbox)
tensor = np.asarray(region).astype(self.input_types[0], copy=False)
tensor = np.expand_dims(tensor, axis=0)
tensor = [tensor]
def preprocess(self, image: np.ndarray, bboxes: np.ndarray):
cutouts = []
for i in range(len(bboxes)):
bbox = np.asarray(bboxes[i])[0:4]
bbox += np.array([-0.5, -0.5, 0.5 - 1e-8, 0.5 - 1e-8])
bbox = bbox.round().astype(np.int32)
region = self.boxcrop.crop_resize_box(image, bbox)
tensor = np.asarray(region).astype(self.input_types[0], copy=False)
cutouts.append(tensor)
if len(bboxes) == 1:
cutouts = np.expand_dims(cutouts[0], axis=0)
else:
cutouts = np.stack(cutouts, axis=0)
tensor = [cutouts]
return tensor
def postprocess(
self, result: List[np.ndarray], image: np.ndarray, bbox: np.ndarray
self, result: List[np.ndarray], image: np.ndarray, bboxes: np.ndarray
):
scores = np.clip(result[1][0], 0, 1)
kp = np.concatenate([result[0][0], np.expand_dims(scores, axis=-1)], axis=-1)
kpts = []
for i in range(len(bboxes)):
scores = np.clip(result[1][i], 0, 1)
kp = np.concatenate(
[result[0][i], np.expand_dims(scores, axis=-1)], axis=-1
)
paddings, scale, bbox, _ = self.boxcrop.calc_params(image.shape, bbox)
kp[:, 0] -= paddings[0]
kp[:, 1] -= paddings[2]
kp[:, 0:2] /= scale
kp[:, 0] += bbox[0]
kp[:, 1] += bbox[1]
kp[:, 0:2] = np.maximum(kp[:, 0:2], 0)
max_w = image.shape[1] - 1
max_h = image.shape[0] - 1
kp[:, 0] = np.minimum(kp[:, 0], max_w)
kp[:, 1] = np.minimum(kp[:, 1], max_h)
paddings, scale, bbox, _ = self.boxcrop.calc_params(image.shape, bboxes[i])
kp[:, 0] -= paddings[0]
kp[:, 1] -= paddings[2]
kp[:, 0:2] /= scale
kp[:, 0] += bbox[0]
kp[:, 1] += bbox[1]
kp[:, 0:2] = np.maximum(kp[:, 0:2], 0)
max_w = image.shape[1] - 1
max_h = image.shape[0] - 1
kp[:, 0] = np.minimum(kp[:, 0], max_w)
kp[:, 1] = np.minimum(kp[:, 1], max_h)
kpts.append(kp)
return kp
return kpts
# ==================================================================================================
@ -444,6 +440,8 @@ class TopDown:
box_min_area: float,
warmup: int = 30,
):
self.batch_poses = bool("Bx" in pose_model_path)
self.det_model = RTMDet(
det_model_path, box_conf_threshold, box_min_area, warmup
)
@ -451,22 +449,29 @@ class TopDown:
def predict(self, image):
boxes = self.det_model(image=image)
if len(boxes) == 0:
return []
results = []
for i in range(boxes.shape[0]):
kp = self.pose_model(image=image, bbox=boxes[i])
results.append(kp)
if self.batch_poses:
results = self.pose_model(image=image, bboxes=boxes)
else:
for i in range(boxes.shape[0]):
kp = self.pose_model(image=image, bboxes=[boxes[i]])
results.append(kp[0])
return results
# ==================================================================================================
def load_model(min_bbox_score=0.3, min_bbox_area=0.1 * 0.1):
def load_model(min_bbox_score=0.3, min_bbox_area=0.1 * 0.1, batch_poses=False):
print("Loading 2D model ...")
model = TopDown(
"/RapidPoseTriangulation/extras/mmdeploy/exports/rtmdet-nano_320x320_fp16_extra-steps.onnx",
"/RapidPoseTriangulation/extras/mmdeploy/exports/rtmpose-m_384x288_fp16_extra-steps.onnx",
"/RapidPoseTriangulation/extras/mmdeploy/exports/rtmdet-nano_1x320x320x3_fp16_extra-steps.onnx",
f"/RapidPoseTriangulation/extras/mmdeploy/exports/rtmpose-m_{'B' if batch_poses else '1'}x384x288x3_fp16_extra-steps.onnx",
box_conf_threshold=min_bbox_score,
box_min_area=min_bbox_area,
warmup=30,