使用ultralytics制作服务化的yolov11-obb旋转目标检测

🕗 发布于 2024-11-05 17:02 YOLO 目标检测 人工智能

前言

本文章的ultralytics版本为Ultralytics v8.1.0

通过官网的yolov11n-obb实现旋转目标检测服务化

创建环境

conda create -n yolov11_obb python=3.11

在这里插入图片描述

安装ultralytics库

这个库基本上囊括了了我们需要的一切，注意源指定国内的源，我这里用的清华源

pip install ultralytics --trusted-host pypi.tuna.tsinghua.edu.cn -i https://pypi.tuna.tsinghua.edu.cn/simple

在这里插入图片描述

安装FastAPI服务框架

pip install fastapi uvicorn pydantic -i https://pypi.tuna.tsinghua.edu.cn/simple

配置开发环境连接远程环境

打开本地电脑的pycharm，添加解释器
在这里插入图片描述
选择ssh

选择conda环境，配置远程同步目录
1.conda环境选择你刚才创建的conda的环境名称
2.下面的同步目录选择你远程的代码目录，到时候你本地运行的代码看到的结果，实际上就是这个远程服务器上的代码。
在这里插入图片描述
有些同学配置错误后，会有各种问题，注意检查你运行代码的时候，这里执行的代码是不是你和这里配置的路径一样，不然会出现你本地改了代码，但是远程服务器上的代码和你不同步。导致结果和预期不一致的问题。
在这里插入图片描述

下载权重

https://github.com/ultralytics/assets/releases
在这里插入图片描述

代码部分

demo.py

这个函数实现主要的推理函数

import base64
import io

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

# set version
model_version = 'test_obb'

# Load a model
model = YOLO("model/yolo11n-obb.pt")  # load an official model
# model = YOLO("path/to/best.pt")  # load a custom model

# Predict with the model
results = model("boats.jpg")  # predict on an image

# 定义数据集名称列表
coco_class_names = [
    'plane', 'ship', 'storage tank', 'baseball diamond', 'tennis court',
    'basketball court', 'ground track field', 'harbor', 'bridge', 'large vehicle',
    'small vehicle', 'helicopter', 'roundabout', 'soccer ball field', 'swimming pool'
]

coco_class_names_cn = [
    '飞机', '船舶', '储罐', '棒球场', '网球场',
    '篮球场', '田径场', '港口', '桥梁', '大型车辆',
    '小型车辆', '直升机', '环岛', '足球场', '游泳池'
]


model_conf=0.2


