yolov5口罩检测模型

1 项目背景及意义

全球范围内的公共卫生安全和人脸识别技术的发展。在面对新型冠状病毒等传染病的爆发和传播风险时,佩戴口罩成为一种重要的防护措施。然而,现有的人脸识别系统在识别戴口罩的人脸时存在一定的困难。

通过口罩识别技术,可以更准确地辨别佩戴口罩的人员身份,有助于提高公共安全水平,减少犯罪行为和恶意攻击的风险,且传统的人脸识别技术在遇到戴口罩的情况下准确度较低。通过开发口罩识别的异构高性能项目,可以提高人脸识别系统在戴口罩情况下的准确性和可靠性。在口罩成为日常生活必备品的背景下,通过口罩识别技术可以实现更便捷的公共服务,如自助购物、自助取票等,提升人们的生活体验。

这一技术的应用有助于应对公共卫生安全挑战,并推动人脸识别技术的进一步发展与应用。

2 环境搭建

2.1 yolov5模型下载

从github将yolov5的模型下载下来,github中yolov5首页如下图所示,通过右上角code下载全部代码。yolov5模型地址github地址
github
可通过release下载不同版本,本项目使用的是yolov5-5.0版本,可自由选择最新版本.
在这里插入图片描述

2.2 下载依赖

在已经有python环境的情况下,可以选用conda来构建虚拟环境,设置一个新的pytorch环境。

下载的yolov5模型根目录下有一个requirement.txt,其中包含了yolo所需的依赖包,通过指令pip install –r requirement.txt快速下载。如下图所示,图中的“already satisfied”是由于依赖已安装
在这里插入图片描述

3 数据集处理

实验中初始数据集为jpg格式图片与对应的txt格式标签文件,需要转为yolo格式的标签文件,以及划分数据集。代码部分主要参考目标检测数据集划分

3.1 txt转为xml

对于没有标签文件的图片可以使用labelImg标注,保存为xml文件,则不需要中txt转为xml,本次实验中没有xml标签,所以需要先转为xml。

文件结构如下图所示,在项目根目录下创建VOCdevkit/VOC2007,分为三个文件夹,JPEGImages用以存放图片,YOLO存放待转txt标签文件,Annotaions目前为空,用以存放转化后的XML文件。
在这里插入图片描述
在根目录下创建txttoxml.py,用以执行代码转txt为xml,详细代码如下

