基于YOLOv8算法的照片角度分类项目实践

一、任务概述

最近一个项目，需要针对库存的车辆照片运用人工智能算法进行照片朝向分类和矫正，算法在设计时需要满足轻量化需求，适合在CPU环境中进行快速推理。在具体实现时，可以将照片分为4个类别：ni_0、ni_90、ni_180、ni_270，分别表示照片经过0度、90度、180度、270度逆向旋转。综和考虑算法精度和速度要求，本文拟采用YOLOv8算法来实现该任务。

YOLOv8 是当前业界领先的感知模型，它建立在以前 YOLO 版本的基础上，引入了新的功能并做了相关改进，提升了性能和灵活性。具体创新包括一个新的骨干网络、一个新的 Ancher-Free 检测头和一个新的损失函数，可以在从 CPU 到 GPU 的各种硬件平台上运行。

二、YOLOv8算法简介

• YOLOv8作者：glenn-jocher
• 项目链接：https://github.com/ultralytics/ultralytics

YOLO 是一种基于图像全局信息进行预测的目标检测系统。自 2015 年 Joseph Redmon、Ali Farhadi 等人提出初代模型以来，领域内的研究者们已经对 YOLO 进行了多次更新迭代，模型性能越来越强大。当前最新版本为YOLOv8。

具体的，YOLOv8 是由小型初创公司 Ultralytics 创建并维护的，值得注意的是 YOLOv5 也是由该公司创建的。

2.1 算法改进

YOLOv8 算法的核心特性和改动可以归纳如下：

• 基本特性：提供了一个全新的 SOTA 模型，包括P5 640 和 P6 1280 分辨率的目标检测网络和基于 YOLACT 的实例分割模型。和 YOLOv5 一样，基于缩放系数也提供了 N/S/M/L/X 尺度的不同大小模型，用于满足不同场景需求。
• 骨干网络和 Neck 网络：参考了 YOLOv7 ELAN 设计思想，将 YOLOv5 的 C3 结构换成了梯度流更丰富的 C2f 结构，并对不同尺度模型调整了不同的通道数，这种精心微调方式大幅提升了模型性能。不过这个C2f 模块中存在 Split 等操作，对特定硬件部署没有之前那么友好了。
• Head网络：相比 YOLOv5 改动较大，换成了目前主流的解耦头结构，将分类和检测头分离，同时也从 Anchor-Based 换成了 Anchor-Free。
• Loss损失函数：采用了 TaskAlignedAssigner 正样本分配策略，并引入了 Distribution Focal Loss。

从上面可以看出，YOLOv8 主要参考了最近提出的诸如 YOLOX、YOLOv6、YOLOv7 和 PPYOLOE 等算法的相关设计，本身的创新点不多，偏向工程实践。

2.2 算法特点

• 对用户友好的 API（命令行 + Python）；
• 模型更快更准确；
• 模型能完成常见的图像感知任务，包括图像分类、目标检测、实例分割、关键点检测和视频跟踪；
• 与先前所有版本的 YOLO 兼容可扩展；

2.3 网络结构

上图所示即为YOLOv8的完整模型结构图。

查看 N/S/M/L/X 等不同大小模型，可以发现 N/S 和 L/X 两组模型只是改了缩放系数，但是 S/M/L 等骨干网络的通道数设置不一样，没有遵循同一套缩放系数。如此设计的原因应该是同一套缩放系数下的通道设置不是最优设计，YOLOv7 网络设计时也没有遵循一套缩放系数作用于所有模型。

Head 部分变化最大，从原先的耦合头变成了解耦头，其结构如下所示：

可以看出，不再有之前的 objectness 分支，只有解耦的分类和回归分支，并且其回归分支使用了 Distribution Focal Loss 中提出的积分形式表示法。

2.4 性能比较

YOLOv8是由小型初创公司 Ultralytics 创建并维护的，不过 Ultralytics 并没有直接将开源库命名为 YOLOv8，而是直接使用 Ultralytics 这个词，原因是 Ultralytics 将这个库定位为算法框架，而非某一个特定算法。

Ultralytics 开源库的两个主要优点是：

• 融合众多当前 SOTA 技术于一体；
• 未来将支持其他 YOLO 系列以及 YOLO 之外的更多算法；

三、工程实践

