95 lines
2.7 KiB
Python
95 lines
2.7 KiB
Python
import base64
|
||
import io
|
||
import logging
|
||
import os
|
||
import time
|
||
|
||
import numpy as np
|
||
from PIL import Image
|
||
from fastapi import FastAPI, Request, status
|
||
from fastapi.responses import RedirectResponse
|
||
|
||
# 导入你的DetectionTool类
|
||
from utils.detection_tool import DetectionTool # 请替换为实际模块路径
|
||
|
||
"""
|
||
pip install ultralytics
|
||
pip install fastapi
|
||
"""
|
||
|
||
# 基础配置
|
||
logging.basicConfig(
|
||
level=logging.INFO,
|
||
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
|
||
)
|
||
logger = logging.getLogger(__name__)
|
||
save_dir = os.getenv("SAVE_DIR", "save_dir")
|
||
|
||
# 初始化检测工具
|
||
detector = DetectionTool()
|
||
|
||
# 创建FastAPI实例
|
||
app = FastAPI()
|
||
|
||
|
||
def image_to_numpy(image_source):
|
||
"""将图片源转换为numpy数组"""
|
||
if isinstance(image_source, str):
|
||
# 处理文件路径或URL
|
||
if image_source.startswith(("http://", "https://")):
|
||
import requests
|
||
response = requests.get(image_source)
|
||
if response.status_code != 200:
|
||
raise ValueError(f"无法获取图片，状态码: {response.status_code}")
|
||
image = Image.open(io.BytesIO(response.content))
|
||
else:
|
||
# 本地文件路径
|
||
image = Image.open(image_source)
|
||
elif isinstance(image_source, bytes):
|
||
# 二进制数据
|
||
image = Image.open(io.BytesIO(image_source))
|
||
else:
|
||
raise TypeError("不支持的图片源类型")
|
||
|
||
# 转换为RGB并转为numpy数组
|
||
return np.array(image.convert("RGB"))
|
||
|
||
|
||
@app.post(f"/v1/predict")
|
||
async def predict(request: Request):
|
||
"""核心预测接口：接收base64编码图片，返回检测结果"""
|
||
data = await request.json()
|
||
try:
|
||
# 解码base64图片为numpy数组
|
||
image_bytes = base64.b64decode(data["image"])
|
||
img_np = image_to_numpy(image_bytes)
|
||
|
||
# 使用你的检测工具获取结果
|
||
image_detected, result = detector.detect_image(img_np, image_format="pillow")
|
||
|
||
# 处理保存逻辑（如果需要）
|
||
if save_dir is not None:
|
||
os.makedirs(save_dir, exist_ok=True)
|
||
# 可根据需要保存输入图片和结果
|
||
timestamp = int(time.time())
|
||
img_save_path = os.path.join(save_dir, f"input_{timestamp}.jpg")
|
||
image_detected.save(img_save_path)
|
||
|
||
return result
|
||
|
||
except Exception as e:
|
||
logger.error(f"预测出错: {str(e)}")
|
||
return str(e)
|
||
|
||
|
||
@app.get("/")
|
||
async def index():
|
||
"""根路径重定向到API文档"""
|
||
return RedirectResponse(url="/docs", status_code=status.HTTP_302_FOUND)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
import uvicorn
|
||
|
||
uvicorn.run(app, host="0.0.0.0", port=8080)
|