def predict(image, conf=model_conf):

    # 这部分逻辑放在服务化的main.py
    # 将base64字符串转换回图像格式
    # try:
    #     # 解码base64字符串
    #     img_bytes = base64.b64decode(image)
    #
    #     # 将字节转换为numpy数组
    #     nparr = np.frombuffer(img_bytes, np.uint8)
    #
    #     # 解码为图像
    #     img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
    #
    #     if img is None:
    #         raise ValueError("Failed to decode image")
    #
    # except Exception as e:
    #     raise ValueError(f"Invalid image data: {str(e)}")
    # # conf默认使用model_conf变量的值,也可以接受传参

    model.conf = conf

    # 对传入的图像数据进行预测
    results = model(image, save=True)  # save=True会保存带标注的图像

    data = []  # 存储推理结果

    # 遍历每个结果
    for result in results:
        obb = result.obb
        if obb is not None:
            # 遍历每个检测框
            for i, (box, conf, cls) in enumerate(zip(obb.data, obb.conf, obb.cls)):
                # 获取框的各个参数
                cx, cy, w, h, angle = box[:5].cpu().numpy()
                cls_id = int(cls.item())

                # 构造每个预测结果的字典
                item = {
                    'index': i,  # 预测框序号
                    'cls': cls_id,  # 类别id
                    'label': coco_class_names[cls_id],  # 类别英文名称
                    'label_cn': coco_class_names_cn[cls_id],  # 类别中文名称
                    'confidence': float(conf.item()),  # 置信度
                    'x': float(cx),  # 框中心x
                    'y': float(cy),  # 框中心y
                    'w': float(w),  # 框宽度
                    'h': float(h),  # 框高度
                    'angle': float(angle * 180 / math.pi),  # 旋转角(度数)
                    'angle_rad': float(angle)  # 旋转角(弧度)
                }
                data.append(item)

    # 获取带标注的图像并转换为base64
    try:
        # 获取结果图像
        annotated_img = results[0].plot()  # 返回带标注的numpy数组

        # 将numpy数组转换为PIL Image
        im = Image.fromarray(annotated_img[..., ::-1])  # BGR转RGB

        # 将图像转换为base64字符串
        buffered = io.BytesIO()
        im.save(buffered, format="JPEG")
            # # 这样是对的
            # im.save('output.jpg', format='JPEG')  # 文件名可以用.jpg
            # im.save('output.jpeg', format='JPEG')  # 文件名可以用.jpeg
            # im.save(buffered, format='JPEG')  # format参数必须用JPEG
            #
            # # 这样是错的
            # im.save(buffered, format='JPG')  # 会报错
        img_str = base64.b64encode(buffered.getvalue()).decode('utf-8')

    except Exception as e:
        print(f"Warning: Failed to generate annotated image: {str(e)}")
        img_str = ""

    # 构造返回的JSON格式结果
    pred_data = {
        'version': model_version,
        'inference_time': f"{results[0].speed['inference'] / 1000:.3f}",  # 直接从results获取推理时间，默认的时间是ms为单位，我们给搞成s
        'results': data,
        'detimg': img_str,  # base64编码的图像
    }

    return pred_data

# 测试predict函数调用
if __name__ == "__main__":

    # 读取本地图片
    image_path = "boats.jpg"
    img = cv2.imread(image_path)

    # 调用predict并输出结果
    results = predict(image=img)
    print(json.dumps(results, indent=2, ensure_ascii=False))

main.py

这个函数实现FastAPI的服务化调用

# -*- coding: utf-8 -*-
# @Time    : 2024年09月05日
# @Author  : Mark White
# @FileName: main.py
# @Software: PyCharm


# 标准库导入
import base64
import io
from typing import Optional

# 第三方库导入
# pip install fastapi uvicorn pydantic
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from PIL import Image
import uvicorn

#自定义库导入
# LOOK AT ME!!这个demo函数是您们自己的实现,采用其他命名也完全可以
import demo  # 导入自定义预测函数模块


app = FastAPI()  # 创建FastAPI应用对象

class PredictRequest(BaseModel): # 固定写法,图像匹配类
    image: str
    conf: Optional[float] = None
    """
    用于目标检测类算法的请求模型。

    属性:
        image (str): Base64 编码的图像字符串。图像应该被编码为 base64 格式。
        conf (Optional[float]): 可选的置信度阈值。
                                用于过滤检测结果，仅返回置信度高于此阈值的结果。
                                如果不提供，将使用demo.py里配置的默认阈值。

    示例:
        {
            "image": "/9j/4AAQSkZJRgABAQAAAQABAAD/2wBDAA...",
            "conf": 0.5
        }

    注意:
        - 这是一个 Pydantic BaseModel，用于验证和序列化/反序列化请求数据。
        - 'image' 字段是必需的，而 'conf' 字段是可选的。
    """

app = FastAPI() #固定写法,创建FastAPI应用对象

@app.post('/predict') #固定写法,目标检测类算法,定义预测接口路径和方法
async def predict(request: PredictRequest):
    try:
        image_data = base64.b64decode(request.image)  # 解码base64编码的图片数据
        image = Image.open(io.BytesIO(image_data))  # 使用PIL库读取图片数据到Image对象

        # LOOK AT ME!! 可以从这里开始修改您们的逻辑
        if request.conf is not None:
            results = demo.predict(image=image,conf=request.conf)  # 带阈值预测
        else:
            results = demo.predict(image=image)  # 不带阈值预测

        return results  # 返回预测结果

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


# FastAPI固定写法,一般不用修改
if __name__ == '__main__':
    # uvicorn.run() 用于启动 ASGI 服务器来运行 FastAPI 应用
    uvicorn.run(
        "main:app",  # 使用字符串 "文件名:app变量名"
        host='0.0.0.0',  # 允许来自任何 IP 地址的连接，使服务器可以被公开访问
        port=8886,  # 指定服务器应监听的端口。如果容器需要开启多个端口，请联系管理员
        workers=4,  # 指定工作进程的数量
        # 如果您的服务是 I/O 密集型，可以考虑增加 workers 数量,若修改此参数,请同步管理员确认服务器资源容量
        reload=False,  # 生产环境建议设置为 False
        log_level="info",  # 日志级别，可选：critical, error, warning, info, debug, trace
    )

# 注意事项：
# 1. 确保 main.py 文件中定义了名为 'app' 的 FastAPI 应用实例
# 2. workers 数量应根据服务器 CPU 核心数和内存情况进行调整
# 3. 对于 I/O 密集型应用，可以设置比 CPU 核心数更多的 workers

构建镜像

ssh连接到服务器远程环境

把env下的yolov11n_obb拷贝到当前的代码目录

使用conda info 查看我们的目录在哪里

conda info

在这里插入图片描述
然后我们把这个目录下的文件拷贝到我们的代码目录

cp -r /home/baijs/miniconda3/envs/yolov11_obb ./env_yolov11_obb

编写Dockerfile

vim Dockerfile

内容如下：

FROM nvidia/cuda:11.4.0-cudnn8-devel-ubuntu20.04

# 创建环境目录和工作目录
RUN mkdir -p /env

# 这里最好加上，安装vim和一些依赖
RUN apt-get update && apt-get install -y \
    vim \
    libgl1-mesa-glx \
    libglib2.0-0 \
    libsm6 \
    libxext6 \
    libxrender-dev \
    && apt-get clean \

WORKDIR /workdir

# 复制conda环境到/env目录
COPY env_yolov11_obb /env/yolov11_obb

# 设置环境变量
ENV PATH=/env/yolov11_obb/bin:$PATH

# 复制当前目录所有文件到工作目录
COPY . /workdir/

# 暴露端口
EXPOSE 8886

# 设置启动命令
CMD ["python", "main.py"]

构建镜像

docker build -t yolov11n_obb:v1 .

启动镜像

docker run -it -p 8886:8886 yolov11n_obb:v1

在这里插入图片描述

测试服务

这个测试代码参考下一节

bash test_curl.txt

在这里插入图片描述

其他工具代码

测试图片 boats.images

这个图片可以通过官方实例代码获得，或者保存我这里的图片也可以

results = model("https://ultralytics.com/images/boats.jpg")  # predict on an image

在这里插入图片描述

test_gen.py

这个代码手动执行，生成一个测试脚本，用来调用本地服务化接口，因为base64字符串很长，所以他会生成一个json文件，然后用curl指令里@文件的方式来发起调用，这个代码会生成request_data.json和test_curl.txt

你需要修改这个代码里的路由函数和端口，要和你的函数实现一致。

import base64
import os


def generate_curl_command():
    # 读取图片并转base64
    with open("boats.jpg", "rb") as f:
        img_base64 = base64.b64encode(f.read()).decode('utf-8')

    # 创建JSON数据文件
    json_data = {
        "image": img_base64,
        "conf": 0.5
    }

    # 保存JSON数据到文件
    with open("request_data.json", "w") as f:
        json.dump(json_data, f)  # 使用json.dump而不是json.dumps

    # 生成使用-d @file的curl命令
    curl_command = """curl -X POST "http://localhost:8886/predict" \\
     -H "Content-Type: application/json" \\
     -d @request_data.json"""

    with open("test_curl.txt", "w") as f:
        f.write(curl_command)

    print("已生成curl命令到 test_curl.txt")
    print("使用方法: bash test_curl.txt")


