425_Gold-YOLO-Body-Head-Hand：Fixed class ID parsing.

PINTO0309 · Nov 30, 2023 · 8306d68 · 8306d68
1 parent 90a8c6e
commit 8306d68
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/425_Gold-YOLO-Body-Head-Hand/demo/demo_goldyolo_onnx.py b/425_Gold-YOLO-Body-Head-Hand/demo/demo_goldyolo_onnx.py
@@ -202,13 +202,14 @@ def __postprocess(
 
             if len(boxes_keep) > 0:
                 for box, score in zip(boxes_keep, scores_keep):
+                    class_id = int(box[1])
                     x_min = int(max(box[2], 0) * image_width / self.input_shapes[0][3])
                     y_min = int(max(box[3], 0) * image_height / self.input_shapes[0][2])
                     x_max = int(min(box[4], self.input_shapes[0][3]) * image_width / self.input_shapes[0][3])
                     y_max = int(min(box[5], self.input_shapes[0][2]) * image_height / self.input_shapes[0][2])
 
                     result_boxes.append(
-                        [x_min, y_min, x_max, y_max]
+                        [x_min, y_min, x_max, y_max, class_id]
                     )
                     result_scores.append(
                         score

diff --git a/425_Gold-YOLO-Body-Head-Hand/demo/demo_goldyolo_onnx_image.py b/425_Gold-YOLO-Body-Head-Hand/demo/demo_goldyolo_onnx_image.py
@@ -204,13 +204,14 @@ def __postprocess(
 
             if len(boxes_keep) > 0:
                 for box, score in zip(boxes_keep, scores_keep):
+                    class_id = int(box[1])
                     x_min = int(max(box[2], 0) * image_width / self.input_shapes[0][3])
                     y_min = int(max(box[3], 0) * image_height / self.input_shapes[0][2])
                     x_max = int(min(box[4], self.input_shapes[0][3]) * image_width / self.input_shapes[0][3])
                     y_max = int(min(box[5], self.input_shapes[0][2]) * image_height / self.input_shapes[0][2])
 
                     result_boxes.append(
-                        [x_min, y_min, x_max, y_max]
+                        [x_min, y_min, x_max, y_max, class_id]
                     )
                     result_scores.append(
                         score