[Fix] Modify the dtype of heatmap_weights to speed up the calculation (…

…#2486)
open-mmlab · Jan 3, 2024 · 6f7411c · 6f7411c
1 parent df0a374
commit 6f7411c
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/mmpose/codecs/spr.py b/mmpose/codecs/spr.py
@@ -138,7 +138,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