★★★ 本文源自AlStudio社区精品项目，【点击此处】查看更多精品内容 >>>

本项目使用VitDet系列模型，在铁路工人安全作业任务上取得了87.64%的mAP，达到了不错的检测效果，但不可避免地遇到了推理延时高的问题，有待改进。

一、项目背景

ViTDet是Meta AI团队（kaiming团队）在MAE之后提出的基于原生ViT模型作为骨干网络的检测模型。

最早的论文Benchmarking Detection Transfer Learning with Vision Transformers中，作者初步研究了以ViT作为骨干网络的检测模型所面临的挑战（架构的不兼容，训练速度慢以及显存占用大等问题），并给出了具体的解决方案，最重要的是发现基于MAE的预训练模型展现了较强的下游任务迁移能力，效果大大超过随机初始化和有监督预训练模型。

该模型设计 ROI heads的box head由原来的2个全连接层变为4个卷积层+1个全连接层； ROI heads的box head和mask head的卷积层之间均采用LayerNorm（最早的版本是采用BatchNorm，但往往需要SyncBN，而LN则不受batch size的影响）。

二、数据集简介

该数据集用于正确检测工人、他们的反光背心和安全帽。该数据集有3222张图片，其中包含三个标签：工人、反光背心和安全帽。AI Studio链接

用途举例：

  1. 可以判断是否有工人正在铁路上作业；
  2. 可以判断工人是否正确佩戴反光背心和安全帽规范作业。

部分数据集图片如下：

三、数据预处理

Step01： 解压数据集

!unzip /home/aistudio/data/data200816/dataset.zip -d /home/aistudio/work/

Step02： 将txt格式标注文件转换成xml格式标注文件，代码如下所示。

from xml.dom.minidom import Document
import os
import cv2
 
 
def makexml(picPath, txtPath, xmlPath):  
    # 标签映射
    dic = {'0': "vest", '1': "helmet", '2': "worker"}
    files = os.listdir(txtPath)
    for i, name in enumerate(files):
        if name == ".ipynb_checkpoints":
            continue
        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")      # .jpg/.png
        Pheight, Pwidth, Pdepth = img.shape
 
        folder = xmlBuilder.createElement("folder")  # folder标签
        foldercontent = xmlBuilder.createTextNode("datasetRGB")
        folder.appendChild(foldercontent)
        annotation.appendChild(folder) 
 
        filename = xmlBuilder.createElement("filename")  # filename标签
        filenamecontent = xmlBuilder.createTextNode(name[0:-4] + ".jpg")
        filename.appendChild(filenamecontent)
        annotation.appendChild(filename) 
 
        size = xmlBuilder.createElement("size")  # size标签
        width = xmlBuilder.createElement("width")  # size子标签width
        widthcontent = xmlBuilder.createTextNode(str(Pwidth))
        width.appendChild(widthcontent)
        size.appendChild(width) 
 
        height = xmlBuilder.createElement("height")  # size子标签height
        heightcontent = xmlBuilder.createTextNode(str(Pheight))
        height.appendChild(heightcontent)
        size.appendChild(height) 
 
        depth = xmlBuilder.createElement("depth")  # size子标签depth
        depthcontent = xmlBuilder.createTextNode(str(Pdepth))
        depth.appendChild(depthcontent)
        size.appendChild(depth) 
 
        annotation.appendChild(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) 
 
            pose = xmlBuilder.createElement("pose")  # pose标签
            posecontent = xmlBuilder.createTextNode("Unspecified")
            pose.appendChild(posecontent)
            object.appendChild(pose) 
 
            truncated = xmlBuilder.createElement("truncated")  # truncated标签
            truncatedContent = xmlBuilder.createTextNode("0")
            truncated.appendChild(truncatedContent)
            object.appendChild(truncated) 
 
            difficult = xmlBuilder.createElement("difficult")  # difficult标签
            difficultcontent = xmlBuilder.createTextNode("0")
            difficult.appendChild(difficultcontent)
            object.appendChild(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) 
 
            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) 
 
            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) 
 
            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)  
 
            object.appendChild(bndbox)  # bndbox标签结束
 
            annotation.appendChild(object) 
 
        f = open(xmlPath + name[0:-4] + ".xml", 'w')
        print(name)
        xmlBuilder.writexml(f, indent='\t', newl='\n', addindent='\t', encoding='utf-8')
        f.close()
 
 
