yolov8识别图片并获取识别信息 - 幽络源

背景

有个项目要求：使用yolov8训练出的模型将图片进行识别后，需要返回识别后的信息，信息中包括 耗时(毫秒)、类别、类别对应的方框坐标，本篇教程便是完成此功能。

前提

有自己的yolov8环境且有训练好的模型

直接上源码

注意：这里提供了两个函数，一个是只返回信息不返回识别后的图片，一个是即返回信息又返回图片

import base64
import os
import shutil
import time
import uuid

'''
输入：图 picId , 图片路径 picPath , 识别模型 identifyYolo
输出：
{   
    "picId" : "图id",
    "wastedTime" : 耗时(毫秒),
    "results" : [
        {
            "cls" : "类别",
            "location" : "方框坐标"
        }
    ]
}
'''
def PicPathToInfo(picId,picPath,identifyYolo):
    #进行识别并记录耗时(毫秒)
    startTime = time.time()
    results = identifyYolo(source=picPath, save=False) #不需要保存识别后的图片，因此将save设置False
    endTime = time.time()
    wastedTime = int((endTime - startTime) * 1000)

    # 获取类别名称列表
    class_names = identifyYolo.names

    #因为只处理一张图，所以取第0个结果就行了
    result=results[0]
    #获取所有检测结果
    detections = result.boxes
    #构造返回结果
    resDict= {"picId":picId,"wastedTime": wastedTime, "results": []}
    #遍历检测结果
    for detection in detections:
        cls_id = int(detection.cls)  # 类别ID
        cls_name = class_names[cls_id]  # 类别名称
        location = detection.xyxy[0].tolist()  # 获取边界框坐标 [x_min, y_min, x_max, y_max]
        obj={
            "cls":cls_name,
            "location":location
        }
        resDict["results"].append(obj)
    return resDict

'''
输入：图片二进制文件流 picBytes , 识别模型 identifyYolo
输出：
{
    "wastedTime" : 耗时(毫秒),
    "identifiedPicBytesBase64" : 识别后的图片二进制文件流的Base54编码
    "results" : [
        {
            "cls" : "类别",
            "location" : "方框坐标"
        }
    ]
}
'''
def PicBytesToIdentifiedPicBytes(picBytes,identifyYolo):
    # 创建原图临时目录以保存图片（如果不存在），后续删掉里面的图片
    tempOriginDir="originTemp"
    if not os.path.exists(tempOriginDir):
        os.makedirs(tempOriginDir)

    uuidName = str(uuid.uuid4()) + '.jpg'  #将图片保存为jpg格式

    # 将二进制数据写入临时目录
    with open(os.path.join(tempOriginDir,uuidName), 'wb') as f:
        f.write(picBytes)

    # 进行识别并记录耗时(毫秒)
    startTime = time.time()
    results = identifyYolo(source=os.path.join(tempOriginDir,uuidName), save=True)
    endTime = time.time()
    wastedTime = int((endTime - startTime) * 1000)
    # 删掉原图，防止占用空间
    os.remove(os.path.join(tempOriginDir,uuidName))

    # 获取类别名称列表
    class_names = identifyYolo.names

    # 因为只处理一张图，所以取第0个结果就行了
    result = results[0]
    # 获取识别后的图片
    identifyPicPath=os.path.join(result.save_dir,uuidName)
    # 将其读取为二进制文件流
    with open(identifyPicPath, 'rb') as picFile:
        identifiedPicBytes = picFile.read()
    #将识别后的图片的所在目录删掉，防止占用空间
    shutil.rmtree(result.save_dir)
    # 将文件流转为Base64编码用于传入json
    identifiedPicBytesBase64=base64.b64encode(identifiedPicBytes).decode('utf-8')

    # 获取所有检测结果
    detections = result.boxes
    # 构造返回结果
    resDict = {"identifiedPicBytesBase64":identifiedPicBytesBase64,"wastedTime": wastedTime, "results": []}
    # 遍历检测结果
    for detection in detections:
        cls_id = int(detection.cls)  # 类别ID
        cls_name = class_names[cls_id]  # 类别名称
        location = detection.xyxy[0].tolist()  # 获取边界框坐标 [x_min, y_min, x_max, y_max]
        obj = {
            "cls": cls_name,
            "location": location
        }
        resDict["results"].append(obj)
    return resDict

使用方法

PicPathToInfo的使用

1.在yolv8源码根目录新建IdentifyUtil.py，将上面的源码粘贴进去，如图

2.新建test.py，然后这里我训练好的模型名为Merged.pt，因此需先加载Merged.pt获取yolo对象，如图

from ultralytics import YOLO

#混合缺陷识别模型
mergedYolo=YOLO(model=r"D:\A01PythonProjects3123\ultralytics-main\Merged.pt",task="detect")

3.调用函数进行识别并获取信息

PicPathToInfo参数解释：

• picId：这个原本是用于后端项目上的，不需要的可以随便传，也可以自行删掉
• picUrl：可以是本地路径，也可以是http链接的图片
• identtifyYolo：加载了模型的yolo对象

如图，我这里因为是本地测试，随便传了个id和本地的图片，最终返回的结果为json数据，非常方便

PicBytesToIdentifiedPicBytes的使用

1.同样的，需要先加载自己的模型获取yolo对象，和上面第2步一样

2.需注意的是，这个方法要求我们传入图片的二进制文件流，那就找一张本地的图片，先获取它的二进制文件流，代码与图如下：

from ultralytics import YOLO
from IdentifyUtil import *

#混合缺陷识别模型
mergedYolo=YOLO(model=r"D:\A01PythonProjects3123\ultralytics-main\Merged.pt",task="detect")

#读取图片为二进制文件流
picPath=r"D:\测试图.jpg"
with open(picPath, 'rb') as f:
    picBytes=f.read()

#调用函数识别图片获取识别信息
mergedResult=PicBytesToIdentifiedPicBytes(picBytes,mergedYolo)
print(mergedResult)

3.执行后可以看到返回的json结果中包括了如下三个

• wastedTime：耗时(毫秒)
• results：包括了识别出的类别和类别对应的方框坐标
• identifiedPicBytesBase64：这个是图片的二进制文件流的Base64编码，我们可以通过这个数据来获取到识别后的图片

4.将返回的Base64码写为图片，代码与图如下：

#将Base64码写为图片
image_data = base64.b64decode(mergedResult["identifiedPicBytesBase64"])
image_bytes = io.BytesIO(image_data)
image = Image.open(image_bytes)
# 保存图片到目录
savePath = "savePath"
if not os.path.exists(savePath):
    os.mkdir(savePath)
output_path = savePath + "/识别后的图.jpg"
image.save(output_path)

你可能的疑问

你或许会问：对于PicBytesToIdentifiedPicBytes函数为什么要传入流再返回识别后的图片的Base64编码？yolov8本身就可以在本地生成识别后的图片啊？

我想说的是这两个方法主要是用于远程调用的，比如我在一台服务器部署了识别服务，我需要在本地用http去调用这个服务，那么此时就需要通过传文件流和图片Base64编码来进行图片的传输了

原文链接：yolov8识别图片并获取识别信息 – 幽络源