醋醋百科网

Good Luck To You!

FastAPI+yolov8构建图像识别服务(图像识别api接口)

近期搞智能体应用,需要使用各种工具模型的服务,包括图像识别,OCR等。于是用FastAPI构建了一系列的工具服务,这是其中之一。

YOLO

UltralyticsYOLOv8基于深度学习和计算机视觉领域的尖端技术,在速度和准确性方面具有无与伦比的性能。其流线型设计使其适用于各种应用,并可轻松适应从边缘设备到云 API 等不同硬件平台。

模型地址:Ultralytics/YOLOv8 · Hugging Face

文档地址:
https://docs.ultralytics.com/

添加图片注释,不超过 140 字(可选)

YOLO(You Only Look Once)是一种流行的物体检测和图像分割模型,由华盛顿大学的约瑟夫-雷德蒙(Joseph Redmon)和阿里-法哈迪(Ali Farhadi)开发。YOLO 于 2015 年推出,因其高速度和高精确度而迅速受到欢迎。

  • 2016 年发布的YOLOv2 通过纳入批量归一化、锚框和维度集群改进了原始模型。
  • 2018 年推出的YOLOv3 使用更高效的骨干网络、多锚和空间金字塔池进一步增强了模型的性能。
  • YOLOv4于 2020 年发布,引入了 Mosaic 数据增强、新的无锚检测头和新的损失函数等创新技术。
  • YOLOv5进一步提高了模型的性能,并增加了超参数优化、集成实验跟踪和自动导出为常用导出格式等新功能。
  • YOLOv6于 2022 年由美团开源,目前已用于该公司的许多自主配送机器人。
  • YOLOv7增加了额外的任务,如 COCO 关键点数据集的姿势估计。
  • YOLOv8是YOLO 的最新版本,由Ultralytics 提供。YOLOv8 YOLOv8 支持全方位的视觉 AI 任务,包括检测、分割、姿态估计、跟踪和分类。这种多功能性使用户能够在各种应用和领域中利用YOLOv8 的功能。
  • YOLOv9引入了可编程梯度信息 (PGI) 和通用高效层聚合网络 (GELAN) 等创新方法。

FastAPI

Git地址:
https://github.com/tiangolo/fastapi

添加图片注释,不超过 140 字(可选)

添加图片注释,不超过 140 字(可选)

FastAPI 是一个现代、快速(高性能)的 Web 框架,用于基于标准 Python 类型提示使用 Python 构建 API。

主要特点是:

  • 快速:非常高的性能,与NodeJS和Go相当(感谢 Starlette 和 Pydantic)。可用的最快的 Python 框架之一。
  • 快速编码:将开发功能的速度提高约 200% 至 300%。*
  • 更少的错误:减少约 40% 的人为(开发人员)引起的错误。*
  • 直观:强大的编辑器支持。到处都完成。调试时间更少。
  • 简单:旨在易于使用和学习。阅读文档的时间更少。
  • 简短:最大限度地减少代码重复。每个参数声明具有多个功能。更少的错误。
  • 健壮:获取可用于生产的代码。具有自动交互式文档。
  • 基于标准:基于(并完全兼容)API 的开放标准:OpenAPI(以前称为 Swagger)和JSON Schema。

服务构建

安装python包

需要安装ultralytics,uvicorn,gunicorn,fastapi,loguru等包,我的环境以前都装过。具体版本没看,大概最新的就行。

代码编写

代码包括main.py和app.py,具体代码如下

main.py

####################################### import模块 #################################
import json
import pandas as pd
from PIL import Image
from loguru import logger
import sys

from fastapi import FastAPI, File, status
from fastapi.responses import RedirectResponse
from fastapi.responses import StreamingResponse
from fastapi.middleware.cors import CORSMiddleware
from fastapi.exceptions import HTTPException

from io import BytesIO

from app import get_image_from_bytes
from app import detect_sample_model
from app import add_bboxs_on_img
from app import get_bytes_from_image

####################################### 日志 #################################

logger.remove()
logger.add(
    sys.stderr,
    colorize=True,
    format="<green>{time:HH:mm:ss}</green> | <level>{message}</level>",
    level=10,
)
logger.add("log.log", rotation="1 MB", level="DEBUG", compression="zip")

###################### FastAPI 设置 #############################

# 标题
app = FastAPI(
    title="Object Detection FastAPI 模板",
    description="""从图像中获取对象值
                    并返回图像和JSON结果""",
    version="2023.1.31",
)

# 如果您希望允许来自特定域(在origins参数中指定)的客户端请求
# 访问FastAPI服务器的资源,并且客户端和服务器托管在不同的域上,则需要此功能。
origins = [
    "http://localhost",
    "http://localhost:8008",
    "*"
]

app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

@app.on_event("startup")
def save_openapi_json():
    '''此函数用于将FastAPI应用程序的OpenAPI文档数据保存到JSON文件中。
    保存OpenAPI文档数据的目的是拥有API规范的永久和离线记录,
    可用于文档目的或生成客户端库。虽然不一定需要,但在某些情况下可能会有帮助。'''
    openapi_data = app.openapi()
    # 将"openapi.json"更改为所需的文件名
    with open("openapi.json", "w") as file:
        json.dump(openapi_data, file)