if __name__ == "__main__":
    picPath = "/home/aistudio/work/dataset/JPEGImages/"  # 图片所在文件夹路径
    txtPath = "/home/aistudio/work/dataset/txt/"  # txt所在文件夹路径
    xmlPath = "/home/aistudio/work/dataset/Annotations/"  # xml文件保存路径
    makexml(picPath, txtPath, xmlPath)

Step03： 可视化转换结果，判断转换是否正确，代码如下所示。

  • 红色框（0，0，255）：代表worker；
  • 绿色框（0，255，0）：代表helmet；
  • 蓝色框（255，0，0）：代表vest。
import xml.etree.ElementTree as ET
import os
import cv2

src_XML_dir = '/home/aistudio/work/dataset/Annotations'  # xml源路径
src_IMG_dir = '/home/aistudio/work/dataset/JPEGImages'  # IMG原路径
IMG_format = '.jpg'    # IMG格式
out_dir = '/home/aistudio/work/output'  # 输出路径

if not os.path.exists(out_dir):
    os.makedirs(out_dir)
xml_file = os.listdir(src_XML_dir)  # 只返回文件名称,带后缀

for each_XML in xml_file:  # 遍历所有xml文件
    # 读入IMG
    if each_XML == ".ipynb_checkpoints":
        continue
    xml_FirstName = os.path.splitext(each_XML)[0]
    img_save_file = os.path.join(out_dir, xml_FirstName+IMG_format)
    img_src_path = os.path.join(src_IMG_dir, xml_FirstName+IMG_format)
    img = cv2.imread(img_src_path)
    # 解析XML
    each_XML_fullPath = src_XML_dir + '/' + each_XML  # 每个xml文件的完整路径
    tree = ET.parse(each_XML_fullPath)  # ET.parse()内要为完整相对路径
    root = tree.getroot()  # 类型为element

    # 画框
    for obj in root.findall('object'):
        if obj.find('bndbox'):
            if obj.find('name').text == 'worker':
                bndbox = obj.find('bndbox')
                xmin = int(bndbox.find('xmin').text)
                xmax = int(bndbox.find('xmax').text)
                ymin = int(bndbox.find('ymin').text)
                ymax = int(bndbox.find('ymax').text)

                cv2.rectangle(img=img,
                            pt1=(xmin,ymin),
                            pt2=(xmax,ymax),
                            color=(0,0,255),
                            thickness=2)

            if obj.find('name').text == 'helmet':
                bndbox = obj.find('bndbox')
                xmin = int(bndbox.find('xmin').text)
                xmax = int(bndbox.find('xmax').text)
                ymin = int(bndbox.find('ymin').text)
                ymax = int(bndbox.find('ymax').text)

                cv2.rectangle(img=img,
                            pt1=(xmin,ymin),
                            pt2=(xmax,ymax),
                            color=(0,255,0),
                            thickness=2)

            if obj.find('name').text == 'vest':
                bndbox = obj.find('bndbox')
                xmin = int(bndbox.find('xmin').text)
                xmax = int(bndbox.find('xmax').text)
                ymin = int(bndbox.find('ymin').text)
                ymax = int(bndbox.find('ymax').text)

                cv2.rectangle(img=img,
                            pt1=(xmin,ymin),
                            pt2=(xmax,ymax),
                            color=(255,0,0),
                            thickness=2)


    cv2.imwrite(filename=img_save_file, img=img)
    print('保存结果{}'.format(xml_FirstName))

部分可视化结果如下：

四、代码实现

4.1 检测数据分析

该数据集总共包含3个标签，各类标签的数量分别为：

  • worker: 7973
  • vest: 7883
  • helmet: 6515
import os
from unicodedata import name
import xml.etree.ElementTree as ET
import glob

