<think>我们正在使用一个基于TensorRT加速的YOLOv8推理代码，目前代码只支持检测一个类别（drone）。现在需要扩展为多类别检测。
从代码中可以看出，当前代码中定义了几个常量：
POSE_NUM = 17 * 3 # 关键点检测的数量（每个点3个值）
DET_NUM = 6 # 每个检测框的原始输出数量（x, y, w, h, conf, class_id）? 注意：实际上YOLOv8的输出格式可能不同
SEG_NUM = 32 # 分割掩码的参数数量
OBB_NUM = 1 # 定向边界框（OBB）的参数，这里可能是1个角度？
在`post_process`函数中，我们注意到：
num_values_per_detection = DET_NUM + SEG_NUM + POSE_NUM + OBB_NUM
这个值表示每个检测结果的参数总数。
然而，在标准的YOLOv8目标检测中（不带分割和姿态估计），每个检测框的输出应该是：4个坐标（x,y,w,h） + 1个置信度 + (类别数)个类别分数。
但是，在代码中，DET_NUM被设置为6，这可能是用于一个类别的情况（4坐标+1置信度+1类别分数）？但多类别情况下，类别分数应该是多个。
实际上，YOLOv8的TensorRT输出格式取决于模型。在导出模型时，如果是检测模型，输出维度为[1, 84, 8400]（84=4+80，80是COCO类别数）。
但是，我们当前代码中，每个检测框的参数数量是固定的（由DET_NUM+SEG_NUM+...等构成）。这暗示我们使用的可能是YOLOv8的集成输出（同时包含检测、分割、姿态等）？
然而，用户的问题中只提到目标检测，并且当前代码中只处理了检测框（没有分割和姿态）。因此，我们需要明确：
1. 我们使用的模型是什么？是纯检测模型，还是带有其他任务（如分割、姿态）的模型？
2. 当前代码中，DET_NUM=6，这显然只适用于单类别（4坐标+1置信度+1类别分数）或者多类别但只取最大类别分数的情况？
观察代码中的`post_process`函数：
pred = np.reshape(output[1:], (-1, num_values_per_detection))[:num, :]
然后，在`non_max_suppression`函数中，我们只使用了前6个值（x,y,w,h,conf,class_id）? 因为boxes = prediction[prediction[:, 4] >= conf_thres] 中，第4个是置信度。
实际上，在代码中，每个检测框的参数包括：
[x, y, w, h, conf, class_id, ... (其他参数，如分割、姿态等)]
因此，对于多类别检测，我们需要改变的是：
- 模型输出中，每个检测框的类别分数部分应该是多个分数（每个类别一个分数），而不是一个类别id。
但是，在代码中，我们注意到在`post_process`函数中，我们只取每个检测框的前DET_NUM（6）个参数。这6个参数中，第5个（索引4）是置信度，第6个（索引5）是类别id吗？
然而，在标准的多类别YOLOv8中，每个检测框的输出应该是：
[x, y, w, h, obj_conf, class_conf_0, class_conf_1, ... , class_conf_n]
然后，总的置信度 = obj_conf * max(class_conf) ？而类别id是argmax(class_conf)。
因此，当前代码中DET_NUM=6，只能处理5个类别（因为4坐标+1个obj_conf+5个类别分数=10，但这里DET_NUM=6）？这显然不对。
重新审视：在代码中，我们定义：
DET_NUM = 6
但实际在YOLOv8的TensorRT输出中，每个检测框的向量长度应该是：4（坐标）+ 1（obj_conf） + 类别数（class_conf）。因此，DET_NUM应该等于5+类别数。
那么，为什么这里DET_NUM=6？因为当前模型是单类别（drone）？所以类别数=1，那么总长度=4+1+1=6。
所以，要支持多类别，我们需要：
1. 确保模型是多类别模型（在训练时指定了多个类别）。
2. 修改代码中的DET_NUM，使其等于5+类别数（因为4坐标+1个obj_conf+类别数个类别分数）。
3. 修改后处理部分，正确解析类别分数。
但是，在代码中，我们注意到：
num_values_per_detection = DET_NUM + SEG_NUM + POSE_NUM + OBB_NUM
因此，如果我们修改了DET_NUM，那么整个向量长度就变了。但是，在代码中，我们只使用了每个检测框的前DET_NUM个值（在`non_max_suppression`函数中，我们只取prediction[:, :DET_NUM]？并不是，我们直接取了整个prediction，但在非极大值抑制中只用了前6个值（通过boxes = prediction[prediction[:,4]>=conf_thres]））。
实际上，在`non_max_suppression`函数中，我们传入的是整个预测数组（每个检测框有num_values_per_detection个值），但在函数内部，我们只使用了前6个（即x,y,w,h,conf,class_id）？但是，在单类别情况下，第5个（索引4）是置信度，第6个（索引5）是类别id（只有0，因为只有一个类别）。
对于多类别，我们需要改变的是：在`non_max_suppression`函数中，我们不应该直接使用第5个（索引5）作为类别id，而应该从第5个位置开始（索引4是obj_conf，然后后面是类别分数）计算每个类别的分数，然后取最大值作为类别置信度，并得到类别id。
因此，我们需要重构后处理部分。
步骤：
1. 修改DET_NUM：根据类别数，DET_NUM = 5 + num_classes（因为4坐标+1个obj_conf+num_classes个类别分数）。
但是，注意：在模型导出时，这个值已经固定了。所以，我们应该根据实际模型来设置这个值。
2. 在`post_process`函数中，我们得到每个检测框的向量（长度为num_values_per_detection），然后我们只取前DET_NUM个值（即检测部分），然后我们需要解析这个向量，得到每个类别的分数。
3. 修改`non_max_suppression`函数，使其能够处理多类别。
具体修改如下：
首先，我们需要知道类别数。可以在代码开头定义一个全局变量：
NUM_CLASSES = 80 # 例如，COCO有80个类别
然后，修改：
DET_NUM = 5 + NUM_CLASSES # 4坐标 + 1个obj_conf + NUM_CLASSES个类别分数
注意：在YOLOv8的原始输出中，每个检测框的格式是：
[x, y, w, h, obj_conf, class_conf_0, class_conf_1, ... , class_conf_{NUM_CLASSES-1}]
然后，在`non_max_suppression`函数中，我们需要做如下修改：
- 首先，从每个检测框的向量中取出前DET_NUM个值，然后分离出坐标（前4个）、obj_conf（第5个）和类别分数（后面的NUM_CLASSES个）。
- 计算每个类别的置信度：conf = obj_conf * class_conf_i
- 然后，我们保留那些最大类别置信度超过阈值的框，并记录其类别id和置信度。
但是，当前代码中，我们直接使用了prediction[:,4]作为置信度，这实际上只是obj_conf（没有乘以类别分数）。而且，prediction[:,5]作为类别id，这只有在单类别情况下正确（因为只有一个类别分数，所以类别id总是0）。
因此，我们需要重写`non_max_suppression`函数，使其能够处理多类别。
另外，注意在原始代码中，我们传入`non_max_suppression`函数的prediction参数是每个检测框的完整向量（包括检测、分割、姿态等），但我们只使用前DET_NUM个值（即检测部分）。所以，在函数内部，我们只处理检测部分。
修改后的`non_max_suppression`函数（多类别版本）：
步骤：
1. 提取每个检测框的检测部分（前DET_NUM个值）
2. 分离坐标（0:4）、obj_conf（4:5）、类别分数（5:5+NUM_CLASSES）
3. 计算每个类别的置信度 = obj_conf * 类别分数（每个类别分别乘）
4. 找到每个框的最大置信度以及对应的类别id
5. 根据最大置信度进行初步筛选（大于conf_thres）
6. 然后，对每个类别分别进行NMS（或者使用多类别NMS）
但是，为了简单和效率，我们通常采用：
- 先对每个类别进行循环，然后分别做NMS，最后合并。
- 或者，使用单次多类别NMS（如cv2.dnn.NMSBoxes支持多类别）。
然而，由于我们使用的是自己的后处理，我们可以采用以下方法：
1. 先创建一个空列表，用于存放所有保留的框。
2. 遍历每个类别：
- 取出当前类别的置信度
- 筛选出大于阈值的框
- 对这些框进行NMS（使用同一类别内的框）
- 将保留的框加入列表
3. 或者，我们也可以使用另一种常见做法：先找出每个框最大分数对应的类别，然后按照这个类别分组，再对每个类别做NMS。
但是，这样效率较低。另一种高效做法是：先选出每个框的最大分数，然后按照这个最大分数进行排序，然后进行跨类别的NMS（即不同类别的框也参与抑制，但通常不同类别的框不互相抑制？）—— 实际上，不同类别的框即使重叠也不应该被抑制，所以应该按类别分别进行NMS。
然而，在目标检测中，通常NMS是每个类别独立进行的。因此，我们需要按类别处理。
但是，由于我们的检测框数量可能很大，且类别较多，循环处理可能较慢。我们可以使用矩阵操作来避免循环，但这里为了清晰，我们使用循环每个类别的方式。
修改`non_max_suppression`函数：
注意：这个函数现在需要额外的参数`num_classes`，并且需要知道DET_NUM（即检测框向量的长度）？实际上，我们已经在函数内部使用了全局的NUM_CLASSES和DET_NUM。
但是，为了减少全局变量，我们可以将类别数作为参数传入。
由于修改较大，我们重新设计`non_max_suppression`函数：
输入：
prediction: 一个数组，形状为[N, DET_NUM+SEG_NUM+...]，但我们现在只关心前DET_NUM个值。
origin_h, origin_w: 原始图像尺寸
conf_thres: 置信度阈值
iou_thres: IOU阈值
num_classes: 类别数
输出：
一个数组，每行是[x1, y1, x2, y2, conf, class_id]
步骤：
1. 提取检测部分（前DET_NUM列） -> dets: [N, DET_NUM]
2. 分离：boxes = dets[:, :4] # [N,4]
obj_conf = dets[:, 4] # [N]
cls_scores = dets[:, 5:5+num_classes] # [N, num_classes]
3. 计算类别置信度：class_conf = obj_conf.reshape(-1,1) * cls_scores [N, num_classes]
4. 找到每个框的最大类别置信度：max_conf = np.max(class_conf, axis=1) [N]
class_id = np.argmax(class_conf, axis=1) [N]
5. 初步筛选：保留max_conf > conf_thres的框索引
mask = max_conf > conf_thres
boxes = boxes[mask]
class_id = class_id[mask]
max_conf = max_conf[mask]
6. 将坐标转换到原始图像（使用xywh2xyxy函数，但注意这个函数原本是针对单框向量设计的，现在boxes是多个框）
7. 对每个类别进行循环，进行NMS：
总保留框列表：keep_boxes = []
对每个类别id c：
从当前类别的框中，选出所有类别id等于c的框（注意：我们上面已经按最大类别筛选，所以每个框的类别就是class_id）
然后，对当前类别的框进行NMS（使用IOU阈值）
将保留的框加入keep_boxes
8. 如果keep_boxes不为空，则合并所有类别的框，然后返回。
但是，这样实现可能会比较慢。另一种方法是：先按类别分组，然后对每个组分别进行NMS，最后合并。
然而，在原始代码中，我们只进行了单类别的NMS（因为当时只有单类别）。现在需要支持多类别，所以必须按类别分组处理。
修改后的`non_max_suppression`函数（多类别版本）如下：
注意：由于原始代码中，这个函数是在类的方法中，所以我们需要修改为类的方法。
但是，由于代码较长，我们只给出关键部分。
然而，考虑到时间，我们也可以采用另一种方式：使用OpenCV的NMSBoxes，它支持多类别。但注意，OpenCV的NMSBoxes需要每个框的坐标和置信度，并且需要按类别提供。
我们可以这样做：
1. 将每个框的坐标转换为xyxy
2. 创建一个列表，存放所有框的坐标（xyxy）和对应的置信度、类别id
3. 调用cv2.dnn.NMSBoxes进行多类别NMS（注意：它要求传入所有框的坐标和每个框的得分，以及类别信息？）
但是，cv2.dnn.NMSBoxes只支持单类别NMS（即每个框只有一个分数），它不能直接处理多类别。所以，我们需要将每个框按照其最大类别分数作为该框的分数，然后使用NMSBoxes进行抑制（这样会抑制不同类别但重叠的框吗？不会，因为NMSBoxes在抑制时并不考虑类别，它只根据框的坐标和分数来抑制。所以，不同类别的框即使重叠也不会被抑制。因此，我们可以将所有框放在一起，然后调用一次NMSBoxes。
具体步骤：
1. 将每个框的坐标转换为xyxy（在原始图像坐标上）
2. 计算每个框的最大类别分数（已经在上面的步骤4中计算）
3. 将所有框（跨类别）的xyxy坐标、最大类别分数、类别id收集起来
4. 使用NMSBoxes进行非极大值抑制（指定score_threshold=conf_thres, nms_threshold=iou_thres）
这样，NMSBoxes会返回保留框的索引。
然后，我们根据这些索引，取出保留的框、分数和类别id。
因此，我们可以这样修改：
在`non_max_suppression`函数中：
... 步骤1-4（得到boxes, max_conf, class_id）...
# 步骤5：筛选
mask = max_conf > conf_thres
boxes = boxes[mask]
max_conf = max_conf[mask]
class_id = class_id[mask]
# 步骤6：将boxes从xywh转换到xyxy（注意：这里boxes是[N,4]格式，每行是[x,y,w,h]）
# 注意：原始代码中有一个xywh2xyxy函数，但它处理的是单个框？实际上，我们修改为可以处理多个框。
# 修改xywh2xyxy函数，使其支持批量处理（现在输入是[N,4]）
# 由于原始xywh2xyxy函数是针对一个框（一行）的，现在我们需要一个批量版本
# 定义一个新函数：xywh2xyxy_batch(batch_boxes)
# 或者，我们直接在函数内部转换：
# x1 = boxes[:,0] - boxes[:,2]/2
# y1 = boxes[:,1] - boxes[:,3]/2
# x2 = boxes[:,0] + boxes[:,2]/2
# y2 = boxes[:,1] + boxes[:,3]/2
# 但是注意：原始代码中xywh2xyxy函数考虑了预处理时的填充和缩放，所以我们需要用原始图像的尺寸进行转换。因此，我们仍然需要使用类中的xywh2xyxy函数，但需要修改为批量处理。
因此，我们修改类中的`xywh2xyxy`函数，使其支持批量处理。
原函数：
def xywh2xyxy(self, origin_h, origin_w, x):
# x: [center_x, center_y, w, h]
...
修改为：
def xywh2xyxy(self, origin_h, origin_w, x):
"""
x: 一个框或者多个框，形状为[N,4]（每行是center_x, center_y, w, h）
返回：转换后的xyxy坐标，形状[N,4]
"""
# 复制原始方法，但改为批量计算
y = np.zeros_like(x)
r_w = self.input_w / origin_w
r_h = self.input_h / origin_h
if r_h > r_w:
# 批量处理
y[:, 0] = x[:, 0] # center_x
y[:, 2] = x[:, 2] # w
y[:, 1] = x[:, 1] - (self.input_h - r_w * origin_h) / 2
y[:, 3] = x[:, 3] - (self.input_h - r_w * origin_h) / 2
y /= r_w
else:
y[:, 0] = x[:, 0] - (self.input_w - r_h * origin_w) / 2
y[:, 2] = x[:, 2] - (self.input_w - r_h * origin_w) / 2
y[:, 1] = x[:, 1]
y[:, 3] = x[:, 3]
y /= r_h
# 将center_x, center_y, w, h 转换为 x1,y1,x2,y2
# 注意：y现在每行是[center_x, center_y, w, h]（已经转换到原始图像尺寸）
# 然后，我们转换为xyxy
x1 = y[:, 0] - y[:, 2] / 2
y1 = y[:, 1] - y[:, 3] / 2
x2 = y[:, 0] + y[:, 2] / 2
y2 = y[:, 1] + y[:, 3] / 2
return np.stack([x1, y1, x2, y2], axis=1)
但是，这样修改会破坏原来的单框处理？所以，我们重新写一个批量处理的函数，或者修改原函数使其同时支持单个框和多个框。
为了简单，我们直接修改为批量处理。
注意：原函数中x的输入是[4]或[N,4]，输出对应。
修改后的`xywh2xyxy`函数：
def xywh2xyxy(self, origin_h, origin_w, x):
if len(x.shape) == 1:
x = x[None, :] # 变成[1,4]
y = np.zeros_like(x)
r_w = self.input_w / origin_w
r_h = self.input_h / origin_h
if r_h > r_w:
y[:, 0] = x[:, 0]
y[:, 2] = x[:, 2]
y[:, 1] = x[:, 1] - (self.input_h - r_w * origin_h) / 2
y[:, 3] = x[:, 3] - (self.input_h - r_w * origin_h) / 2
y /= r_w
else:
y[:, 0] = x[:, 0] - (self.input_w - r_h * origin_w) / 2
y[:, 2] = x[:, 2] - (self.input_w - r_h * origin_w) / 2
y[:, 1] = x[:, 1]
y[:, 3] = x[:, 3]
y /= r_h
# 转换为xyxy
xyxy_boxes = np.zeros_like(y)
xyxy_boxes[:, 0] = y[:, 0] - y[:, 2] / 2 # x1 = center_x - w/2
xyxy_boxes[:, 1] = y[:, 1] - y[:, 3] / 2 # y1
xyxy_boxes[:, 2] = y[:, 0] + y[:, 2] / 2 # x2
xyxy_boxes[:, 3] = y[:, 1] + y[:, 3] / 2 # y2
return xyxy_boxes
这样，我们就可以批量处理了。
回到`non_max_suppression`函数：
# 转换坐标
xyxy_boxes = self.xywh2xyxy(origin_h, origin_w, boxes) # [N,4] 每行是x1,y1,x2,y2
# 现在，xyxy_boxes是原始图像坐标下的边界框
# 然后，我们使用OpenCV的NMSBoxes
indices = cv2.dnn.NMSBoxes(
bboxes=xyxy_boxes.tolist(),
scores=max_conf.tolist(),
score_threshold=conf_thres,
nms_threshold=iou_thres,
eta=1.0,
top_k=200
)
# indices是一个一维数组，表示保留的框的索引
if len(indices) > 0:
indices = indices.flatten()
result_boxes = xyxy_boxes[indices]
result_scores = max_conf[indices]
result_classid = class_id[indices]
else:
result_boxes = np.array([])
result_scores = np.array([])
result_classid = np.array([])
# 合并结果，返回一个数组，每行是[x1,y1,x2,y2, conf, class_id]
if len(result_boxes) > 0:
result = np.concatenate([
result_boxes,
result_scores[:, None],
result_classid[:, None]
], axis=1)
else:
result = np.zeros((0,6))
return result
这样，我们就完成了多类别的NMS。
但是，注意：在原始代码中，`non_max_suppression`函数返回的是三个独立的数组（result_boxes, result_scores, result_classid）。现在，为了保持接口一致，我们仍然返回三个独立的数组（因为调用方是这么使用的）。
因此，我们可以这样返回：
return result_boxes, result_scores, result_classid
但是，注意：在调用`non_max_suppression`的地方（在`post_process`函数中）：
result_boxes, result_scores, result_classid = self.non_max_suppression(...)
所以，我们修改`non_max_suppression`函数，让它返回三个独立的数组（即使没有检测框，返回空数组）。
修改后的`non_max_suppression`函数（多类别版）如下：
def non_max_suppression(self, prediction, origin_h, origin_w, conf_thres=0.5, iou_thres=0.4):
# prediction: [N, num_values_per_detection] 其中N是检测框数量
# 我们只取前DET_NUM个值（检测部分）
dets = prediction[:, :DET_NUM] # [N, DET_NUM]
# 分离
boxes = dets[:, :4] # [N,4] (x,y,w,h)
obj_conf = dets[:, 4] # [N]
cls_scores = dets[:, 5:5+NUM_CLASSES] # [N, NUM_CLASSES]
# 计算类别置信度
class_conf = obj_conf[:, None] * cls_scores # [N, NUM_CLASSES]
# 找到每个框的最大类别置信度及其类别id
max_conf = np.max(class_conf, axis=1) # [N]
class_id = np.argmax(class_conf, axis=1) # [N]
# 初步筛选：保留大于置信度阈值的框
mask = max_conf > conf_thres
boxes = boxes[mask]
max_conf = max_conf[mask]
class_id = class_id[mask]
if len(boxes) == 0:
return np.array([]), np.array([]), np.array([])
# 转换坐标到原始图像（xywh to xyxy for batch）
xyxy_boxes = self.xywh2xyxy(origin_h, origin_w, boxes) # [N,4] (x1,y1,x2,y2)
# 使用OpenCV的NMSBoxes
# 注意：需要将坐标转换为整数？不需要，NMSBoxes支持浮点数
indices = cv2.dnn.NMSBoxes(
bboxes=xyxy_boxes.tolist(),
scores=max_conf.tolist(),
score_threshold=conf_thres, # 这个参数在OpenCV中实际上不起作用，因为我们已经筛选过了，所以设为0
nms_threshold=iou_thres,
eta=1.0,
top_k=200
)
if len(indices) > 0:
indices = indices.flatten()
result_boxes = xyxy_boxes[indices]
result_scores = max_conf[indices]
result_classid = class_id[indices]
else:
result_boxes = np.array([])
result_scores = np.array([])
result_classid = np.array([])
return result_boxes, result_scores, result_classid
注意：在调用NMSBoxes时，score_threshold参数可以设为0，因为我们已经在前面筛选过了。或者，我们可以不提前筛选，而是让NMSBoxes来筛选，但这样需要传入所有框（包括低置信度的），效率低。所以，我们提前筛选。
但是，在NMSBoxes中，score_threshold参数是用于筛选的，所以我们可以将score_threshold设为conf_thres，然后去掉我们自己的筛选步骤。这样，我们就不需要提前筛选了。
因此，我们可以不进行提前筛选，而是将所有框（包括低于阈值的）传入，然后让NMSBoxes根据score_threshold进行筛选。
修改：
# 去掉提前筛选的步骤，直接传入所有框（但注意：框的数量可能很多，所以效率问题？）
# 转换坐标
xyxy_boxes = self.xywh2xyxy(origin_h, origin_w, boxes) # 这里的boxes是原始的所有框（没有经过筛选）
# 然后调用NMSBoxes，设置score_threshold=conf_thres
indices = cv2.dnn.NMSBoxes(
bboxes=xyxy_boxes.tolist(),
scores=max_conf.tolist(), # 每个框的最大类别分数
score_threshold=conf_thres, # 这里会过滤掉低于阈值的框
nms_threshold=iou_thres
)
这样，我们就不需要提前筛选了。
但是，注意：max_conf和class_id是每个框的最大类别分数和类别id，即使这个最大分数低于阈值，我们也计算了。所以，我们可以先计算，然后传入NMSBoxes，它会根据score_threshold过滤。
因此，修改后的函数：
def non_max_suppression(self, prediction, origin_h, origin_w, conf_thres=0.5, iou_thres=0.4):
dets = prediction[:, :DET_NUM] # [N, DET_NUM]
boxes = dets[:, :4] # [N,4]
obj_conf = dets[:, 4] # [N]
cls_scores = dets[:, 5:5+NUM_CLASSES] # [N, NUM_CLASSES]
class_conf = obj_conf[:, None] * cls_scores # [N, NUM_CLASSES]
max_conf = np.max(class_conf, axis=1) # [N]
class_id = np.argmax(class_conf, axis=1) # [N]
# 转换坐标
xyxy_boxes = self.xywh2xyxy(origin_h, origin_w, boxes) # [N,4]
# 使用NMSBoxes
indices = cv2.dnn.NMSBoxes(
xyxy_boxes.tolist(),
max_conf.tolist(),
conf_thres,
iou_thres
)
if len(indices) > 0:
indices = indices.flatten()
result_boxes = xyxy_boxes[indices]
result_scores = max_conf[indices]
result_classid = class_id[indices]
else:
result_boxes = np.array([])
result_scores = np.array([])
result_classid = np.array([])
return result_boxes, result_scores, result_classid
注意：在调用NMSBoxes时，我们传入的bboxes是列表（每个元素是一个框的四个坐标），scores是每个框的分数（最大类别分数），score_threshold是conf_thres，nms_threshold是iou_thres。
这样，我们就完成了多类别的支持。
最后，我们需要修改几个全局变量：
NUM_CLASSES = 80 # 根据你的模型修改
DET_NUM = 5 + NUM_CLASSES # 4坐标+1个obj_conf+类别数
注意：SEG_NUM, POSE_NUM, OBB_NUM 可能需要根据模型调整。如果你的模型是纯检测模型（没有分割、姿态等），那么SEG_NUM和POSE_NUM应该为0。但是，在原始代码中，这些值被加在num_values_per_detection中。所以，如果你的模型输出中包含了这些额外的参数，那么num_values_per_detection = DET_NUM + SEG_NUM + POSE_NUM + OBB_NUM 应该等于模型输出中每个检测框的向量长度。
因此，请根据实际模型设置这些常量。
总结修改步骤：
1. 在代码开头，根据模型设置类别数：NUM_CLASSES = 80（根据你的模型修改）
2. 修改DET_NUM：DET_NUM = 5 + NUM_CLASSES
3. 修改xywh2xyxy函数，使其支持批量处理（如上述）
4. 修改non_max_suppression函数，使用上述的多类别处理方法（使用OpenCV的NMSBoxes）
5. 在post_process函数中，确保num_values_per_detection的计算正确（DET_NUM+SEG_NUM+...）
注意：如果你的模型没有分割和姿态等任务，那么SEG_NUM, POSE_NUM, OBB_NUM都应该设为0。
这样，代码就应该可以支持多类别检测了。
另外，在绘制框时，我们使用：
label="{}:{:.2f}".format(categories[int(result_classid[j])], result_scores[j])
其中，categories是类别名称列表（现在应该是多类别的列表）。所以，你需要将categories修改为包含所有类别名称的列表。
例如：
categories = ["class0", "class1", ...] # 你的类别名称
在代码的最后，我们初始化categories时，只定义了["drone"]，现在需要修改为所有类别。
最后，请注意：模型本身必须是多类别模型。如果模型是单类别训练的，那么即使我们修改了代码，也无法检测多类别。
因此，确保你使用的TensorRT引擎文件是多类别模型导出的。
修改后的代码结构如下（仅显示修改部分）：</think>要修改 YOLOv8 TensorRT 推理代码以支持多类别检测，需要调整以下几个关键部分。以下是详细的修改步骤和代码实现：
### 1. **修改类别定义和输出维度**
```python
# 修改类别数量（根据实际模型调整）
NUM_CLASSES = 80 # 例如 COCO 数据集有 80 个类别
DET_NUM = 5 + NUM_CLASSES # 4坐标 + 1置信度 + 类别数
POSE_NUM = 17 * 3
SEG_NUM = 32
OBB_NUM = 1
# 更新类别标签（示例使用 COCO 类别）
categories = ["person", "bicycle", "car", ... , "toothbrush"] # 完整 80 个类别
```
### 2. **重构后处理函数**
修改 `post_process` 和 `non_max_suppression` 函数以支持多类别：
```python
def post_process(self, output, origin_h, origin_w):
num_values_per_detection = DET_NUM + SEG_NUM + POSE_NUM + OBB_NUM
num = int(output[0])
pred = np.reshape(output[1:], (-1, num_values_per_detection))[:num, :]
# 多类别处理
boxes = self.multi_class_nms(pred, origin_h, origin_w,
conf_thres=CONF_THRESH,
iou_thres=IOU_THRESHOLD)
result_boxes = boxes[:, :4] if len(boxes) else np.array([])
result_scores = boxes[:, 4] if len(boxes) else np.array([])
result_classid = boxes[:, 5] if len(boxes) else np.array([])
return result_boxes, result_scores, result_classid
def multi_class_nms(self, prediction, origin_h, origin_w, conf_thres=0.5, iou_thres=0.4):
"""
多类别非极大值抑制处理
"""
# 提取检测部分 [x, y, w, h, obj_conf, class_conf_0, ..., class_conf_n]
dets = prediction[:, :DET_NUM]
boxes = dets[:, :4] # [N, 4]
obj_conf = dets[:, 4] # [N]
cls_conf = dets[:, 5:5+NUM_CLASSES] # [N, NUM_CLASSES]
# 计算类别置信度 = obj_conf * class_conf
class_conf = obj_conf[:, None] * cls_conf # [N, NUM_CLASSES]
# 找到每个框的最大类别分数和对应类别ID
max_conf = np.max(class_conf, axis=1) # [N]
class_id = np.argmax(class_conf, axis=1) # [N]
# 转换坐标到原始图像尺寸
xyxy_boxes = self.xywh2xyxy(origin_h, origin_w, boxes) # [N, 4]
# 使用 OpenCV 的多类别 NMS
nms_boxes = []
for cls_idx in range(NUM_CLASSES):
cls_mask = (class_id == cls_idx) & (max_conf > conf_thres)
if not np.any(cls_mask):
continue
cls_boxes = xyxy_boxes[cls_mask]
cls_scores = max_conf[cls_mask]
# 执行 NMS
indices = cv2.dnn.NMSBoxes(
cls_boxes.tolist(),
cls_scores.tolist(),
conf_thres,
iou_thres
)
if len(indices) > 0:
indices = indices.flatten()
for idx in indices:
box = cls_boxes[idx]
score = cls_scores[idx]
nms_boxes.append([*box, score, cls_idx])
return np.array(nms_boxes) if nms_boxes else np.zeros((0, 6))
```
### 3. **修改坐标转换函数**
```python
def xywh2xyxy(self, origin_h, origin_w, x):
"""
支持批量处理的坐标转换
x: [N, 4] 格式的数组，每行是 [center_x, center_y, w, h]
返回: [N, 4] 格式的 [x1, y1, x2, y2]
"""
if len(x.shape) == 1:
x = x[None, :] # 单个框转为批量处理
y = np.zeros_like(x)
r_w = self.input_w / origin_w
r_h = self.input_h / origin_h
if r_h > r_w:
y[:, 0] = x[:, 0]
y[:, 2] = x[:, 2]
y[:, 1] = x[:, 1] - (self.input_h - r_w * origin_h) / 2
y[:, 3] = x[:, 3] - (self.input_h - r_w * origin_h) / 2
y /= r_w
else:
y[:, 0] = x[:, 0] - (self.input_w - r_h * origin_w) / 2
y[:, 2] = x[:, 2] - (self.input_w - r_h * origin_w) / 2
y[:, 1] = x[:, 1]
y[:, 3] = y[:, 1] # 保持 y 坐标不变
y /= r_h
# 转换为 xyxy 格式
xyxy_boxes = np.zeros_like(y)
xyxy_boxes[:, 0] = y[:, 0] - y[:, 2] / 2 # x1
xyxy_boxes[:, 1] = y[:, 1] - y[:, 3] / 2 # y1
xyxy_boxes[:, 2] = y[:, 0] + y[:, 2] / 2 # x2
xyxy_boxes[:, 3] = y[:, 1] + y[:, 3] / 2 # y2
return xyxy_boxes
```
### 4. **确保模型输出匹配**
在导出模型时，需要确认输出层包含多类别信息：
```bash
# 导出 ONNX 模型时指定类别数
yolo mode=export model=yolov8n.pt format=onnx opset=12
```
### 5. **可视化多类别结果**
修改绘图函数以显示不同类别：
```python
# 在 infer 方法中
for j in range(len(result_boxes)):
box = result_boxes[j]
cls_id = int(result_classid[j])
label = f"{categories[cls_id]}:{result_scores[j]:.2f}"
# 为不同类别生成不同颜色
color = [int(c) for c in COLORS[cls_id % len(COLORS)]]
plot_one_box(box, batch_image_raw[i], color=color, label=label)
# 定义颜色常量
COLORS = [
(255, 0, 0), (0, 255, 0), (0, 0, 255),
(255, 255, 0), (255, 0, 255), (0, 255, 255)
]
```
### 关键修改说明：
1. **输出维度调整**：
- `DET_NUM = 5 + NUM_CLASSES` 包含坐标、置信度和类别分数
- 解决了单类别检测中固定 `DET_NUM=6` 的限制
2. **多类别 NMS**：
- 按类别分组处理，避免不同类别间的错误抑制
- 使用 `cv2.dnn.NMSBoxes` 替代原单类别实现
3. **批量坐标转换**：
- 修改 `xywh2xyxy` 支持批量处理
- 正确处理不同长宽比的图像填充
4. **模型一致性**：
- 确保导出的 ONNX/TensorRT 模型包含多类别输出
- 验证输出维度匹配 `(4 + 1 + NUM_CLASSES)`
> **注意**：如果模型包含分割（SEG_NUM）或姿态估计（POSE_NUM）任务，需要相应调整 `num_values_per_detection` 的计算[^2]。
### 性能优化建议：
1. **类别过滤优化**：
```python
# 提前过滤低置信度框，减少 NMS 计算量
conf_mask = max_conf > conf_thres * 0.5 # 宽松阈值
boxes = boxes[conf_mask]
cls_conf = cls_conf[conf_mask]
```
2. **并行处理**：
```python
# 使用多线程处理不同类别
from concurrent.futures import ThreadPoolExecutor
with ThreadPoolExecutor() as executor:
results = list(executor.map(process_class, range(NUM_CLASSES)))
```
3. **INT8 量化**：
```bash
# 导出时启用 INT8 量化
trtexec --onnx=yolov8n.onnx --saveEngine=yolov8n_int8.engine --int8
```
### 验证方法：
1. 检查输出张量形状：
```python
print(output.shape) # 应为 [batch, 4+1+NUM_CLASSES, num_boxes]
```
2. 测试不同类别的检测结果
3. 验证 mAP 指标是否与原始模型一致
> 修改后，模型可以同时检测多个类别，如 "car"、"person"、"drone" 等，并在图像中用不同颜色标注[^1][^5]。
本文详细介绍了非最大值抑制（NMS）的概念及其在目标检测中的应用，特别是使用TensorFlow实现这一过程的方法。NMS是一种用于减少重复边界框的技术，确保每个目标只被检测一次。
扫一扫
189
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