# 重定向
@app.get("/", include_in_schema=False)
async def redirect():
    return RedirectResponse("/docs")

@app.get('/healthcheck', status_code=status.HTTP_200_OK)
def perform_healthcheck():
    '''
    它发送一个GET请求到该路由,并希望得到一个"200"响应代码。
    未能返回200响应代码将使GitHub Actions回滚到项目处于"工作状态"的最后一个版本。
    它作为最后一道防线,以防发生问题。
    此外,它还以JSON格式返回响应,形式为:
    {
        'healthcheck': '一切正常!'
    }
    '''
    return {'healthcheck': '一切正常!'}

######################### 支持函数 #################################

def crop_image_by_predict(image: Image, predict: pd.DataFrame(), crop_class_name: str,) -> Image:
    """根据图像中某个对象的检测结果裁剪图像。
    
    参数:
        image: 要裁剪的图像。
        predict (pd.DataFrame): 包含对象检测模型预测结果的数据框。
        crop_class_name (str, 可选): 要根据其裁剪图像的对象类名称。如果未提供,函数将返回图像中找到的第一个对象。
    
    返回:
        Image: 裁剪后的图像或None
    """
    crop_predicts = predict[(predict['name'] == crop_class_name)]

    if crop_predicts.empty:
        raise HTTPException(status_code=400, detail=f"照片中未找到{crop_class_name}")

    # 如果有多个检测结果,选择置信度更高的那个
    if len(crop_predicts) > 1:
        crop_predicts = crop_predicts.sort_values(by=['confidence'], ascending=False)

    crop_bbox = crop_predicts[['xmin', 'ymin', 'xmax','ymax']].iloc[0].values
    # 裁剪
    img_crop = image.crop(crop_bbox)
    return(img_crop)

######################### 主功能 #################################

@app.post("/img_object_detection_to_json")
def img_object_detection_to_json(file: bytes = File(...)):
    """
    从图像中进行对象检测。

    参数:
        file (bytes): 以字节格式的图像文件。
    返回:
        dict: 包含对象检测结果的JSON格式。
    """
    # 步骤1:用None值初始化结果字典
    result = {'detect_objects': None}

    # 步骤2:将图像文件转换为图像对象
    input_image = get_image_from_bytes(file)

    # 步骤3:从模型中进行预测
    predict = detect_sample_model(input_image)

    # 步骤4:选择检测对象返回信息
    # 您可以在此选择要发送到结果中的数据
    detect_res = predict[['name', 'confidence']]
    objects = detect_res['name'].values

    result['detect_objects_names'] = ', '.join(objects)
    result['detect_objects'] = json.loads(detect_res.to_json(orient='records'))

    # 步骤5:日志记录和返回
    logger.info("结果: {}", result)
    return result

@app.post("/img_object_detection_to_img")
def img_object_detection_to_img(file: bytes = File(...)):
    """
    从图像中进行对象检测并在图像上绘制边界框

    参数:
        file (bytes): 以字节格式的图像文件。
    返回:
        Image: 带有边界框注释的字节格式图像。
    """
    # 从字节获取图像
    input_image = get_image_from_bytes(file)

    # 模型预测
    predict = detect_sample_model(input_image)

    # 在图像上添加边界框
    final_image = add_bboxs_on_img(image=input_image, predict=predict)

    # 以字节格式返回图像
    return StreamingResponse(content=get_bytes_from_image(final_image), media_type="image/jpeg")

app.py

from PIL import Image
import io
import pandas as pd
import numpy as np

from typing import Optional

from ultralytics import YOLO
from ultralytics.utils.plotting import Annotator, colors


# 初始化模型
model_sample_model = YOLO("./models/sample_model/yolov8n.pt")


def get_image_from_bytes(binary_image: bytes) -> Image:
    """将字节格式的图像转换为PIL RGB格式
    
    参数:
        binary_image (bytes): 图像的字节表示
    
    返回:
        PIL.Image: PIL RGB格式的图像
    """
    input_image = Image.open(io.BytesIO(binary_image)).convert("RGB")
    return input_image


def get_bytes_from_image(image: Image) -> bytes:
    """
    将PIL图像转换为字节
    
    参数:
    image (Image): PIL图像实例
    
    返回:
    bytes : 包含JPEG格式图像的BytesIO对象,质量为85
    """
    return_image = io.BytesIO()
    image.save(return_image, format='JPEG', quality=85)  # 将图像保存为JPEG格式,质量为85
    return_image.seek(0)  # 将指针设置到文件开头
    return return_image

