Fixed running custom onnx models.

2024-11-29 15:18:57 +01:00
parent f6d13ea5a7
commit 93d4611a91
6 changed files with 158 additions and 32 deletions
@@ -177,18 +177,22 @@ def get_real_keypoints(keypoints: np.ndarray, heatmaps: np.ndarray, img_size: Se
    return keypoints


-def simcc_decoder(simcc_x: np.ndarray,
-                  simcc_y: np.ndarray,
-                  input_size: Sequence[int],
-                  dx: int,
-                  dy: int,
-                  scale: float):
+def simcc_decoder(
+    simcc_x: np.ndarray,
+    simcc_y: np.ndarray,
+    input_size: Sequence[int],
+    dx: int,
+    dy: int,
+    scale: float,
+):
+    # See: /mmpose/codecs/utils/post_processing.py - get_simcc_maximum()
+
    x = np.argmax(simcc_x, axis=-1, keepdims=True).astype(np.float32)
    y = np.argmax(simcc_y, axis=-1, keepdims=True).astype(np.float32)

    x_conf = np.max(simcc_x, axis=-1, keepdims=True)
    y_conf = np.max(simcc_y, axis=-1, keepdims=True)
-    conf = (x_conf + y_conf) / 2
+    conf = np.minimum(x_conf, y_conf)

    x /= simcc_x.shape[-1]
    y /= simcc_y.shape[-1]