if __name__ == "__main__":
    import json

    generate_curl_command()

request_data.json

由test_gen.py生成
内容长这样

{"image": "/9j/4i超级长的字符串，我缩略了uZ1No", "conf": 0.5}

test_curl.txt

由test_gen.py生成
内容长这样：

curl -X POST "http://localhost:38886/predict" \
     -H "Content-Type: application/json" \
     -d @request_data.json

然后执行的时候就如我上述测试的一样

bash test_curl.txt

原文地址：https://blog.csdn.net/crazyjinks/article/details/143479755

免责声明：本站文章内容转载自网络资源，如本站内容侵犯了原著者的合法权益，可联系本站删除。更多内容请关注自学内容网（zxcms.com）！

上一篇：基于SSM+小程序的高校寻物平台管理系统（失物1）
下一篇：Rust常用属性及应用

linux shell脚本学习（1）：shell脚本基本概念与操作
shebang指的是出现在shell脚本第一行的前2个字符，unix系统的程序会分析shebang后的内容，以作为解释器的指令，例如：（1）/bin/sh，程序会调用bash解释器。/bin/sh是/
阅读更多2024-11-05
批量图片转PDF文件的多种方法详解
根据需求和可用工具，可以选择以上任一方法来将批量图片转换为PDF文件。在线工具简单快捷，桌面应用程序功能强大，而编程方式则适合批量处理和自动化。
阅读更多2024-11-05
linux 上使用 D-Bus 控制Onboard
linux 屏幕键盘 onboard 控制
阅读更多2024-11-05
使用JavaScript实现新窗口打开并设置sessionStorage的简单指南
我们将编写一个名为navigateTo的函数，该函数接受一个参数id，用于构建一个新的URL，并在新窗口中打开该URL。同时，在新窗口加载完成后，我们将设置其sessionStorage。通过本文，我
阅读更多2024-11-05
项目模块1~12总结：服务器大模块梳理
分析完通信套接字的回调函数后分析监听套接字的读回调，监听套接字需要对接收的通信套接字进行处理（用 Connection 管理，设置通信套接字的5种回调，开启超时连接销毁，启动连接，添加连接），处理的过
阅读更多2024-11-05
服务器对于企业业务有哪些影响？
企业服务器可以充当网络服务提供商，为企业和客户提供各种服务，同时服务器一般都会配备强大的计算能力和处理能力，可以处理企业中的大量数据，如数据库服务器可以用来存储和管理企业的数据，为企业的应用程序提供数
阅读更多2024-11-05
SRS：构建实时免费视频服务器的全方位指南
SRS是一个功能强大的实时视频服务器，它支持多种流媒体协议和功能，如推流、拉流、转码、截图和流转发。通过详细的配置和测试，我们可以利用SRS构建一个高效、稳定的直播系统。希望这篇文章能帮助你更好地理解
阅读更多2024-11-05
数据库基础(2) . 安装MySQL
添加到鼠标右键运行 reg文件在注册表中添加信息这样在右键菜单中就有以管理员身份打开命令行的选项了。
阅读更多2024-11-05
TapData 发布官方性能测试报告，针对各流行数据源，在多项指标中表现拔群
TapData 性能测试报告已上线
阅读更多2024-11-05
【深度学习|地学应用】人工智能技术的发展历程与现状：探讨深度学习在遥感地学中的应用前景
【深度学习|地学应用】人工智能技术的发展历程与现状：探讨深度学习在遥感地学中的应用前景
阅读更多2024-11-05

使用ultralytics制作服务化的yolov11-obb旋转目标检测

前言

创建环境

安装ultralytics库

安装FastAPI服务框架

配置开发环境连接远程环境

下载权重

代码部分

demo.py

main.py

构建镜像

把env下的yolov11n_obb拷贝到当前的代码目录

编写Dockerfile

启动镜像

测试服务

其他工具代码

测试图片 boats.images

test_gen.py

request_data.json

test_curl.txt

相关文章