下面开始针对实际工程任务进行操作。

3.1 安装算法框架库ultralytics

YOLOv8算法位于开源框架库ultralytics中，因此先要安装ultralytics。
安装方式如下：

 pip install ultralytics

为了方便后续配置和使用，可以将github上的ultralytics源码拉取到本地：

git clone https://github.com/ultralytics/ultralytics.git@main

3.2 库存照片预处理

3.2.1 提取所有图片

假设库存照片位于名为“第一批原始照片”的文件夹中，由于库存照片目录结构混乱，图像格式不统一，因此需要将所有图片提取到一个统一的文件夹中，并且所有图片以jpg格式保存，这样方便后续使用。

在同目录下创建文件夹car_data/1，然后使用下面的脚本完成图片提取和转换。

import os
import cv2
import numpy as np


def getFileList(dir, Filelist, ext=None):
    """
    获取文件夹及其子文件夹中文件列表
    输入 dir:文件夹根目录
    输入 ext:扩展名
    返回： 文件路径列表
    """
    newDir = dir
    if os.path.isfile(dir):
        if ext is None:
            Filelist.append(dir)
        else:
            if ext in dir[-3:]:
                Filelist.append(dir)

    elif os.path.isdir(dir):
        for s in os.listdir(dir):
            newDir = os.path.join(dir, s)
            getFileList(newDir, Filelist, ext)

    return Filelist


org_img_folder = "./第一批原始照片"

# 检索文件
imglist = getFileList(org_img_folder, [], "jpg")
print("本次执行检索到 " + str(len(imglist)) + " 张图像\n")

imgIndex = 1
for imgpath in imglist:
    print(imgpath)
    try:
        img = cv2.imdecode(np.fromfile(imgpath, dtype=np.uint8), -1)
        if img is None:
            print('读取失败')
            continue
        if len(img.shape) == 2:
            img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
        elif img.shape[2] == 1:
            img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
        elif img.shape[2] == 4:
            img = cv2.cvtColor(img, cv2.COLOR_BGRA2BGR)
        savepath = os.path.join("car_data/1", 'jianyan_' + str(imgIndex) + ".jpg")
        cv2.imwrite(savepath, img)
        imgIndex += 1
    except:
        print("异常")
    else:
        pass

print("完成")

上述脚本用来提取库存中的jpg照片，如果库存中还存在png或bmp图片，那么就修改代码：

imglist = getFileList(org_img_folder, [], "jpg")

将其中的jpg修改为png或bmp，同时修改对应的imgIndex起始标签值。

所有图片提取完以后都存放在car_data/1文件夹中，以jpg格式存储。库存总图片数达到89786张。

3.2.2 去除冗余的相同照片

库存照片中可能存在相同照片多次存放的问题，因此需要将完全相同的图像剔除掉，减少冗余。本文使用哈希比对算法来实现，具体脚本代码如下：

import os
import cv2
import shutil
from PIL import Image
import imagehash

srcFolder = 'car_data/1'
dstFolder = 'delimgs'
imgnames = os.listdir(srcFolder) 

# 计算所有图像哈希值
hashlst = []
for imgname in imgnames:
    print('计算哈希值中  '+imgname)
    hash_size = 16
    imgpath = os.path.join(srcFolder,imgname)
    hash = imagehash.dhash(Image.open(imgpath),hash_size=hash_size)
    hashlst.append(hash)

# 检索相似图片
for curIndex in range(len(imgnames)-1):
    hash1 = hashlst[curIndex]
    print('比对中  '+imgnames[curIndex])
    for compIndex in range(curIndex+1,len(imgnames)):
        hash2 =  hashlst[compIndex]
        if hash1==hash2:
            imgname = imgnames[curIndex]
            dstpath = os.path.join(dstFolder, imgname)
            shutil.move(os.path.join(srcFolder,imgnames[curIndex]), dstpath)
            imgname = imgname.split('.')[0]
            dstpath = os.path.join(dstFolder, imgname+'_compare.jpg')
            shutil.copyfile(os.path.join(srcFolder,imgnames[compIndex]), dstpath)
            print('找到相同文件')
            break
print('完成')

去除冗余后的库存总图片数达到74125张。

3.2.3 去除无车辆照片