from xml.dom.minidom import Document
import os
import cv2
 
 
# def makexml(txtPath, xmlPath, picPath):  # txt所在文件夹路径,xml文件保存路径,图片所在文件夹路径
def makexml(picPath, txtPath, xmlPath):  # txt所在文件夹路径,xml文件保存路径,图片所在文件夹路径
    """此函数用于将yolo格式txt标注文件转换为voc格式xml标注文件
    在自己的标注图片文件夹下建三个子文件夹,分别命名为picture、txt、xml
    """
    dic = {'0': "face",  
           '1': "mask",
           }
    files = os.listdir(txtPath)
    for i, name in enumerate(files):
        xmlBuilder = Document()
        annotation = xmlBuilder.createElement("annotation")  # 创建annotation标签
        xmlBuilder.appendChild(annotation)
        txtFile = open(txtPath + name)
        txtList = txtFile.readlines()
        img = cv2.imread(picPath + name[0:-4] + ".jpg")
        Pheight, Pwidth, Pdepth = img.shape
 
        folder = xmlBuilder.createElement("folder")  # folder标签
        foldercontent = xmlBuilder.createTextNode("driving_annotation_dataset")
        folder.appendChild(foldercontent)
        annotation.appendChild(folder)  # folder标签结束
 
        filename = xmlBuilder.createElement("filename")  # filename标签
        filenamecontent = xmlBuilder.createTextNode(name[0:-4] + ".jpg")
        filename.appendChild(filenamecontent)
        annotation.appendChild(filename)  # filename标签结束
 
        size = xmlBuilder.createElement("size")  # size标签
        width = xmlBuilder.createElement("width")  # size子标签width
        widthcontent = xmlBuilder.createTextNode(str(Pwidth))
        width.appendChild(widthcontent)
        size.appendChild(width)  # size子标签width结束
 
        height = xmlBuilder.createElement("height")  # size子标签height
        heightcontent = xmlBuilder.createTextNode(str(Pheight))
        height.appendChild(heightcontent)
        size.appendChild(height)  # size子标签height结束
 
        depth = xmlBuilder.createElement("depth")  # size子标签depth
        depthcontent = xmlBuilder.createTextNode(str(Pdepth))
        depth.appendChild(depthcontent)
        size.appendChild(depth)  # size子标签depth结束
 
        annotation.appendChild(size)  # size标签结束
 
        for j in txtList:
            oneline = j.strip().split(" ")
            object = xmlBuilder.createElement("object")  # object 标签
            picname = xmlBuilder.createElement("name")  # name标签
            namecontent = xmlBuilder.createTextNode(dic[oneline[0]])
            picname.appendChild(namecontent)
            object.appendChild(picname)  # name标签结束
 
            pose = xmlBuilder.createElement("pose")  # pose标签
            posecontent = xmlBuilder.createTextNode("Unspecified")
            pose.appendChild(posecontent)
            object.appendChild(pose)  # pose标签结束
 
            truncated = xmlBuilder.createElement("truncated")  # truncated标签
            truncatedContent = xmlBuilder.createTextNode("0")
            truncated.appendChild(truncatedContent)
            object.appendChild(truncated)  # truncated标签结束
 
            difficult = xmlBuilder.createElement("difficult")  # difficult标签
            difficultcontent = xmlBuilder.createTextNode("0")
            difficult.appendChild(difficultcontent)
            object.appendChild(difficult)  # difficult标签结束
 
            bndbox = xmlBuilder.createElement("bndbox")  # bndbox标签
            xmin = xmlBuilder.createElement("xmin")  # xmin标签
            mathData = int(((float(oneline[1])) * Pwidth + 1) - (float(oneline[3])) * 0.5 * Pwidth)
            xminContent = xmlBuilder.createTextNode(str(mathData))
            xmin.appendChild(xminContent)
            bndbox.appendChild(xmin)  # xmin标签结束
 
            ymin = xmlBuilder.createElement("ymin")  # ymin标签
            mathData = int(((float(oneline[2])) * Pheight + 1) - (float(oneline[4])) * 0.5 * Pheight)
            yminContent = xmlBuilder.createTextNode(str(mathData))
            ymin.appendChild(yminContent)
            bndbox.appendChild(ymin)  # ymin标签结束
 
            xmax = xmlBuilder.createElement("xmax")  # xmax标签
            mathData = int(((float(oneline[1])) * Pwidth + 1) + (float(oneline[3])) * 0.5 * Pwidth)
            xmaxContent = xmlBuilder.createTextNode(str(mathData))
            xmax.appendChild(xmaxContent)
            bndbox.appendChild(xmax)  # xmax标签结束
 
            ymax = xmlBuilder.createElement("ymax")  # ymax标签
            mathData = int(((float(oneline[2])) * Pheight + 1) + (float(oneline[4])) * 0.5 * Pheight)
            ymaxContent = xmlBuilder.createTextNode(str(mathData))
            ymax.appendChild(ymaxContent)
            bndbox.appendChild(ymax)  # ymax标签结束
 
            object.appendChild(bndbox)  # bndbox标签结束
 
            annotation.appendChild(object)  # object标签结束
 
        f = open(xmlPath + name[0:-4] + ".xml", 'w')
        xmlBuilder.writexml(f, indent='\t', newl='\n', addindent='\t', encoding='utf-8')
        f.close()
 
if __name__ == "__main__":
    picPath = "VOCdevkit/VOC2007/JPEGImages/" 
    txtPath = "VOCdevkit/VOC2007/YOLO/" 
    xmlPath = "VOCdevkit/VOC2007/Annotations/"  
    makexml(picPath, txtPath, xmlPath)

3.2 xml转为yolo

yolo标签文件也是txt格式的,但要直接使用txt很可能会出错,所以需要先转xml再转yolo。

实验通过文件命名划分数据集,将test_00003到test_00004作为验证集,大约1000张图片。test_00004以后作为测试集,大约800张。余下作为训练集,大约6000张,如果需要修改这个,可用正则。

import xml.etree.ElementTree as ET
import pickle
import os
from os import listdir, getcwd
from os.path import join
import random
import re
from shutil import copyfile
 
classes = ["face", "mask"]
#classes=["ball"]
 
TRAIN_RATIO = 80
 
def clear_hidden_files(path):
    dir_list = os.listdir(path)
    for i in dir_list:
        abspath = os.path.join(os.path.abspath(path), i)
        if os.path.isfile(abspath):
            if i.startswith("._"):
                os.remove(abspath)
        else:
            clear_hidden_files(abspath)
 
def convert(size, box):
    dw = 1./size[0]
    dh = 1./size[1]
    x = (box[0] + box[1])/2.0
    y = (box[2] + box[3])/2.0
    w = box[1] - box[0]
    h = box[3] - box[2]
    x = x*dw
    w = w*dw
    y = y*dh
    h = h*dh
    return (x,y,w,h)
 
def convert_annotation(image_id):
    in_file = open('VOCdevkit/VOC2007/Annotations/%s.xml' %image_id)
    out_file = open('VOCdevkit/VOC2007/YOLO/%s.txt' %image_id, 'w')
    tree=ET.parse(in_file)
    root = tree.getroot()
    size = root.find('size')
    w = int(size.find('width').text)
    h = int(size.find('height').text)
 
    for obj in root.iter('object'):
        difficult = obj.find('difficult').text
        cls = obj.find('name').text
        if cls not in classes or int(difficult) == 1:
            continue
        cls_id = classes.index(cls)
        xmlbox = obj.find('bndbox')
        b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text), float(xmlbox.find('ymax').text))
        bb = convert((w,h), b)
        out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')
    in_file.close()
    out_file.close()
 
wd = os.getcwd()
wd = os.getcwd()
data_base_dir = os.path.join(wd, "VOCdevkit/")
if not os.path.isdir(data_base_dir):
    os.mkdir(data_base_dir)
work_sapce_dir = os.path.join(data_base_dir, "VOC2007/")
if not os.path.isdir(work_sapce_dir):
    os.mkdir(work_sapce_dir)
annotation_dir = os.path.join(work_sapce_dir, "Annotations/")
if not os.path.isdir(annotation_dir):
        os.mkdir(annotation_dir)
clear_hidden_files(annotation_dir)
image_dir = os.path.join(work_sapce_dir, "JPEGImages/")
if not os.path.isdir(image_dir):
        os.mkdir(image_dir)
clear_hidden_files(image_dir)
yolo_labels_dir = os.path.join(work_sapce_dir, "YOLO/")
if not os.path.isdir(yolo_labels_dir):
        os.mkdir(yolo_labels_dir)
clear_hidden_files(yolo_labels_dir)
yolov5_images_dir = os.path.join(data_base_dir, "images/")
if not os.path.isdir(yolov5_images_dir):
        os.mkdir(yolov5_images_dir)
clear_hidden_files(yolov5_images_dir)
yolov5_labels_dir = os.path.join(data_base_dir, "labels/")
if not os.path.isdir(yolov5_labels_dir):
        os.mkdir(yolov5_labels_dir)
clear_hidden_files(yolov5_labels_dir)
yolov5_images_train_dir = os.path.join(yolov5_images_dir, "train/")
if not os.path.isdir(yolov5_images_train_dir):
        os.mkdir(yolov5_images_train_dir)
clear_hidden_files(yolov5_images_train_dir)
yolov5_images_val_dir = os.path.join(yolov5_images_dir, "val/")
if not os.path.isdir(yolov5_images_val_dir):
        os.mkdir(yolov5_images_val_dir)
clear_hidden_files(yolov5_images_val_dir)
yolov5_images_test_dir = os.path.join(yolov5_images_dir, "test/")
if not os.path.isdir(yolov5_images_test_dir):
        os.mkdir(yolov5_images_test_dir)
clear_hidden_files(yolov5_images_test_dir)
yolov5_labels_train_dir = os.path.join(yolov5_labels_dir, "train/")
if not os.path.isdir(yolov5_labels_train_dir):
        os.mkdir(yolov5_labels_train_dir)
clear_hidden_files(yolov5_labels_train_dir)
yolov5_labels_val_dir = os.path.join(yolov5_labels_dir, "val/")
if not os.path.isdir(yolov5_labels_val_dir):
        os.mkdir(yolov5_labels_val_dir)
clear_hidden_files(yolov5_labels_val_dir)
yolov5_labels_test_dir = os.path.join(yolov5_labels_dir, "test/")
if not os.path.isdir(yolov5_labels_test_dir):
        os.mkdir(yolov5_labels_test_dir)
clear_hidden_files(yolov5_labels_test_dir)
 
train_file = open(os.path.join(wd, "yolov5_train.txt"), 'w')
val_file = open(os.path.join(wd, "yolov5_val.txt"), 'w')
test_file = open(os.path.join(wd, "yolov5_test.txt"), 'w')
train_file.close()
val_file.close()
test_file.close()
train_file = open(os.path.join(wd, "yolov5_train.txt"), 'a')
val_file = open(os.path.join(wd, "yolov5_val.txt"), 'a')
test_file = open(os.path.join(wd, "yolov5_test.txt"), 'a')
list_imgs = os.listdir(image_dir) # list image files
for i in range(0,len(list_imgs)):
    path = os.path.join(image_dir,list_imgs[i])
    if os.path.isfile(path):
        image_path = image_dir + list_imgs[i]
        voc_path = list_imgs[i]
        (nameWithoutExtention, extention) = os.path.splitext(os.path.basename(image_path))
        (voc_nameWithoutExtention, voc_extention) = os.path.splitext(os.path.basename(voc_path))
        annotation_name = nameWithoutExtention + '.xml'
        annotation_path = os.path.join(annotation_dir, annotation_name)
        label_name = nameWithoutExtention + '.txt'
        label_path = os.path.join(yolo_labels_dir, label_name)
    if(nameWithoutExtention.startswith("test_00003")): # test dataset
        if os.path.exists(annotation_path):
            val_file.write(image_path + '\n')
            convert_annotation(nameWithoutExtention) # convert label
            copyfile(image_path, yolov5_images_val_dir + voc_path)
            copyfile(label_path, yolov5_labels_val_dir + label_name)
    elif(nameWithoutExtention.startswith("test_00004")):
        if os.path.exists(annotation_path):
            test_file.write(image_path + '\n')
            convert_annotation(nameWithoutExtention) # convert label
            copyfile(image_path, yolov5_images_test_dir + voc_path)
            copyfile(label_path, yolov5_labels_test_dir + label_name)
             
    else: # train dataset
        if os.path.exists(annotation_path):
            train_file.write(image_path + '\n')
            convert_annotation(nameWithoutExtention) # convert label
            copyfile(image_path, yolov5_images_train_dir + voc_path)
            copyfile(label_path, yolov5_labels_train_dir + label_name)

train_file.close()
val_file.close()
test_file.close()

划分后文件结构如图所示,在VOCdevkit下生成了images与labels文件夹分别存放图片和yolo标签文件。
在这里插入图片描述
images下文件结构如图所示,test、train、val文件夹用以存放不同类型的数据集图片,labels文件夹同理

在这里插入图片描述

4 模型训练

4.1 下载yolov5预训练模型

同样在github的yolov5官网release页面可找到发布的预训练模型,如图所示。使用预训练模型能加快训练速度。实验中选择最基础的yolo5s.pt,精度更高的预训练模型需要更长的训练时间,实验中选择5s即可。
在这里插入图片描述
将下载后的模型放入weights目录下,如图所示
在这里插入图片描述

4.2 更改配置文件

4.2.1 修改模型配置文件

本着不破坏原文件的原则,将我们需要修改的models目录下的yolov5s.yaml复制,重命名为mask.yaml,修改其中nc为2,即改类型数量(number of classes)为2,如图所示
在这里插入图片描述

4.2.2 修改数据集配置文件

