From 5fc3c249f23b200a8a9c3e8709539503d860e5c5 Mon Sep 17 00:00:00 2001 From: "Tanner, Gilbert" Date: Wed, 6 May 2026 01:47:31 +0200 Subject: [PATCH] Fix: single-person inference picks largest-area detection instead of highest-confidence Signed-off-by: Tanner, Gilbert --- main/inference.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/main/inference.py b/main/inference.py index fb2a5a1..f5872e9 100644 --- a/main/inference.py +++ b/main/inference.py @@ -13,7 +13,6 @@ import cv2 from tqdm import tqdm import json -from typing import Literal, Union from mmdet.apis import init_detector, inference_detector from utils.inference_utils import process_mmdet_results, non_max_suppression @@ -36,6 +35,7 @@ def parse_args(): parser.add_argument('--multi_person', action="store_true") parser.add_argument('--iou_thr', type=float, default=0.5) parser.add_argument('--bbox_thr', type=int, default=50) + parser.add_argument('--min_score', type=float, default=0.3) args = parser.parse_args() return args @@ -82,7 +82,14 @@ def main(): ## mmdet inference mmdet_results = inference_detector(model, img_path) - mmdet_box = process_mmdet_results(mmdet_results, cat_id=0, multi_person=True) + person_dets = [d for d in mmdet_results[0] if (len(d) >= 5 and d[4] >= args.min_score)] + if not multi_person and person_dets: + # Sort by score + person_dets_sorted = sorted(person_dets, key=lambda d: -d[4]) + mmdet_box = [person_dets_sorted] + else: + mmdet_results = [np.asarray(person_dets)] + list(mmdet_results[1:]) + mmdet_box = process_mmdet_results(mmdet_results, cat_id=0, multi_person=True) # save original image if no bbox if len(mmdet_box[0])<1: