使用craft-text-detector时遇cv2.polylines参数错误的解决求助

问题背景 #

• 本地通过图片路径调用craft-text-detector实现文本裁剪功能完全正常
• 改为FastAPI服务接收Base64编码的图片字节流后，部分图片处理时触发cv2.polylines参数错误

错误信息 #

INFO:     127.0.0.1:36386 - "POST /detect-text-and-make-crops HTTP/1.1" 500 Internal Server Error
ERROR:    Exception in ASGI application
Traceback (most recent call last):
  File "../python3.10/site-packages/uvicorn/protocols/http/h11_impl.py", line 408, in run_asgi
    result = await app(  # type: ignore[func-returns-value]
  File "../python3.10/site-packages/uvicorn/middleware/proxy_headers.py", line 84, in __call__
    return await self.app(scope, receive, send)
  File "../project.py", line 73, in detect_text
    export_extra_results(
  File "../python3.10/site-packages/craft_text_detector/file_utils.py", line 246, in export_extra_results
    cv2.polylines(
cv2.error: OpenCV(4.5.4) :-1: error: (-5:Bad argument) in function 'polylines'
> Overload resolution failed:
>  - Layout of the output array img is incompatible with cv::Mat
>  - Expected Ptr<cv::UMat> for argument 'img'

解决方案 #

错误根源是部分图片为RGBA四通道格式，cv2.polylines对这类格式的兼容性不佳，同时需要确保图像数组内存连续。具体修复步骤如下：

1. 转换图像为RGB三通道格式：将cv2.imdecode读取的图像统一转为RGB格式，避免四通道导致的兼容问题
2. 确保数组内存连续：使用np.ascontiguousarray处理图像数组，满足OpenCV对数组内存布局的要求
3. 优化模型加载：将模型加载逻辑移到函数外部，避免每次请求重复加载模型，提升服务性能

修改后的代码 #

import base64
import logging
import cv2
from fastapi import FastAPI, HTTPException
from fastapi.encoders import jsonable_encoder
from fastapi.responses import JSONResponse
from craft_text_detector import (
    read_image,
    load_craftnet_model,
    load_refinenet_model,
    get_prediction,
    export_detected_regions,
    export_extra_results,
    empty_cuda_cache
)
import numpy as np
from pydantic import BaseModel, ValidationError

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI()

# 提前加载模型，避免每次请求重复加载
refine_net = load_refinenet_model(cuda=False)
craft_net = load_craftnet_model(cuda=False)

class ImageRequest(BaseModel):
    image_bytes: bytes
    output_dir: str

class TextDetectionResponse(BaseModel):
    message: str
    output_dir: str


def detect_text(image_bytes: bytes, output_dir: str):    
    image_array = np.frombuffer(image_bytes, dtype=np.uint8)
    mat_like_image = cv2.imdecode(image_array, cv2.IMREAD_COLOR)  # 直接读取为RGB三通道
    
    # 确保图像数组内存连续
    mat_like_image = np.ascontiguousarray(mat_like_image)
    
    image = read_image(mat_like_image)

    prediction_result = get_prediction(
        image=image,
        craft_net=craft_net,
        refine_net=refine_net,
        text_threshold=0.7,
        link_threshold=0.4,
        low_text=0.4,
        cuda=False,
        long_size=1280
    )

    exported_file_paths = export_detected_regions(
        image=image,
        regions=prediction_result["boxes"],
        output_dir=output_dir,
        rectify=True
    )

    export_extra_results(
        image=image,
        regions=prediction_result["boxes"],
        heatmaps=prediction_result["heatmaps"],
        output_dir=output_dir
    )

    empty_cuda_cache()


@app.post("/detect-text-and-make-crops")
async def detect_text_and_make_crops(request: ImageRequest):
    try:
        output_dir = request.output_dir
        updated_image_bytes = base64.b64decode(request.image_bytes)
        detect_text(updated_image_bytes, output_dir)
        
        text_detection_response = TextDetectionResponse(
            message='Text Detection Completed, and Crops are Saved',
            output_dir=output_dir
        )
        return JSONResponse(content=jsonable_encoder(text_detection_response))
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

if __name__ == '__main__':
    import uvicorn
    uvicorn.run("__main__:app", host="0.0.0.0", port=8090)

关键修改点说明 #

• 将cv2.imdecode的参数从cv2.IMREAD_UNCHANGED改为cv2.IMREAD_COLOR，强制读取为三通道RGB图像
• 添加np.ascontiguousarray(mat_like_image)确保数组内存连续，符合OpenCV函数的要求
• 将模型加载逻辑移到全局作用域，避免每次请求重复加载模型，大幅提升服务响应速度
• 扩展异常捕获范围，捕获所有类型的异常并返回更详细的错误信息

内容的提问来源于stack exchange，提问作者Zaryab Ali