improve calculate_centroid

PINTO0309 · Dec 15, 2024 · edd03a1 · edd03a1
1 parent d1a2ae9
commit edd03a1
Showing 1 changed file with 3 additions and 7 deletions.
diff --git a/462_Gaze-LLE/demo/demo_yolov9_onnx_gazelle.py b/462_Gaze-LLE/demo/demo_yolov9_onnx_gazelle.py
@@ -1406,13 +1406,9 @@ def calculate_centroid(heatmap: np.ndarray) -> Tuple[int, int, float]:
                 Peak score for heat map score
             """
             # 1. ピーク値を求める
-            peak_value = np.max(heatmap)
-            # 2. ピーク値との差の絶対値を計算
-            diff = np.abs(heatmap - peak_value)
-            # 3. 差が最小の要素のインデックスを取得（1Dインデックス）
-            min_index = np.argmin(diff)
-            # 4. 1Dインデックスを2Dインデックス (y, x) に変換
-            y, x = np.unravel_index(min_index, heatmap.shape)
+            max_index = np.argmax(heatmap)
+            # 2. 1Dインデックスを2Dインデックス (y, x) に変換
+            y, x = np.unravel_index(max_index, heatmap.shape)
             return int(x), int(y), heatmap[y, x]
 
         for head_box, heatmap in zip(head_boxes, heatmaps):