接着修改数据集配置文件,复制data下的voco.yaml,也重命名为mask.yaml,修改数据集对应位置,类型数量,类型名字,如图所示
在这里插入图片描述

4.2.3 修改训练代码参数

最后修改train.py中的参数加载修改后的配置,并设置需要的配置,那么训练的配置就已配置完毕

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--weights', type=str, default='weights/yolov5s.pt', help='initial weights path')
    parser.add_argument('--cfg', type=str, default='models/mask.yaml', help='model.yaml path')
    parser.add_argument('--data', type=str, default='data/mask.yaml', help='data.yaml path')
    parser.add_argument('--hyp', type=str, default='data/hyp.scratch.yaml', help='hyperparameters path')
    parser.add_argument('--epochs', type=int, default=200)
    parser.add_argument('--batch-size', type=int, default=8, help='total batch size for all GPUs')
    parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='[train, test] image sizes')
    parser.add_argument('--rect', action='store_true', help='rectangular training')
    parser.add_argument('--resume', nargs='?', const=True, default=False, help='resume most recent training')
    parser.add_argument('--nosave', action='store_true', help='only save final checkpoint')
    parser.add_argument('--notest', action='store_true', help='only test final epoch')
    parser.add_argument('--noautoanchor', action='store_true', help='disable autoanchor check')
    parser.add_argument('--evolve', action='store_true', help='evolve hyperparameters')
    parser.add_argument('--bucket', type=str, default='', help='gsutil bucket')
    parser.add_argument('--cache-images', action='store_true', help='cache images for faster training')
    parser.add_argument('--image-weights', action='store_true', help='use weighted image selection for training')
    parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
    parser.add_argument('--multi-scale', action='store_true', help='vary img-size +/- 50%%')
    parser.add_argument('--single-cls', action='store_true', help='train multi-class data as single-class')
    parser.add_argument('--adam', action='store_true', help='use torch.optim.Adam() optimizer')
    parser.add_argument('--sync-bn', action='store_true', help='use SyncBatchNorm, only available in DDP mode')
    parser.add_argument('--local_rank', type=int, default=-1, help='DDP parameter, do not modify')
    parser.add_argument('--workers', type=int, default=8, help='maximum number of dataloader workers')
    parser.add_argument('--project', default='runs/train', help='save to project/name')
    parser.add_argument('--entity', default=None, help='W&B entity')
    parser.add_argument('--name', default='exp', help='save to project/name')
    parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
    parser.add_argument('--quad', action='store_true', help='quad dataloader')
    parser.add_argument('--linear-lr', action='store_true', help='linear LR')
    parser.add_argument('--label-smoothing', type=float, default=0.0, help='Label smoothing epsilon')
    parser.add_argument('--upload_dataset', action='store_true', help='Upload dataset as W&B artifact table')
    parser.add_argument('--bbox_interval', type=int, default=-1, help='Set bounding-box image logging interval for W&B')
    parser.add_argument('--save_period', type=int, default=-1, help='Log model after every "save_period" epoch')
    parser.add_argument('--artifact_alias', type=str, default="latest", help='version of dataset artifact to be used')
    opt = parser.parse_args()

4.3 训练模型

配置好参数后运行train.py即可开始训练,由于yolo训练时间较长,6000张的训练集,且云端GPU资源有限,CPU与GPU交替训练,实验只训练了80次,但也断断续续花费了10多个小时,部分训练情况如图所示。
在这里插入图片描述
通过参数可及时查看训练情况

训练集

  • box:bounding box的平均损失值
  • obj: objectness的平均损失值
  • cls: 分类的平均损失值
  • total: 损失值的总和,即box+obj+cls
  • labels: 每个batch中标注物体数量的平均值
  • img_size: 图像大小

验证集

  • Images: 类别中图片数量
  • Labels: 类别中真实标注数量
  • P: 该类别的预测精准度
  • R: 找回率
  • mAP@.5: 平均精度均值,即在loU阈值为0.5时的平均精度
  • mAP@.5:.95: loU阈值从0.5到0.95之间的平均精度均值

