From 38a43fae760b4b7b51cc5143897fef81415bfa17 Mon Sep 17 00:00:00 2001 From: Yinlei Sun Date: Mon, 26 Jun 2023 20:52:46 +0800 Subject: [PATCH] Modify the dtype of heatmap_weights to speed up the calculation --- mmpose/codecs/spr.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mmpose/codecs/spr.py b/mmpose/codecs/spr.py index add6f5715b..b974ce9ebe 100644 --- a/mmpose/codecs/spr.py +++ b/mmpose/codecs/spr.py @@ -131,7 +131,7 @@ def _get_heatmap_weights(self, Returns: np.ndarray: Heatmap weight array in the same shape with heatmaps """ - heatmap_weights = np.ones(heatmaps.shape) * bg_weight + heatmap_weights = np.ones(heatmaps.shape, dtype=np.float32) * bg_weight heatmap_weights[heatmaps > 0] = fg_weight return heatmap_weights