FastAPI+TF-Serving部署深度学习模型：输入尺寸不兼容问题解决

我在StackOverflow上的第一个问题：尝试用FastAPI和TensorFlow Serving部署深度学习模型，运行代码时出现如下错误：

ValueError: Input 0 of layer "model" is incompatible with the layer: expected shape=(None, 224, 224, 3), found shape=(None, 3088, 3088, 3)

需求：编写代码将任意输入图片调整为模型训练时使用的(224, 224)尺寸，解决上述错误。

现有代码文件 #

main.py #

from fastapi import FastAPI, File, UploadFile
from fastapi.middleware.cors import CORSMiddleware
import uvicorn
import numpy as np
from io import BytesIO
from PIL import Image
import tensorflow as tf

app = FastAPI()

origins = [
    "http://localhost",
    "http://localhost:3000",
]

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

MODEL = tf.keras.models.load_model("../saved_models/1")

CLASS_NAMES = ["diseased cotton leaf", "diseased cotton plant", "fresh cotton leaf", "fresh cotton plant"]

@app.get("/ping")
async def ping():
    return "Hello, I am alive"

def read_file_as_image(data) -> np.ndarray:
    image = np.array(Image.open(BytesIO(data)))
    return image

@app.post("/predict")
async def predict(
    file: UploadFile = File(...)
):
    image = read_file_as_image(await file.read())
    img_batch = np.expand_dims(image, 0)
    
    predictions = MODEL.predict(img_batch)
    
    predicted_class = CLASS_NAMES[np.argmax(predictions[0])]
    confidence = np.max(predictions[0])
    return {
        'class': predicted_class
    }

if __name__ == "__main__":
    uvicorn.run(app, host = 'localhost', port = 8000)

main-tf-serving.py #

from fastapi import FastAPI, File, UploadFile
import uvicorn
import numpy as np
from io import BytesIO
from PIL import Image
import tensorflow as tf
import requests

app = FastAPI()

endpoint = "http://localhost:8502/v1/models/cotton-models:predict"

CLASS_NAMES = ["diseased cotton leaf", "diseased cotton plant", "fresh cotton leaf", "fresh cotton plant"]

@app.get("/ping")
async def ping():
    return "Hello, I am alive"

def read_file_as_image(data) -> np.ndarray:
    image = np.array(Image.open(BytesIO(data)))
    return image

@app.post("/predict")
async def predict(
    file: UploadFile = File(...)
):
    image = read_file_as_image(await file.read())
    img_batch = np.expand_dims(image, 0)

    json_data = {
        "instances" : img_batch.tolist()
    }
    
    response = requests.post(endpoint, json=json_data)
    prediction= np.array(response.json()["predictions"][0])

    predicted_class = CLASS_NAMES[np.argmax(prediction)]
    confidence = np.max(prediction)
    
    return {
        'class': predicted_class
    }

if __name__ == "__main__":
    uvicorn.run(app, host = 'localhost', port = 8000)

解决方案：调整输入图片尺寸 #

错误核心是模型训练时固定输入尺寸为(224,224,3)，但上传的图片尺寸不符合要求。只需在读取图片后添加尺寸调整步骤，用PIL库的resize方法即可实现。

修改read_file_as_image函数 #

将原函数替换为以下代码：

def read_file_as_image(data) -> np.ndarray:
    # 打开图片文件
    image = Image.open(BytesIO(data))
    # 确保图片为RGB三通道（兼容灰度图上传场景）
    if image.mode != 'RGB':
        image = image.convert('RGB')
    # 调整图片尺寸到模型要求的(224,224)
    image_resized = image.resize((224, 224))
    # 转换为numpy数组供模型使用
    return np.array(image_resized)

可选：同步训练时的预处理步骤 #

如果训练模型时对图片做了归一化（比如将像素值缩放到[0,1]区间），需要在调整尺寸后同步该操作，示例如下：

def read_file_as_image(data) -> np.ndarray:
    image = Image.open(BytesIO(data))
    if image.mode != 'RGB':
        image = image.convert('RGB')
    image_resized = image.resize((224, 224))
    # 归一化像素值（根据训练时的预处理逻辑调整）
    image_array = np.array(image_resized) / 255.0
    return image_array

将上述修改应用到main.py和main-tf-serving.py的对应函数中，重启服务后即可解决输入尺寸不匹配的问题。

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