使用飞腾派PRO开始学习机器视觉: 教飞腾派PRO检测安全帽

1. 目标检测

目标检测（Object Detection）计算视觉中一项常见的任务，其具体是指将检测图像中目标目标，并用带标签的边界框框出来。（如下图所示）

YOLO系列算法是目前最流行的目标检测算法，其检测速度快，精度高，在医疗，工业等领域被广泛使用。YOLO的核心思想在于对整张图像分割为数个格子（grid）,预测每个网格的类概率和边界框（bounding box）,将目标检测问题转化为回归问题。由于YOLO的实现细节较为复杂，而本文目的在于介绍飞腾派PRO的实际部署，因此算法细节本文不一一赘述。感兴趣的各位可以参考YOLO论文1进行更深入的了解。在此本篇文章将使用YOLO v8n作为安全帽识别的检测算法。

2. 环境准备

在PC端安装Ultralytics至少需要满足以下要求：

Python 3.8-3.11

Ubuntu 16.04 or later / Windows 10

NVIDIA GPU Driver (如果需要GPU)

Ultralytics使用PyTorch作为其框架

2.1 安装PyTorch

1.创建conda环境

使用Win + x，唤起开始菜单，选择运行，键入pwsh，进入powershell。

进入miniconda安装位置，激活conda环境

$ cd D:\miniconda\shell\condabin

$ .\conda-hook.ps1

$ conda activate base

2.创建torch环境。

(base) PS C:\Users> conda create --name torch python=3.10

(base) PS C:\Users> conda activate torch

(torch) PS C:\Users>

3.安装PyTorch并验证

(torch) PS C:\Users> conda install pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch -c nvidia

(torch) PS C:\Users> python -c "import torch; print(torch.cuda.is_available())"

True

2.2 安装Ultralytics

(torch) PS C:\Users> python -m pip install ultralytics

3. 代码时间

3.1 转换数据

本文使用百度AI Studio提供的安全帽数据集，此数据集拥有5000张带标注的安装帽图片，拥有安全帽,头两种标注。

由于此数据集使用的VOC格式，为了Ultralytics可以正常训练，还需要将其转换为coco格式

安全帽数据集可在Baidu AI Studio注册账号，搜索下载

转换代码如下：

import os

import re

from tqdm.auto import tqdm

from shutil import copy

import yaml

import xml.etree.ElementTree as ET

def convert_box(size, box):

"""

convert into (x, y, w, h) from (xmin, ymin, xmax, ymax)

"""

dw, dh = 1. / size[0], 1. / size[1]

x, y, w, h = (box[0] + box[1]) / 2.0 - 1, (box[2] + box[3]) / 2.0 - 1, box[1] - box[0], box[3] - box[2]

return x * dw, y * dh, w * dw, h * dh

def convert_label(path, lb_path, year, image_id):

in_file = open(path / f'VOC{year}/Annotations/{image_id}.xml')

out_file = open(lb_path, '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)

names = list(yaml['names'].values()) # names list

for obj in root.iter('object'):

cls = obj.find('name').text

if cls in names and int(obj.find('difficult').text) != 1:

xmlbox = obj.find('bndbox')

bb = convert_box((w, h), [float(xmlbox.find(x).text) for x in ('xmin', 'xmax', 'ymin', 'ymax')])

cls_id = names.index(cls) # class id

out_file.write(" ".join([str(a) for a in (cls_id, *bb)]) + '\n')

# Path to images and annotations

path_images = "./HelmetDetection/images/"

path_annot = "./HelmetDetection/annotations/"

target_dir = "./dataset"

# Get all XML file paths in path_annot and sort them

xml_files = sorted(

[

os.path.join(path_annot, file_name)

for file_name in os.listdir(path_annot)

if file_name.endswith(".xml")

]

)

# Get all JPEG image file paths in path_images and sort them

jpg_files = sorted(

[

os.path.join(path_images, file_name)

for file_name in os.listdir(path_images)

if file_name.endswith(".jpg")

]

)

# split dataset

ratio = 0.8

N = len(xml_files)

train_set = xml_files[:int(N * ratio)]

val_set = xml_files[int(N*ratio):]

dataset = {

"train": train_set,

"val": val_set

}

# convert coco

class_ids = [

"helmet",

"head",

"person",

]

class_mapping = dict(zip(range(len(class_ids)), class_ids))

for k, v in dataset.items():

for xml_file in tqdm(v):

tree = ET.parse(xml_file)

root = tree.getroot()

image_names = root.find("filename").text

# write coco label

label_txt = re.sub("png", "txt", image_names)

out_file = open(f"{target_dir}/labels/{k}/{label_txt}", 'w+')

size = root.find('size')

w = int(size.find('width').text)

h = int(size.find('height').text)

for obj in root.iter("object"):

cls = obj.find("name").text

xmlbox = obj.find('bndbox')

bb = convert_box((w, h), [float(xmlbox.find(x).text) for x in ('xmin', 'xmax', 'ymin', 'ymax')])

cls_id = class_ids.index(cls) # class id

out_file.write(" ".join([str(a) for a in (cls_id, *bb)]) + '\n')

copy(os.path.join(path_images, image_names), os.path.join(target_dir, 'images', k, image_names))

out_file.close()

转换之后，数据集格式如下：

dataset

|--- images

|--- train

|--- xx.png

|--- val

|--- xx.png

|--- labels

|--- train

|--- xx.txt

|--- val

|--- xx.txt

3.2 训练YOLOv8模型

创建数据集描述文件helmet.yml，helmet.yml描述了数据集的存放位置以及数据类别信息

train: C:/Users/dataset/images/train # train images (relative to 'path') 4 images

val: C:/Users/dataset/images/val # val images (relative to 'path') 4 images

test: # test images (optional)

# Classes

names:

0: helmet

1: head

2: person

模型训练：

需要GPU，否则训练非常慢

$ yolo task=detect mode=train model=yolov8n.pt data=C:/Users/helmet.yml device=0 batch=16 epochs=50