YOLOv7 模型融合

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

首先需要两种模型，一个是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_load

def 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":
                continue
            class_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}" ]]; then
  echo "${outFrames} doesn't exist. Creating it.";
  mkdir -p ${outFrames}
fi

cd ${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/" ]]; then
  echo "./runs/detect/ doesn't exist. Creating it.";
  mkdir -p "./runs/detect/"
fi

cd ./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}" ]]; then
  echo "${merged_labels_path} doesn't exist. Creating it.";
  mkdir -p ${merged_labels_path}
fi


python merged_labels.py "${classes1}" "${classes2}" "${merged_labels_path}"