109 lines
3.0 KiB
Python
109 lines
3.0 KiB
Python
from fastapi import FastAPI, HTTPException, UploadFile, File
|
|
from pydantic import BaseModel
|
|
import uvicorn
|
|
import json
|
|
import logging
|
|
import base64
|
|
from io import BytesIO
|
|
from .ai_algorithm import ImageRecognizer
|
|
from .config import settings
|
|
|
|
# 配置日志
|
|
logging.basicConfig(
|
|
level=logging.INFO,
|
|
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
|
|
)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# 初始化FastAPI应用
|
|
app = FastAPI(
|
|
title="图像识别服务",
|
|
description="提供图像识别功能的AI服务",
|
|
version="1.0.0"
|
|
)
|
|
|
|
# 初始化识别器
|
|
recognizer = ImageRecognizer()
|
|
|
|
# 定义请求模型
|
|
class PredictRequest(BaseModel):
|
|
input_data: list
|
|
params: dict = {}
|
|
|
|
# 定义响应模型
|
|
class PredictResponse(BaseModel):
|
|
predictions: list
|
|
status: str
|
|
|
|
@app.post("/predict", response_model=PredictResponse)
|
|
async def predict(request: PredictRequest):
|
|
"""算法预测接口"""
|
|
try:
|
|
logger.info(f"Received prediction request for {len(request.input_data)} images")
|
|
predictions = recognizer.recognize(request.input_data, request.params)
|
|
logger.info(f"Prediction completed: {predictions}")
|
|
return PredictResponse(
|
|
predictions=predictions,
|
|
status="success"
|
|
)
|
|
except Exception as e:
|
|
logger.error(f"Prediction error: {str(e)}")
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
@app.post("/predict/file")
|
|
async def predict_file(file: UploadFile = File(...)):
|
|
"""通过文件上传进行预测"""
|
|
try:
|
|
logger.info(f"Received file upload: {file.filename}")
|
|
|
|
# 读取文件内容
|
|
contents = await file.read()
|
|
|
|
# 转换为base64
|
|
image_base64 = base64.b64encode(contents).decode('utf-8')
|
|
|
|
# 调用识别器
|
|
predictions = recognizer.recognize([image_base64])
|
|
|
|
logger.info(f"File prediction completed: {predictions}")
|
|
return {
|
|
"predictions": predictions,
|
|
"status": "success",
|
|
"filename": file.filename
|
|
}
|
|
except Exception as e:
|
|
logger.error(f"File prediction error: {str(e)}")
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
@app.get("/health")
|
|
async def health_check():
|
|
"""健康检查接口"""
|
|
return {
|
|
"status": "healthy",
|
|
"service": "image-recognition",
|
|
"version": "1.0.0"
|
|
}
|
|
|
|
@app.get("/info")
|
|
async def service_info():
|
|
"""服务信息接口"""
|
|
return {
|
|
"name": "图像识别服务",
|
|
"description": "提供图像识别功能的AI服务",
|
|
"version": "1.0.0",
|
|
"endpoints": {
|
|
"/predict": "POST - 图像识别预测",
|
|
"/predict/file": "POST - 通过文件上传进行预测",
|
|
"/health": "GET - 健康检查",
|
|
"/info": "GET - 服务信息"
|
|
}
|
|
}
|
|
|
|
if __name__ == "__main__":
|
|
uvicorn.run(
|
|
"main:app",
|
|
host=settings.HOST,
|
|
port=settings.PORT,
|
|
reload=settings.DEBUG
|
|
)
|