FastAPI集成自定义YOLOv4 Darknet模型时遇IndexError问题求助

问题详情 #

尝试将基于Darknet自定义训练的YOLOv4模型集成到FastAPI中，调用接口时控制台抛出IndexError: invalid index to scalar variable，FastAPI本地服务器返回Internal Server Error，预期接口返回目标检测的边界框与标签。

API代码 #

import cv2
import numpy as np
import tensorflow as tf

# 加载YOLO v4模型
model = cv2.dnn.readNetFromDarknet("yolov4_test.cfg", "yolov4_train_final.weights")

from fastapi import FastAPI, File, UploadFile
from typing import List, Tuple

app = FastAPI()

@app.post("/detect_objects")
async def detect_objects(image: UploadFile = File(...)) -> List[Tuple[str, Tuple[int, int, int, int]]]:
    # 读取图片文件
    image_bytes = await image.read()
    nparr = np.frombuffer(image_bytes, np.uint8)
    img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)

    # 运行YOLO v4模型
    blob = cv2.dnn.blobFromImage(img, 1/255.0, (608, 608), swapRB=True, crop=False)
    model.setInput(blob)
    layer_names = model.getLayerNames()
    output_layers = [layer_names[i[0] - 1] for i in model.getUnconnectedOutLayers()]
    outputs = model.forward(output_layers)

    # 提取边界框和类别标签
    boxes = []
    for output in outputs:
        for detection in output:
            scores = detection[5:]
            class_id = np.argmax(scores)
            confidence = scores[class_id]
            if confidence > 0.5:
                center_x = int(detection[0] * img.shape[1])
                center_y = int(detection[1] * img.shape[0])
                width = int(detection[2] * img.shape[1])
                height = int(detection[3] * img.shape[0])
                left = int(center_x - width / 2)
                top = int(center_y - height / 2)
                boxes.append((class_id, (left, top, width, height)))

    # 映射类别ID到标签
    classes = ["Gun"]
    results = []
    for box in boxes:
        class_label = classes[box[0]]
        bbox = box[1]
        results.append((class_label, bbox))
        
    return results

控制台报错堆栈 #

Traceback (most recent call last):
  File "c:\users\raafeh\desktop\fyp\envfast\lib\site-packages\uvicorn\protocols\http\h11_impl.py", line 428, in run_asgi
    result = await app(  # type: ignore[func-returns-value]
  File "c:\users\raafeh\desktop\fyp\envfast\lib\site-packages\uvicorn\middleware\proxy_headers.py", line 78, in __call__
    return await self.app(scope, receive, send)
  File "c:\users\raafeh\desktop\fyp\envfast\lib\site-packages\fastapi\applications.py", line 276, in __call__
    await super().__call__(scope, receive, send)
  File "c:\users\raafeh\desktop\fyp\envfast\lib\site-packages\starlette\applications.py", line 122, in __call__
    await self.middleware_stack(scope, receive, send)
  File "c:\users\raafeh\desktop\fyp\envfast\lib\site-packages\starlette\middleware\errors.py", line 184, in __call__
    raise exc
  File "c:\users\raafeh\desktop\fyp\envfast\lib\site-packages\starlette\middleware\errors.py", line 162, in __call__
    await self.app(scope, receive, _send)
  File "c:\users\raafeh\desktop\fyp\envfast\lib\site-packages\starlette\middleware\exceptions.py", line 79, in __call__
    raise exc
  File "c:\users\raafeh\desktop\fyp\envfast\lib\site-packages\starlette\middleware\exceptions.py", line 68, in __call__
    await self.app(scope, receive, sender)
  File "c:\users\raafeh\desktop\fyp\envfast\lib\site-packages\fastapi\middleware\asyncexitstack.py", line 21, in __call__
    raise e
  File "c:\users\raafeh\desktop\fyp\envfast\lib\site-packages\fastapi\middleware\asyncexitstack.py", line 18, in __call__
    await self.app(scope, receive, send)
  File "c:\users\raafeh\desktop\fyp\envfast\lib\site-packages\starlette\routing.py", line 718, in __call__
    await route.handle(scope, receive, send)
  File "c:\users\raafeh\desktop\fyp\envfast\lib\site-packages\starlette\routing.py", line 276, in handle
    await self.app(scope, receive, send)
  File "c:\users\raafeh\desktop\fyp\envfast\lib\site-packages\starlette\routing.py", line 66, in app
    response = await func(request)
  File "c:\users\raafeh\desktop\fyp\envfast\lib\site-packages\fastapi\routing.py", line 237, in app
    raw_response = await run_endpoint_function(
  File "c:\users\raafeh\desktop\fyp\envfast\lib\site-packages\fastapi\routing.py", line 163, in run_endpoint_function
    return await dependant.call(**values)
  File "C:\Users\Raafeh\Desktop\FYP\main.py", line 27, in detect_objects
    output_layers = [layer_names[i[0] - 1] for i in model.getUnconnectedOutLayers()]
  File "C:\Users\Raafeh\Desktop\FYP\main.py", line 27, in <listcomp>
    output_layers = [layer_names[i[0] - 1] for i in model.getUnconnectedOutLayers()]
IndexError: invalid index to scalar variable.

解决方案 #

错误出在获取输出层的代码行，原因是不同OpenCV版本中getUnconnectedOutLayers()的返回格式不同：

• 旧版本OpenCV返回二维数组，因此需要用i[0]取值
• 新版本OpenCV返回一维数组，每个元素是标量，直接使用i即可

修改以下代码行：

# 原代码
output_layers = [layer_names[i[0] - 1] for i in model.getUnconnectedOutLayers()]

# 修改后代码
output_layers = [layer_names[i - 1] for i in model.getUnconnectedOutLayers()]

如果要兼容新旧版本，可以增加格式判断：

unconnected_layers = model.getUnconnectedOutLayers()
# 检查返回是否为二维数组
if len(unconnected_layers.shape) > 1:
    output_layers = [layer_names[i[0] - 1] for i in unconnected_layers]
else:
    output_layers = [layer_names[i - 1] for i in unconnected_layers]

修改后重新启动FastAPI服务，即可正常获取检测结果。

内容的提问来源于stack exchange，提问作者Raafeh Sajjad