def transform_predict_to_df(results: list, labeles_dict: dict) -> pd.DataFrame:
    """
    将yolov8的预测结果(torch.Tensor)转换为pandas DataFrame。

    参数:
        results (list): 包含yolov8预测输出的列表,形式为torch.Tensor。
        labeles_dict (dict): 包含标签名称的字典,其中键为类id,值为标签名称。
        
    返回:
        predict_bbox (pd.DataFrame): 包含边界框坐标、置信度分数和类别标签的DataFrame。
    """
    # 将Tensor转换为numpy数组
    predict_bbox = pd.DataFrame(results[0].to("cpu").numpy().boxes.xyxy, columns=['xmin', 'ymin', 'xmax','ymax'])
    # 将预测的置信度添加到DataFrame
    predict_bbox['confidence'] = results[0].to("cpu").numpy().boxes.conf
    # 将预测的类别添加到DataFrame
    predict_bbox['class'] = (results[0].to("cpu").numpy().boxes.cls).astype(int)
    # 用labeles_dict中的类名替换类编号
    predict_bbox['name'] = predict_bbox["class"].replace(labeles_dict)
    return predict_bbox

def get_model_predict(model: YOLO, input_image: Image, save: bool = False, image_size: int = 1248, conf: float = 0.5, augment: bool = False) -> pd.DataFrame:
    """
    获取模型对输入图像的预测结果。
    
    参数:
        model (YOLO): 训练好的YOLO模型。
        input_image (Image): 模型将进行预测的图像。
        save (bool, 可选): 是否保存带有预测结果的图像。默认为False。
        image_size (int, 可选): 模型接收的图像大小。默认为1248。
        conf (float, 可选): 预测的置信度阈值。默认为0.5。
        augment (bool, 可选): 是否对输入图像应用数据增强。默认为False。
    
    返回:
        pd.DataFrame: 包含预测结果的DataFrame。
    """
    # 进行预测
    predictions = model.predict(
                        imgsz=image_size, 
                        source=input_image, 
                        conf=conf,
                        save=save, 
                        augment=augment,
                        flipud= 0.0,
                        fliplr= 0.0,
                        mosaic = 0.0,
                        )
    
    # 将预测结果转换为pandas DataFrame
    predictions = transform_predict_to_df(predictions, model.model.names)
    return predictions


################################# 边界框功能 #####################################

def add_bboxs_on_img(image: Image, predict: pd.DataFrame()) -> Image:
    """
    在图像上添加边界框

    参数:
    image (Image): 输入图像
    predict (pd.DataFrame): 模型的预测结果

    返回:
    Image: 带有边界框的图像
    """
    # 创建注释器对象
    annotator = Annotator(np.array(image))

    # 按xmin值对预测结果进行排序
    predict = predict.sort_values(by=['xmin'], ascending=True)

    # 迭代预测结果的每一行
    for i, row in predict.iterrows():
        # 创建要显示在图像上的文本
        text = f"{row['name']}: {int(row['confidence']*100)}%"
        # 获取边界框坐标
        bbox = [row['xmin'], row['ymin'], row['xmax'], row['ymax']]
        # 在图像上添加边界框和文本
        annotator.box_label(bbox, text, color=colors(row['class'], True))
    # 将带注释的图像转换为PIL图像
    return Image.fromarray(annotator.result())


################################# 模型 #####################################


def detect_sample_model(input_image: Image) -> pd.DataFrame:
    """
    使用sample_model进行预测。
    基于YoloV8

    参数:
        input_image (Image): 输入图像。

    返回:
        pd.DataFrame: 包含对象位置的DataFrame。
    """
    predict = get_model_predict(
        model=model_sample_model,
        input_image=input_image,
        save=False,
        image_size=640,
        augment=False,
        conf=0.5,
    )
    return predict


将main.py和app.py保存在同一个目录下。运行:

uvicorn main:app --reload --host 0.0.0.0 --port 8001

即可启动服务。第一次启动的时候会自动下载模型。

Downloading https://github.com/ultralytics/assets/releases/download/v8.2.0/yolov8n.pt to 'models\sample_model\yolov8n.pt

访问
http://localhost:8001/docs#/可以查看API介绍

添加图片注释,不超过 140 字(可选)

效果测试:

测试图片:

添加图片注释,不超过 140 字(可选)

调用
img_object_detection_to_json:

import requests

input_image_name = 'cat1.png'
api_host = 'http://127.0.0.1:8001/'
type_rq = 'img_object_detection_to_json'

files = {'file': open(input_image_name, 'rb')}

response = requests.post(api_host+type_rq, files=files)

data = response.json()     
print(data)

添加图片注释,不超过 140 字(可选)

调用
img_object_detection_to_img:

import requests
from PIL import Image
from io import BytesIO
import matplotlib.pyplot as plt

input_image_name = 'cat1.png'
api_host = 'http://127.0.0.1:8001/'
type_rq = 'img_object_detection_to_img'

files = {'file': open(input_image_name, 'rb')}

response = requests.post(api_host+type_rq, files=files)

img = Image.open(BytesIO(response.content)) 
plt.imshow(img)

返回:

添加图片注释,不超过 140 字(可选)


有试了其他一些图片:

添加图片注释,不超过 140 字(可选)

添加图片注释,不超过 140 字(可选)

效果都很不错。

控制面板
您好,欢迎到访网站!
  查看权限
网站分类
最新留言