库存照片中存在大量无车辆的错误照片，因此需要写一个脚本将无车辆照片剔除掉。这里使用预先在coco数据集上训练好的yolov8算法来实现。具体脚本代码如下：

import cv2, os, shutil
from ultralytics import YOLO

# 检索文件夹
folderpath = "./car_data/1"
dstFolder = './delimgs'
imgnames = os.listdir(folderpath)

# 加载模型
model = YOLO("models/yolov8m-seg.pt")

# 循环处理
for imgname in imgnames:
    # 读取图像
    imgpath = os.path.join(folderpath, imgname)
    print(imgpath)
    img = cv2.imread(imgpath)
    if img is None:
        os.remove(imgpath)
        continue

    # 车辆检索
    result = model(img, imgsz=640, conf=0.5)[0]
    boxes = result.boxes
    isfind = False
    for box in boxes:
        classlabel = box.cls.cpu().numpy()[0]
        if classlabel == 1 or classlabel == 2 or classlabel == 3 or classlabel == 5 or classlabel == 7:
            isfind = True
            break
        
    # 没找到车辆，删除图像
    if not isfind:
        dstpath = os.path.join(dstFolder, imgname)
        shutil.move(imgpath, dstpath)

print('完成')

在delimgs文件夹中存放着剔除掉的图像，由于算法存在一定的漏检率，因此有些存在车辆的照片被错误的移动到这个delimgs文件夹中，需要人工复核，将这些照片“捞回去”。

去除无车辆照片后，库存总图片数达到51180张。

3.2.4 随机提取指定数量的图片

库存照片数据量庞大，本文只需要提取2万多张图片用来训练算法即可。

import os
import shutil,random

srcFolder =  'car_data/1'
dstFolder = 'car_data/2'

if not os.path.exists(dstFolder):
    os.makedirs(dstFolder)

picIndex = 1
imgnames = os.listdir(srcFolder)
random.shuffle(imgnames)
for imgname in imgnames:
    if picIndex > 25000:
        continue
    imgpath = os.path.join(srcFolder, imgname)
    dstpath = os.path.join(dstFolder, imgname)
    shutil.move(imgpath, dstpath)
    picIndex += 1

print('完成')  

提取好的图片位于car_data/2文件夹中，总数25000张。

3.3 照片朝向分类

3.3.1 数据准备

首先从预处理后的库存照片中精心挑选照片朝向正确的图像共计2万张整，然后分别对这2万张图像进行旋转，得到对应的逆90、逆180、逆270度角的三个类别图像，这样就组成了可以用来分类的图像库photo_direction，共计8万张图像，分4个类别。

完整生成脚本如下：

import os
from PIL import Image

ni0_folder = "./dataset/car_data/ni0"
ni90_folder = "./dataset/car_data/ni90"
ni180_folder = "./dataset/car_data/ni180"
ni270_folder = "./dataset/car_data/ni270"

# 创建文件夹
if not os.path.exists(ni90_folder):
    os.makedirs(ni90_folder)
if not os.path.exists(ni180_folder):
    os.makedirs(ni180_folder)
if not os.path.exists(ni270_folder):
    os.makedirs(ni270_folder)

# 检索图像
img_names = os.listdir(ni0_folder)
for img_name in img_names:
    img_path = os.path.join(ni0_folder, img_name)
    print(img_path)
    # 读取图像
    img = Image.open(img_path)
    # 逆时针旋转90
    img90 = img.transpose(Image.ROTATE_90)
    save_path = os.path.join(ni90_folder, img_name)
    img90.save(save_path)
    # 逆时针旋转180
    img180 = img.transpose(Image.ROTATE_180)
    save_path = os.path.join(ni180_folder, img_name)
    img180.save(save_path)
    # 逆时针旋转270
    img270 = img.transpose(Image.ROTATE_270)
    save_path = os.path.join(ni270_folder, img_name)
    img270.save(save_path)

print("完成")

其中ni0、ni90、ni180、ni270分别存储了逆时针0°、90°、180°、270°对应的图像。最后从每个文件夹中随机抽取2000张图片作为测试集用来评估算法。

最终数据集目录结构整理如下：