这些指标的意义在于能够在训练过程中查看训练情况,及时调整训练参数。yolo支持中断训练的继续,并且经过通过实验表明中断训练只会在继续训练的前期对模型超参数造成影响,几次训练后将会恢复。

通过tensorboard可查看训练情况,100次训练已经足够,mAP参数与精准率和召回率都已经接近收敛,口罩检测模型训练完毕。
在这里插入图片描述
再通过测试集测试得到的结果如图13所示,用GPU进行训练,每张图片的推理时间只有3.7秒,F1平均分数在0.86864,F1最高分数为0.95556。图14为测试集中F1分数与置信度的函数变化,置信度在0.1到0.8之间时,F1分数变化比较稳定,说明没有问题
在这里插入图片描述
在这里插入图片描述

4.4 优化模型

4.4.1 模型转化

当直接使用训练得到的结果去CPU上运行时,每张图片的推理时间达到了69.5ms,是GPU的19倍左右,F1平均分数依旧为0.868,如图所示。
在这里插入图片描述
yolo直接在CPU上运行的效果不是很好,所以我选择使用oneapi提供的openvino去部署yolov5模型。此前需要将yolo的pt模型文件转化成openvino的IR模型。

通过执行export代码将pt模型导出为onnx模型。

修改export.py

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--weights', type=str, default='runs/train/exp7/weights/best.pt', help='weights path')  # from yolov5/models/
    parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='image size')  # height, width
    parser.add_argument('--batch-size', type=int, default=1, help='batch size')
    parser.add_argument('--dynamic', action='store_true', help='dynamic ONNX axes')
    parser.add_argument('--grid', action='store_true', help='export Detect() layer grid')
    parser.add_argument('--device', default='cpu', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
    opt = parser.parse_args()
    opt.img_size *= 2 if len(opt.img_size) == 1 else 1  # expand
    print(opt)
    set_logging()
    t = time.time()

通过命令行转化成onnx

python export.py --weights runs/weights/exp7/best.pt --include onnx

在下载openvino后,会提供转onnx为IR的工具mo_onnx.py,如图所示,通过mo_onnx.py转onnx为IR模型。
在这里插入图片描述
配置好openvino后通过命令转化

python mo_onnx.py --input_model 输入模型  --output_dir 输出地址

转化后weight目录下就有了onnx模型文件与IR模型
在这里插入图片描述

4.4.2 优化后模型检测

创建一个python文件进行图片检测。

import cv2
import numpy as np
import os
from openvino.inference_engine import IECore

def preprocess_image(image, input_shape):
    resized_image = cv2.resize(image, (input_shape[3], input_shape[2]))
    normalized_image = resized_image.astype(np.float32) / 255.0
    return np.transpose(normalized_image, [2, 0, 1])

def postprocess_output(output, image_shape, confidence_threshold=0.4):
    boxes, class_ids, confidences = [], [], []

    for detection in output[0][0]:
        score = detection[2][0]  # 调整以获取标量值
        if score > confidence_threshold:
            class_id = int(detection[1][0])  # 调整以获取标量值
            confidence = score
            box = detection[3:7]

            # 将边界框坐标缩放到原始图像大小
            h, w = image_shape[:2]
            box = [int(e[0] * min(w, h)) for e in box]

            boxes.append(box)
            class_ids.append(class_id)
            confidences.append(float(confidence))

    return boxes, class_ids, confidences

# COCO数据集标签
class_names = ["face", "mask"]

# OpenVINO IR模型文件
ir_filename = "./runs/train/exp7/weights/best.xml"

# 创建OpenVINO Core
ie = IECore()

# 读取网络并加载权重
net = ie.read_network(model=ir_filename, weights=ir_filename.replace(".xml", ".bin"))
exec_net = ie.load_network(network=net, device_name="CPU", num_requests=1)

# 获取输入和输出张量信息
input_image_tensor = next(iter(net.input_info))
input_shape = net.input_info[input_image_tensor].input_data.shape
output_image_tensor = next(iter(net.outputs))
output_shape = net.outputs[output_image_tensor].shape

# 图像文件夹路径
image_folder = "./VOCdevkit/images/test"
image_files = [f for f in os.listdir(image_folder) if f.endswith((".jpg", ".jpeg", ".png"))]
image_files = [os.path.join(image_folder, f) for f in image_files]

detections_list = []
f1_scores = []
inference_times = []

# 循环遍历每个图像文件
for image_file in image_files:
    # 读取图像文件
    image = cv2.imread(image_file)

    start = cv2.getTickCount()

    # 预处理图像并填充输入张量
    input_data = preprocess_image(image, input_shape)
    exec_net.requests[0].input_blobs[input_image_tensor].buffer[:] = input_data

    # 进行推断
    exec_net.requests[0].infer()

    # 获取推断结果
    output_tensor = exec_net.requests[0].output_blobs[output_image_tensor]
    det_output = output_tensor.buffer.reshape(output_shape)

    detection_time = (cv2.getTickCount() - start) / cv2.getTickFrequency()

    # 后处理输出
    boxes, class_ids, confidences = postprocess_output(det_output, image.shape)

    # 计算精确度、召回率和F1分数

    indexes = cv2.dnn.NMSBoxes(boxes, confidences, 0.25, 0.45)
    true_positives = len(indexes)
    false_positives = len(boxes) - true_positives
    false_negatives = 0  # 假设所有真实目标都被检测到

    # 计算精确度
    precision = true_positives / (true_positives + false_positives) if (true_positives + false_positives) != 0 else 0

    # 计算召回率
    recall = true_positives / (true_positives + false_negatives) if (true_positives + false_negatives) != 0 else 0

    # 避免除以零
    f1_score = 2 * (precision * recall) / (precision + recall) if (precision + recall) != 0 else 0

    # 计算检测时间

    detections_list.append((image_file, f1_score, detection_time))
    f1_scores.append(f1_score)
    inference_times.append(detection_time)

# 计算平均F1分数和平均推断时间
avg_f1_score = sum(f1_scores) / len(f1_scores)
avg_inference_time = sum(inference_times) / len(inference_times)

# 打印结果
print(f"average f1 score: {avg_f1_score:.5f}")
print(f"average inference time per images: {avg_inference_time*1000} ms")

# cv2.waitKey(0)
# cv2.destroyAllWindows()

最后编写oepnvino的测试图片代码进行测试,最后的测试结果如图18所示,平均F1分数有所下降,对整个测试结果影响不大。每张推理结果从69.5ms优化到了28.4ms每张,提升较大
在这里插入图片描述

5 系统测试

已训练模型会被放在runs/train/exp/weights文件夹中,其中best是最好的权重参数文件,选择best.pt作为最终的模型。

下载的yolov5模型中包含有测试模块与检测模块,测试模块将运行test.py文件对测试集进行测试以及结果分析。检测模块可以对待检测图片或视频进行检测,实验只测试图片检测功能。

在运行前需要配置相关参数,test.py修改模型文件路径与数据配置文件以及batch-size即可
在这里插入图片描述
配置detect.py参数,如图20所示,将weights参数改为实验中训练后模型,将data参数改为刚创建的detectImgs文件夹。

在这里插入图片描述
对于优化后模型,通过opencv绘制结果。

def postprocess_output(output, image_shape, confidence_threshold=0.4):
    boxes, class_ids, confidences = [], [], []

    for detection in output[0][0]:
        score = detection[2][0]  # 调整以获取标量值
        if score > confidence_threshold:
            class_id = int(detection[1][0])  # 调整以获取标量值
            confidence = score
            box = detection[3:7]

            # 将边界框坐标缩放到原始图像大小
            h, w = image_shape[:2]
            box = [int(e[0] * min(w, h)) for e in box]

            boxes.append(box)
            class_ids.append(class_id)
            confidences.append(float(confidence))

    # 在图像上绘制边界框和标签
    for i, box in enumerate(boxes):
        x1, y1, x2, y2 = box
        cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
        label = f"{class_names[class_ids[i]]}: {confidences[i]:.2f}"
        cv2.putText(image, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)

    return boxes, class_ids, confidences

v8

  • 23
    点赞
  • 23
    收藏
    觉得还不错? 一键收藏
  • 0
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值