def count_num(indir):
    # 提取xml文件列表
    os.chdir(indir)
    annotations = os.listdir('.')
    annotations = glob.glob(str(annotations) + '*.xml')

    dict = {} # 新建字典，用于存放各类标签名及其对应的数目
    for i, file in enumerate(annotations): # 遍历xml文件
       
        # actual parsing
        in_file = open(file, encoding = 'utf-8')
        tree = ET.parse(in_file)
        root = tree.getroot()

        # 遍历文件的所有标签
        for obj in root.iter('object'):
            name = obj.find('name').text
            if(name in dict.keys()): dict[name] += 1 # 如果标签不是第一次出现，则+1
            else: dict[name] = 1 # 如果标签是第一次出现，则将该标签名对应的value初始化为1

    # 打印结果
    print("各类标签的数量分别为：")
    for key in dict.keys(): 
        print(key + ': ' + str(dict[key]))            

indir='/home/aistudio/work/dataset/Annotations/'   # xml文件所在的目录
count_num(indir) # 调用函数统计各类标签数目

图像尺寸分析： 通过图像尺寸分析，我们可以看到该数据集图片的尺寸，均为[1920, 1080]。

import os
from unicodedata import name
import xml.etree.ElementTree as ET
import glob

def Image_size(indir):
    # 提取xml文件列表
    os.chdir(indir)
    annotations = os.listdir('.')
    annotations = glob.glob(str(annotations) + '*.xml')
    width_heights = []

    for i, file in enumerate(annotations): # 遍历xml文件
        # actual parsing
        in_file = open(file, encoding = 'utf-8')
        tree = ET.parse(in_file)
        root = tree.getroot()
        width = int(root.find('size').find('width').text)
        height = int(root.find('size').find('height').text)
        if [width, height] not in width_heights: width_heights.append([width, height])
    print("数据集中，有{}种不同的尺寸，分别是：".format(len(width_heights)))
    for item in width_heights:
        print(item)

indir='/home/aistudio/work/dataset/Annotations/'   # xml文件所在的目录
Image_size(indir)

4.2 安装PaddleDetection

# 克隆PaddleDetection仓库
#!git clone https://github.com/PaddlePaddle/PaddleDetection.git

# 安装其他依赖
%cd /home/aistudio/PaddleDetection/
!pip install -r requirements.txt --user

# 编译安装paddledet
!python setup.py install

4.3 数据集划分

首先安装PaddleX。

!pip install paddlex

然后，我们通过paddlex中的split_dataset命令按照0.9：0.1的比例划分训练集和验证集。

!paddlex --split_dataset --format VOC --dataset_dir /home/aistudio/work/dataset --val_value 0.1

4.4 模型选择

本项目选择以ViT作为backbone的PP-YOLOE模型，而原生的PP-YOLOE以CSPResNet作为backbone。ViT作为backbone具有以下几个优点：

  • 高效性：Vit使用自注意力机制，而不是卷积操作，使得计算效率更高，训练速度更快。
  • 通用性：Vit不受输入图片大小的限制，因此可以处理任意尺寸的图片，而且同样适用于各种类型的图像分类任务。
  • 性能：在许多视觉任务中，如ImageNet上的图像分类任务，Vit的性能已经相当接近传统的基于卷积神经网络的模型。
  • 精度：由于Vit的自注意力机制结合了整个图像中的信息，因此在某些情况下，它能够产生更高的准确性。例如，在处理存在遮挡者的图像时，它能够识别到不同区域的物体。
  • 可解释性：Vit中的注意力机制还能够为每个位置生成注意力分数，从而提高可解释性，使得人们更容易理解模型是如何做出预测的。

以下是以ViT作为backbone的Mask-RCNN模型的结构图：

4.5 模型训练

Step01： 将数据集移动到/home/aistudio/PaddleDetection/dataset目录下。

!mv /home/aistudio/work/dataset /home/aistudio/PaddleDetection/dataset/

Step02： 单卡训练

!python tools/train.py -c configs/vitdet/ppyoloe_vit_base_csppan_cae_36e_coco.yml --eval --use_vdl True --vdl_log_dir vdl_log_dir/scalar

损失函数如图所示：

4.6 模型评估

通过如下命令在单个GPU上评估我们的验证集。

!python tools/eval.py -c configs/vitdet/ppyoloe_vit_base_csppan_cae_36e_coco.yml -o weights=output/ppyoloe_vit_base_csppan_cae_36e_coco/best_model.pdparams