dataset/car_data/
|
|-- train/
|   |-- ni0/
|   |   |-- 10008.jpg
|   |   |-- 10009.jpg
|   |   |-- ...
|   |
|   |-- ni90/
|   |   |-- 1000.jpg
|   |   |-- 1001.jpg
|   |   |-- ...
|   |
|   |-- ni180/
|   |   |-- 10014.jpg
|   |   |-- 10015.jpg
|   |   |-- ...
|   |
|   |-- ni270/
|   |   |-- 10014.jpg
|   |   |-- 10015.jpg
|   |   |-- ...
|   |
|
|-- test/
|   |-- ni0/
|   |   |-- 10.jpg
|   |   |-- 11.jpg
|   |   |-- ...
|   |
|   |-- ni90/
|   |   |-- 12.jpg
|   |   |-- 13.jpg
|   |   |-- ...
|   |
|   |-- ni180/
|   |   |-- 14.jpg
|   |   |-- 15.jpg
|   |   |-- ...
|   |
|   |-- ni270/
|   |   |-- 16.jpg
|   |   |-- 17.jpg
|   |   |-- ...

上述结构就是ultralytics的图像分类所需要的目录结构，整个数据集分为train和test两个文件夹，其中每个种类的图片都放在一起，每个种类的文件夹名称即为对应的类别名称。

3.3.2 训练和评估

找到ultralytics/cfg/models/v8中找到yolov8-cls.yaml文件，拷贝一份到ultralytics/configs目录下面，并重命名为yolov8-cls-photodirection.yaml，修改该文件中的nc参数为4，表示共有4个类别。

训练代码如下：

from ultralytics import YOLO

# 加载预训练模型和配置文件
model = YOLO('./configs/yolov8m-cls-photodirection.yaml').load('yolov8m-cls.pt')

# 训练模型
results = model.train(data='./dataset/car_data', epochs=100, imgsz=64, batch=32, device='0,1')

启动训练后如果本地没有预训练模型yolov8-cls.pt，则ultralytics框架会自动从github上进行下载。需要注意的是创建的yaml名称为yolov8-cls-photodirection.yaml，而在代码中调用的是yolov8s-cls-photodirection.yaml，这是ultralytics框架提供的一个功能，我们只需要配置一份yaml文件，即可适配不同规模任务的分类模型，包括：

• YOLOv8n-cls
• YOLOv8s-cls
• YOLOv8m-cls
• YOLOv8l-cls
• YOLOv8x-cls

本文共使用7万多张照片在2个GPU上进行训练，测试集为2000张图片，总耗时约17个小时。在测试集上的最佳top1准确率为0.985。

3.3.3 预测

训练好模型以后，使用下面的代码可以对单张图片进行预测和矫正：

from ultralytics import YOLO
import cv2
from PIL import Image
import numpy as np

# 加载模型
model = YOLO('./runs/classify/train/weights/best.pt')

# 预测模型
img = cv2.imread('./imgs/7.jpg')
results = model(img)
label = int(results[0].probs.top1) # 标签类别
labelconf = results[0].probs.top1conf.cpu().numpy() # 置信度
print(label)
print(labelconf)

# 矫正
Thr = 0.8
if labelconf > Thr:
    if label == 1: # 逆时针180度
        img = Image.fromarray(cv2.cvtColor(img,cv2.COLOR_BGR2RGB))
        img = img.transpose(Image.ROTATE_180)
        img = cv2.cvtColor(np.asarray(img),cv2.COLOR_RGB2BGR)
        save_path = './imgs/result.jpg'
        cv2.imwrite(save_path,img)
    elif label == 2: # 逆时针270度
        img = Image.fromarray(cv2.cvtColor(img,cv2.COLOR_BGR2RGB))
        img = img.transpose(Image.ROTATE_90)
        img = cv2.cvtColor(np.asarray(img),cv2.COLOR_RGB2BGR)
        save_path = './imgs/result.jpg'
        cv2.imwrite(save_path,img)
    elif label == 3: # 逆时针90度
        img = Image.fromarray(cv2.cvtColor(img,cv2.COLOR_BGR2RGB))
        img = img.transpose(Image.ROTATE_270)
        img = cv2.cvtColor(np.asarray(img),cv2.COLOR_RGB2BGR)
        save_path = './imgs/result.jpg'
        cv2.imwrite(save_path,img)

最终输出的是分类标签、置信度以及矫正过后的车辆照片。