如何用MLFlow部署YOLOv8目标检测模型？求相关参考方案

问题核心 #

本地通过mlflow.pyfunc.load_model加载模型时，可直接传入图像路径或NumPy数组字符串，但部署为API服务后无法沿用相同格式传递图像——原因是MLFlow部署的模型服务默认接收JSON格式输入，必须适配图像的编码与解码逻辑。

调整模型Wrapper类 #

修改YOLOWrapper的predict方法，同时兼容本地调用和API服务场景的输入格式：

import base64
import numpy as np
import cv2
from ultralytics import YOLO
import mlflow.pyfunc

class YOLOWrapper(mlflow.pyfunc.PythonModel):
    def load_context(self, context):
        self.model = YOLO(context.artifacts["yolo_model"])

    def predict(self, context, model_input):
        # 处理API传入的Base64编码图像
        if isinstance(model_input, str) and model_input.startswith("data:image/"):
            base64_data = model_input.split(",")[1]
            img_bytes = base64.b64decode(base64_data)
            img_np = np.frombuffer(img_bytes, np.uint8)
            img = cv2.imdecode(img_np, cv2.IMREAD_COLOR)
        # 处理本地调用的图像路径
        elif isinstance(model_input, str) and model_input.endswith((".jpg", ".png", ".jpeg")):
            img = model_input
        # 处理NumPy数组字符串格式
        elif isinstance(model_input, str):
            img = np.array(eval(model_input))
        else:
            raise ValueError("不支持的输入格式，请传入图像路径、Base64编码图像或NumPy数组字符串")
        
        results = self.model(img)
        # 将YOLO结果转为可JSON序列化的格式（避免部署时返回无法序列化的对象）
        output = []
        for result in results:
            pred = {
                "boxes": result.boxes.xyxy.cpu().numpy().tolist(),
                "confidences": result.boxes.conf.cpu().numpy().tolist(),
                "class_ids": result.boxes.cls.cpu().numpy().tolist(),
                "class_names": [result.names[int(cls)] for cls in result.boxes.cls.cpu().numpy()]
            }
            output.append(pred)
        return output

重新注册模型 #

修改Wrapper后，重新打包并注册模型到MLFlow，确保artifacts目录包含YOLOv8的模型文件（如best.pt）。

部署后的API调用示例 #

部署模型后，通过HTTP请求调用时需将图像转为Base64编码传入：

Python请求示例 #

import requests
import base64

def image_to_base64(img_path):
    with open(img_path, "rb") as f:
        img_bytes = f.read()
    return f"data:image/jpeg;base64,{base64.b64encode(img_bytes).decode('utf-8')}"

# 替换为你的模型服务地址
url = "http://localhost:5000/invocations"
headers = {"Content-Type": "application/json"}

img_base64 = image_to_base64("test_image.jpg")
data = {"dataframe_split": {"columns": ["input"], "data": [[img_base64]]}}

response = requests.post(url, json=data, headers=headers)
print(response.json())

额外建议 #

• 启动MLFlow模型服务时，指定正确的模型URI和端口，例如：mlflow models serve -m models:/YOLOv8/latest -p 5000
• 禁止在predict方法中返回YOLO原生的Results对象，必须转换为列表、字典等基础数据类型，否则会出现JSON序列化错误
• 可按需扩展输入格式，比如支持直接传入图像二进制数据（需调整请求的Content-Type为application/octet-stream并修改Wrapper的处理逻辑）

内容的提问来源于stack exchange，提问作者atharva parikh