指标如下：

  • mAP(0.50, 11point) = 87.64%
  • Total sample number: 309, average FPS: 13.920338377531873**

4.7 模型推理

我们可以通过以下命令在单张GPU上推理文件中的所有图片。

!python tools/infer.py -c configs/vitdet/ppyoloe_vit_base_csppan_cae_36e_coco.yml -o weights=output/ppyoloe_vit_base_csppan_cae_36e_coco/best_model.pdparams --infer_dir=dataset/dataset/JPEGImages --output_dir infer_output/

部分可视化结果如下：

4.8 模型导出

PP-YOLOE在GPU上部署或者速度测试需要通过tools/export_model.py导出模型。

!python tools/export_model.py -c configs/vitdet/ppyoloe_vit_base_csppan_cae_36e_coco.yml -o weights=output/ppyoloe_vit_base_csppan_cae_36e_coco/best_model.pdparams

4.9 FastDeploy快速部署

环境准备： 本项目的部署环节主要用到的套件为飞桨部署工具FastDeploy，因此我们先安装FastDeploy。

!pip install fastdeploy-gpu-python -f https://www.paddlepaddle.org.cn/whl/fastdeploy.html --user

部署模型：

导入飞桨部署工具FastDepoy包，创建Runtimeoption，具体实现如下代码所示。

import fastdeploy as fd
import cv2
import os

def build_option(device='cpu', use_trt=False):
    option = fd.RuntimeOption()

    if device.lower() == "gpu":
        option.use_gpu()

    if use_trt:
        option.use_trt_backend()
        option.set_trt_input_shape("image", [1, 3, 640, 640])
        option.set_trt_input_shape("scale_factor", [1, 2])

    return option

配置模型路径，创建Runtimeoption，指定部署设备和后端推理引擎，代码实现如下所示。

# 配置模型路径
model_path = '/home/aistudio/PaddleDetection/output_inference/ppyoloe_vit_base_csppan_cae_36e_coco'
image_path = '/home/aistudio/PaddleDetection/dataset/dataset/JPEGImages/00000350.jpg'
model_file = os.path.join(model_path, "model.pdmodel")
params_file = os.path.join(model_path, "model.pdiparams")
config_file = os.path.join(model_path, "infer_cfg.yml")

# 创建RuntimeOption
runtime_option = build_option(device='cpu', use_trt=False)

# 创建PPYOLOE模型
model = fd.vision.detection.PPYOLO(model_file,
                                   params_file,
                                   config_file,
                                   runtime_option=runtime_option)

# 预测图片检测结果
im = cv2.imread(image_path)
result = model.predict(im.copy())
print(result)

# 预测结果可视化
vis_im = fd.vision.vis_detection(im, result, score_threshold=0.5)
cv2.imwrite("/home/aistudio/work/visualized_result.jpg", vis_im)
print("Visualized result save in ./visualized_result.jpg")

可视化结果如下：

五、总结提高

在铁路安全作业任务中，我使用了ViTDet系列模型，取得了不错了的效果。但还是避不开一个缺点就是推理速度太慢。

后续我们可以做以下改进：

  1. 使用轻量化模型重新训练，在不降低模型精度的情况下，实现更低的推理延时。
  2. 采用剪枝等模型压缩的方式减小模型规模。
  3. 反光背心和安全帽基本都是在工人的标注框中间，我们可以调整模型结构，集中关注工人标注框中间的信息。

此文章为转载
原文链接

# 人工智能 # 深度学习 # 机器学习
Logo
百度飞桨AI Studio社区

学大模型，用大模型上飞桨星河社区！每天8点V100G算力免费领！免费领取ERNIE 4.0 100w Token >>>

更多推荐

cover

Python调用文心一言的API

avatar 百度飞桨星河社区
cover

文心一言API调用（python）

avatar 百度飞桨星河社区

[PaddleSpeech 原神] 音色克隆之胡桃

[PaddleSpeech] 音色克隆之原神角色 ，项目代码已完整放在aistudio,地址如下：快去体验吧，可以直接运行。

avatar 百度飞桨星河社区