YOLO目标检测：实现物体固定唯一ID分配的技术需求

当前基于YOLO的目标检测代码存在问题：每次处理检测帧时都会重置ID计数器，导致物体离开画面后，剩余物体的ID会被重新分配（比如原本ID为2、3、4的物体，在ID1的物体离开后会被改为1、2、3）。需要实现的效果是：已存在的物体保持固定ID，新进入画面的物体分配全新的唯一ID。

修改后的完整代码 #

import cv2
import numpy as np
import os
import yaml
from yaml.loader import SafeLoader
from scipy.optimize import linear_sum_assignment

class YOLO_Pred():
    def __init__(self, onnx_model, data_yaml):
        # 加载YAML配置
        with open(data_yaml, mode='r') as f:
            data_yaml = yaml.load(f, Loader=SafeLoader)

        self.labels = data_yaml['names']
        self.nc = data_yaml['nc']
        self.class_counts = {}

        # 加载YOLO模型
        self.yolo = cv2.dnn.readNetFromONNX(onnx_model)
        self.yolo.setPreferableBackend(cv2.dnn.DNN_BACKEND_OPENCV)
        self.yolo.setPreferableTarget(cv2.dnn.DNN_TARGET_CPU)

        # 跟踪相关变量初始化
        self.tracked_objects = {}  # 存储跟踪物体: {id: [bbox, last_seen_frame]}
        self.next_object_id = 1
        self.current_frame = 0
        self.max_disappeared = 5  # 物体消失超过此帧数后移除跟踪

    def calculate_iou(self, boxA, boxB):
        # 计算两个边界框的IOU（交并比）
        xA = max(boxA[0], boxB[0])
        yA = max(boxA[1], boxB[1])
        xB = min(boxA[0] + boxA[2], boxB[0] + boxB[2])
        yB = min(boxA[1] + boxA[3], boxB[1] + boxB[3])

        interArea = max(0, xB - xA + 1) * max(0, yB - yA + 1)
        boxAArea = (boxA[2] + 1) * (boxA[3] + 1)
        boxBArea = (boxB[2] + 1) * (boxB[3] + 1)

        iou = interArea / float(boxAArea + boxBArea - interArea)
        return iou

    def update_tracks(self, current_detections):
        # 匹配当前检测框与已跟踪物体，更新跟踪状态
        tracked_ids = list(self.tracked_objects.keys())
        tracked_boxes = [self.tracked_objects[id][0] for id in tracked_ids]

        # 构建IOU代价矩阵（用1-IOU表示匹配代价）
        cost_matrix = []
        for det_box in current_detections:
            iou_scores = [self.calculate_iou(det_box, track_box) for track_box in tracked_boxes]
            cost_matrix.append([1 - score for score in iou_scores])

        # 匈牙利算法完成最优匹配
        det_indices, track_indices = linear_sum_assignment(cost_matrix)

        # 更新匹配成功的物体状态
        matched_ids = []
        for det_idx, track_idx in zip(det_indices, track_indices):
            track_id = tracked_ids[track_idx]
            self.tracked_objects[track_id] = [current_detections[det_idx], self.current_frame]
            matched_ids.append(track_id)

        # 为未匹配的新检测框分配唯一ID
        for det_idx in range(len(current_detections)):
            if det_idx not in det_indices:
                self.tracked_objects[self.next_object_id] = [current_detections[det_idx], self.current_frame]
                matched_ids.append(self.next_object_id)
                self.next_object_id += 1

        # 移除长时间未出现的无效跟踪物体
        to_remove = []
        for track_id in tracked_ids:
            if self.current_frame - self.tracked_objects[track_id][1] > self.max_disappeared:
                to_remove.append(track_id)
        for track_id in to_remove:
            del self.tracked_objects[track_id]

        return matched_ids

    def predictions(self, image):
        self.current_frame += 1
        row, col, d = image.shape
        # 转换为正方形输入图像适配YOLO要求
        max_rc = max(row, col)
        input_image = np.zeros((max_rc, max_rc, 3), dtype=np.uint8)
        input_image[0:row, 0:col] = image
        # YOLO模型推理
        INPUT_WH_YOLO = 640
        blob = cv2.dnn.blobFromImage(input_image, 1/255, (INPUT_WH_YOLO, INPUT_WH_YOLO), swapRB=True, crop=False)
        self.yolo.setInput(blob)
        preds = self.yolo.forward()

        # 过滤低置信度检测结果
        detections = preds[0]
        boxes = []
        confidences = []
        classes = []

        image_w, image_h = input_image.shape[:2]
        x_factor = image_w / INPUT_WH_YOLO
        y_factor = image_h / INPUT_WH_YOLO

        for i in range(len(detections)):
            row = detections[i]
            confidence = row[4]
            if confidence > 0.4:
                class_score = row[5:].max()
                class_id = row[5:].argmax()
                if class_score > 0.25:
                    cx, cy, w, h = row[0:4]
                    left = int((cx - 0.5 * w) * x_factor)
                    top = int((cy - 0.5 * h) * y_factor)
                    width = int(w * x_factor)
                    height = int(h * y_factor)
                    box = np.array([left, top, width, height])
                    confidences.append(confidence)
                    boxes.append(box)
                    classes.append(class_id)

        # NMS（非极大值抑制）去除重复检测框
        boxes_np = np.array(boxes).tolist()
        confidences_np = np.array(confidences).tolist()
        index = np.array(cv2.dnn.NMSBoxes(boxes_np, confidences_np, 0.25, 0.45)).flatten()

        # 提取NMS处理后的有效检测结果
        nms_boxes = [boxes_np[ind] for ind in index]
        nms_classes = [classes[ind] for ind in index]
        nms_confidences = [confidences_np[ind] for ind in index]

        # 更新跟踪状态，获取每个检测框对应的固定ID
        matched_ids = self.update_tracks(nms_boxes)

        # 绘制边界框、ID和类别信息
        for idx, ind in enumerate(index):
            x, y, w, h = boxes_np[ind]
            bb_conf = int(nms_confidences[idx] * 100)
            classes_id = nms_classes[idx]
            class_name = self.labels[classes_id]
            color = (0, 0, 255)

            # 头盔（ID0）和背心（ID2）用绿色框，其他用红色框
            if classes_id == 0 or classes_id == 2:
                color = (0, 255, 0)
            else:
                # 保存非头盔/背心类别的检测图像
                folder_name = "SAVED_DATA"
                if not os.path.exists(folder_name):
                    os.makedirs(folder_name)
                count = self.class_counts.get(class_name, 0) + 1
                self.class_counts[class_name] = count
                file_name = f"{folder_name}/{class_name}_{count}.jpg"
                while os.path.exists(file_name):
                    count += 1
                    file_name = f"{folder_name}/{class_name}_{count}.jpg"
                cv2.imwrite(file_name, image)

            object_id = matched_ids[idx]
            id_text = f'ID: {object_id}'

            cv2.rectangle(image, (x, y), (x + w, y + h), color, 5)
            cv2.rectangle(image, (x - 120, y - 30), (x, y), color, -1)
            cv2.putText(image, id_text, (x - 100, y - 10), cv2.FONT_HERSHEY_PLAIN, 1.5, (0, 255, 255), 2)
            cv2.putText(image, f'{class_name}: {bb_conf}%', (x, y - 10), cv2.FONT_HERSHEY_PLAIN, 1.5, (0, 0, 34), 2)

        return image

核心修改说明 #

• 添加跟踪状态存储：在初始化方法中创建tracked_objects字典，保存每个跟踪物体的ID、边界框和最后出现的帧号；next_object_id记录下一个可用的唯一ID，避免重复分配。
• IOU匹配与匈牙利算法：实现calculate_iou计算边界框重叠度，用linear_sum_assignment（匈牙利算法）完成当前检测框与历史跟踪物体的最优匹配，确保同一物体被分配固定ID。
• 跟踪更新逻辑：
  • 匹配成功的物体更新其边界框和最后出现帧号
  • 未匹配的新检测框分配全新ID
  • 移除超过max_disappeared帧未出现的物体，清理无效跟踪数据
• 移除ID重置逻辑：删除原代码中每次检测都重置object_id的操作，改用跟踪模块分配的固定ID。

内容的提问来源于stack exchange，提问作者Ashutosh Gairola