Skip to content

Commit 2528127

Browse files
committed
Support all resize types in kp model
1 parent 5030c5a commit 2528127

File tree

1 file changed

+13
-1
lines changed

1 file changed

+13
-1
lines changed

src/python/model_api/models/keypoint_detection.py

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -55,7 +55,19 @@ def postprocess(
5555
orig_h, orig_w = meta["original_shape"][:2]
5656
kp_scale_h = orig_h / self.h
5757
kp_scale_w = orig_w / self.w
58-
batch_keypoints = batch_keypoints.squeeze() * np.array([kp_scale_w, kp_scale_h])
58+
59+
batch_keypoints = batch_keypoints.squeeze()
60+
61+
if self.resize_type in ["fit_to_window", "fit_to_window_letterbox"]:
62+
inverted_scale = max(kp_scale_h, kp_scale_w)
63+
kp_scale_h = kp_scale_w = inverted_scale
64+
if self.resize_type == "fit_to_window_letterbox":
65+
pad_left = (self.w - round(orig_w / inverted_scale)) // 2
66+
pad_top = (self.h - round(orig_h / inverted_scale)) // 2
67+
batch_keypoints -= np.array([pad_left, pad_top])
68+
69+
batch_keypoints *= np.array([kp_scale_w, kp_scale_h])
70+
5971
return DetectedKeypoints(batch_keypoints, batch_scores.squeeze())
6072

6173
@classmethod

0 commit comments

Comments
 (0)