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

芯查查已开通飞腾派交流社区，此文为社区文章摘选，欢迎大家到社区交流

在前面的文章中，介绍了在飞腾派上实现数字识别，本篇文章更进一步与实际需求结合，更进一步在飞腾派上实现安全帽识别。
在社会智能化的发展趋势之下，越来越多的传统行业开始向着数字化的方向转型，而建筑行业也正经历着通过人工智能技术实现的改革。安全帽是保护建筑工人的重要装备，在建筑工地、煤矿等场景中，佩戴安全帽是保护工人生命安全的必要措施。然而，有些工人或管理人员可能会心存侥幸，或者因为各种原因不佩戴安全帽，而巡查人员也不可能时刻注意是否有员工没佩戴安全帽，此时，安全帽识别技术就成了一个极其便利的科技手段，通过检测和识别工人是否佩戴安全帽，可以有效提高安全生产监管水平，减少安全事故的发生。
本篇文章将借助安全帽识别这一实际问题，介绍在飞腾派上部署目标检测算法实现对象识别。

1. 目标检测

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

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

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
Epoch    GPU_mem   box_loss   cls_loss   dfl_loss  Instances       Size
       1/10      1.24G      1.624      1.789      1.276         31        640: 100%|██████████| 500/500 [00:57<00:00,  8.62it/s]
                 Class     Images  Instances      Box(P          R      mAP50  mAP50-95): 100%|██████████| 63/63 [00:09<00:00,  6.52it/s]
                   all       1000       5272      0.867      0.492      0.534      0.302

      Epoch    GPU_mem   box_loss   cls_loss   dfl_loss  Instances       Size
       2/10      1.35G      1.548      1.206      1.239         18        640: 100%|██████████| 500/500 [00:53<00:00,  9.37it/s]
                 Class     Images  Instances      Box(P          R      mAP50  mAP50-95): 100%|██████████| 63/63 [00:09<00:00,  6.77it/s]
                   all       1000       5272      0.919      0.515      0.575      0.318

      Epoch    GPU_mem   box_loss   cls_loss   dfl_loss  Instances       Size
       3/10      1.34G      1.525      1.058      1.234         30        640: 100%|██████████| 500/500 [00:49<00:00, 10.05it/s]
                 Class     Images  Instances      Box(P          R      mAP50  mAP50-95): 100%|██████████| 63/63 [00:08<00:00,  7.32it/s]
                   all       1000       5272      0.928      0.528      0.587      0.335

      Epoch    GPU_mem   box_loss   cls_loss   dfl_loss  Instances       Size
       4/10      1.33G      1.498     0.9742      1.218         42        640: 100%|██████████| 500/500 [00:49<00:00, 10.05it/s]
                 Class     Images  Instances      Box(P          R      mAP50  mAP50-95): 100%|██████████| 63/63 [00:08<00:00,  7.66it/s]
                   all       1000       5272      0.928      0.535      0.592      0.322

      Epoch    GPU_mem   box_loss   cls_loss   dfl_loss  Instances       Size
       5/10      1.33G      1.489     0.9235      1.214         30        640: 100%|██████████| 500/500 [00:47<00:00, 10.59it/s]
                 Class     Images  Instances      Box(P          R      mAP50  mAP50-95): 100%|██████████| 63/63 [00:08<00:00,  7.49it/s]

3.3 导出tflite模型

上篇文章中已经在飞腾派中安装了tflite，因此为了在飞腾派部署yolo，还需要导出tflite模型

安装依赖项

$ python -m pip install tensorflow>2.4.1 onnx>=1.12.0 onnxsim>=0.4.1 'sng4onnx>=1.0.1' 'onnx_graphsurgeon>=0.3.26'\
$ 'onnx2tf>=1.15.4,<=1.17.5' 'tflite_support' 'onnxruntime' -i https://mirror.baidu.com/pypi/simple

导出模型

$ yolo export model=./runs/detect/train/weights/best.pt format=tflite

3.4 在飞腾派部署YOLO模型

现在让我们教飞腾派检测安全帽吧

• 安装Ultralytics

user@phytiumpi:~/Documents/cv/dl$ python -m pip install ultralytics -i https://mirror.baidu.com/pypi/simple

推理

user@phytiumpi:~/Documents/cv/dl/yolo$ yolo detect predict model=/home/user/Documents/cv/dl/yolo/best_float16.tflite source=/home/user/Documents/cv/dl/yolo/hard_hat_workers1.png
Ultralytics YOLOv8.0.196 ???? Python-3.9.2 torch-2.1.0 CPU (aarch64)
Loading /home/user/Documents/cv/dl/yolo/best_float16.tflite for TensorFlow Lite inference...
INFO: Created TensorFlow Lite XNNPACK delegate for CPU.

image 1/1 /home/user/Documents/cv/dl/yolo/hard_hat_workers1.png: 640x640 9 helmets, 959.0ms
Speed: 17.2ms preprocess, 959.0ms inference, 6.7ms postprocess per image at shape (1, 3, 640, 640)
Results saved to runs/detect/predict

看起来飞腾派已经可以成功的识别安全帽了！（虽然纯CPU有点慢）

5. 总结

本篇文章中将主要介绍了借助YOLO使得飞腾派可以完成识别安全帽，虽然纯CPU有点慢，但是可以借助飞腾派上的 Mini PCIe 连接TPU对神经网络进行加速（ PS. 如果 芯查查 有这样的模块就更棒了！ ）。接下来的文章将简要的介绍使用飞腾派进行音频方面的识别。

由于Float32模型较大，因此附件只放置了Int8量化模型 yolo_helmet_int8.tflite 以供测试