在FastAPI中对图像进行分类时返回标签

from fastapi import FastAPI, UploadFile, File
import tensorflow as tf
import numpy as np

model = tf.keras.models.load_model('path_to_model')

app = FastAPI()

@app.post("/classify_image")
async def classify_image(file: UploadFile = File(...)):
    # 读取上传的图像文件
    image = tf.image.decode_image(await file.read(), channels=3)
    image = tf.image.resize(image, (224, 224))  # 调整图像大小

    # 预处理图像
    image = image / 255.0  # 归一化
    image = tf.expand_dims(image, axis=0)  # 添加批次维度

    # 使用模型进行图像分类
    predictions = model.predict(image)
    predicted_label = np.argmax(predictions[0])

    # 返回分类结果
    return {"label": predicted_label}

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