将两种训练模型的检测结果融合

首先需要两种模型，一个是yolov7原有的模型，另一个是你训练出来的模型

其中yolov7.pt是官方提供的模型，有80个类别
其中yolov7_3.7HRW.pt是我们自己训练的模型，有三个分类，举手、看书、写字

开始前需要确保ffmpeg是否安装

conda install x264 ffmpeg -c conda-forge -y

还需要建立一个python脚本
get_classes.py
写入如下内容：

import torch
from pathlib import Path
from models.experimental import attempt_loaddef get_classes(model_file_path):# 加载模型model = attempt_load(model_file_path, map_location='cpu')# 获取分类标签classes = model.module.names if hasattr(model, 'module') else model.names# 返回分类标签return classes

然后创建：
merged_labels.py

import sys
from numpy import random
import os# 获取第一个参数的值
class1 = sys.argv[1].split("[")[-1].split("]")[0].split(", ")# 获取第二个参数的值
class2 = sys.argv[2].split("[")[-1].split("]")[0].split(", ")# 获取第三个参数的值
merged_labels_path = sys.argv[3]merge_class = class1+class2# 随机生成颜色
colors = [[random.randint(0, 255) for _ in range(3)] for _ in merge_class]path1 = './runs/detect/exp/labels/'
path2 = './runs/detect/exp2/labels/'for file_name in os.listdir(path1):file_path1 = os.path.join(path1, file_name)merged_file_path = os.path.join(merged_labels_path, file_name)with open(file_path1) as f1, open(merged_file_path, 'a') as fw:lines1 = f1.readlines()for line in lines1:fw.write(line)

import os
import argparse
import cv2# 设置不同类别的颜色和标签
class_colors = [(255, 0, 0), (0, 255, 0), (0, 0, 255), (200, 0, 0), (0, 200, 0), (0, 0, 200),(150, 0, 0), (0, 150, 0), (0, 0, 150), (100, 0, 0), (0, 100, 0)] # 对应11个类别的颜色
class_label1 = ['ball', 'sandpile', 'ruler', 'weight', 'pit', 'weigh', 'height measure', 'drop ball', 'size measure', 'record'] # 对应10个类别的标签
class_label2 = ['person']
class_labels = class_label1 + class_label2
# 定义parse对象，用于输入命令行参数
parser = argparse.ArgumentParser(description='Merge the detection results of two YOLOv7 models and draw bounding boxes on images.')
parser.add_argument('path1', type=str, help='the path of the detection results from the first YOLOv7 model')
parser.add_argument('path2', type=str, help='the path of the detection results from the second YOLOv7 model')
parser.add_argument('image_path', type=str, help='the path of the images to be processed')
args = parser.parse_args()# 分别获取两个检测结果文件的路径，并将它们的内容合并
path1 = os.path.join(args.path1, 'labels')
path2 = os.path.join(args.path2, 'labels')
merged_labels_path = os.path.join(args.path1, 'merged_labels')os.makedirs(merged_labels_path, exist_ok=True)for file_name in os.listdir(path1):file_path1 = os.path.join(path1, file_name)merged_file_path = os.path.join(merged_labels_path, file_name)with open(file_path1) as f1, open(merged_file_path, 'a') as fw:lines1 = f1.readlines()for line in lines1:fw.write(line)for file_name in os.listdir(path2):file_path2 = os.path.join(path2, file_name)merged_file_path = os.path.join(merged_labels_path, file_name)with open(file_path2) as f2, open(merged_file_path, 'a') as fw:lines2 = f2.readlines()for line in lines2:class_id, x_center, y_center, w, h = line.split()if class_id != "0":continueclass_id = str(len(class_label1) + int(class_id))line = class_id + ' ' + x_center + ' ' + y_center + ' ' + w + ' ' + h + '\n'fw.write(line)# 处理每张图像，画上检测框
image_path = args.image_path
output_path = 'output'os.makedirs(output_path, exist_ok=True)for file_name in os.listdir(image_path):image_file_path = os.path.join(image_path, file_name)label_file_path = os.path.join(merged_labels_path, file_name[:-4] + '.txt')# label_file_path = os.path.join(merged_labels_path, file_name.split("_")[0] + "_" + str(int(file_name.split("_")[1].split('.')[0])) + '.txt')assert os.path.exists(label_file_path), f"Label file not found: {label_file_path}"img = cv2.imread(image_file_path)with open(label_file_path) as f:for line in f:# 使用yolo格式标签在图片上画矩阵class_id, x_center, y_center, w, h = line.split()x_center, y_center, w, h = float(x_center), float(y_center), float(w), float(h)x_min = int((x_center - w / 2) * img.shape[1])y_min = int((y_center - h / 2) * img.shape[0])x_max = int((x_center + w / 2) * img.shape[1])y_max = int((y_center + h / 2) * img.shape[0])class_id = int(class_id)# 取一个颜色和标签并应用在当前的类别上color = class_colors[class_id % len(class_colors)]label = class_labels[class_id % len(class_labels)]cv2.rectangle(img, (x_min, y_min), (x_max, y_max), color, 2)# 在矩形框上显示类别标签cv2.putText(img, label, (x_min, y_min), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)output_file_path = os.path.join(output_path, file_name)cv2.imwrite(output_file_path, img)print('Done.')

首先是一个bash脚本，用来控制所有程序运行。

bash fusionDetection.sh 待检测的视频路径 每秒裁剪帧数 输出视频帧的路径 第一种模型权重的路径 第二种权重的路径 合并的检测结果（txt）路径
如：
bash fusionDetection.sh ./inVideo/1.mp4 1 ./outFrames ./yolov7.pt ./yolov7_4.2k_HRW.pt ./merged_labels

inVideo=$1
frameRate=$2
outFrames=$3
weight1=$4
weight2=$5# 先对输出视频帧的路径进行清空处理
CRTDIR=$(pwd)if [[ ! -d "${outFrames}" ]]; thenecho "${outFrames} doesn't exist. Creating it.";mkdir -p ${outFrames}
ficd ${outFrames}
rm -rf *
cd ${CRTDIR}# 对视频进行裁剪
out_name="${outFrames}/out_%06d.jpg"
ffmpeg -i "${inVideo}" -r ${frameRate} -q:v 1 "${out_name}"# 进入 yolov7 的目录# 先清空之前的检测结果
if [[ ! -d "./runs/detect/" ]]; thenecho "./runs/detect/ doesn't exist. Creating it.";mkdir -p "./runs/detect/"
ficd ./runs/detect/
rm -rf *
cd ../../python detect.py --weights ${weight1} --conf 0.25 --img-size 640 --source ${outFrames} --save-txt
python detect.py --weights ${weight2} --conf 0.25 --img-size 640 --source ${outFrames} --save-txt# 读取权重种的类别# 调用 Python 脚本获取分类标签
classes1=$(python -c "from get_classes import get_classes; print(get_classes('${weight1}'))")classes2=$(python -c "from get_classes import get_classes; print(get_classes('${weight2}'))")# 打印分类标签
echo "Classes: ${classes1}"
echo "Classes: ${classes2}"# 融合的xt的位置
merged_labels_path = './mergedLabels/'if [[ ! -d "${merged_labels_path}" ]]; thenecho "${merged_labels_path} doesn't exist. Creating it.";mkdir -p ${merged_labels_path}
fipython merged_labels.py "${classes1}" "${classes2}" "${merged_labels_path}"