YOLOv7 模型融合

news/2024/7/11 0:11:09 标签: YOLO, ffmpeg

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

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

http://www.niftyadmin.cn/n/420528.html

相关文章

【Docker】docker部署前Springboot-vue后端分离项目【实战篇】

【Docker】docker部署前Springboot-vue后端分离项目【实战篇】

文章目录 1、docker 安装jdk2、docker 安装mysql通过Docker命令进入Mysql容器内部初始化数据sqlDbx连接查看 3、docker build构建后端镜像修改配置数据库JDBC链接IP为虚拟机服务器IPmaven clean package打包后端jar并上传到服务器编写Dockfilebuild 构建查看构建的后端镜像app …
阅读更多...
【ABAP】数据类型（二）「预定义数据类型」

【ABAP】数据类型（二）「预定义数据类型」

💂作者简介： THUNDER王，一名热爱财税和SAP ABAP编程以及热爱分享的博主。目前于江西师范大学本科在读，同时任汉硕云（广东）科技有限公司ABAP开发顾问。在学习工作中，我通常使用偏后端的开发语言A…
阅读更多...
Java反射与注解

Java反射与注解

文章目录 一、 注解1.简介2. 元注解3. 自定义注解 二、 反射1. 简介2. 理解Class类并获取Class实例3. 类的加载与初始化4. 类加载器ClassLoader5. 获取运行时类的完整结构6. 动态创建对象执行方法7. 反射操作泛型8. 反射操作注解 一、 注解 1.简介 Annotation是JDK5.0开始引入…
阅读更多...
问麻了…阿里一面索命27问，过了就60W+

问麻了…阿里一面索命27问，过了就60W+

前言 在40岁老架构师尼恩的（50）读者社区中，经常有小伙伴，需要面试阿里、 百度、头条、美团、京东等大厂。 下面是一个小伙伴成功拿到通过了阿里三次技术面试，小伙伴通过三个多小时技术拷问，最终拿到 offe…
阅读更多...
轻松解密PDF，解除你的文档

轻松解密PDF，解除你的文档

最简单的办法去除PDF密码、找回PDF密码、解密PDF，不用下载软件，只需在百度搜索“密码帝官网”并点击“立即开始”，在用户中心上传需要解密的文件，稍等片刻就能找回密码或解除限制。即使PDF文档不能编辑、打印、复制，也…
阅读更多...
Django、数据库----------ORM框架

Django、数据库----------ORM框架

目录 1.安装第三方模块 2.ORM 1.自己创建数据库 2.django连接数据库 3.django操作表 创建表：在models.py文件中 删除表以及修改表 在表中新增列时，由于已存在列中可能已有数据，所以新增列必须要指定新增列对应的数据 Django开发操作数…
阅读更多...
让AI写今年高考作文，满分60你给AI打多少分呢？

让AI写今年高考作文，满分60你给AI打多少分呢？

随着高考的落幕，各地高考作文已成为我们讨论的话题。 下面是2023年全国甲卷作文题目 我们就让AI根据要求，写上两篇作文。作文标题也是AI起的 一 、① 标题：与时间赛跑，不为时间所困 ​ 时间，这位无情的审判长&#…
阅读更多...
【沐风老师】详解3DMAX一键破损插件PolyDamage使用方法

【沐风老师】详解3DMAX一键破损插件PolyDamage使用方法

3DMAX一键破损插件教程 3dMax一键破损插件PolyDamage是一种快速添加模型损坏和缺陷的工具。现实生活中没有什么是完美的，所以给你的模型增加一些破损会看起来更逼真。PolyDamage是使用其他软件包手动雕刻损伤的快速替代方案。PolyDamage为混凝土墙、地面、石头和柱子增添了…
阅读更多...
最新文章

Copyright @ 2022~2023