139 lines
3.8 KiB
Python
139 lines
3.8 KiB
Python
from typing import Optional
|
||
from pathlib import Path
|
||
import tempfile
|
||
|
||
from fastapi import FastAPI, UploadFile, File, Query, HTTPException
|
||
from fastapi.middleware.cors import CORSMiddleware
|
||
|
||
from app.schemas import PredictResponse
|
||
from app.utils import decode_image_from_bytes
|
||
from app.pipeline import BirdCascadePipeline
|
||
from app.image_pipeline import save_image_outputs
|
||
from app.video_pipeline import process_video_file
|
||
|
||
|
||
app = FastAPI(
|
||
title="Bird Cascade Service",
|
||
description="YOLOv5 detection + target crop + ConvNeXt fine-grained bird classification",
|
||
version="1.0.0",
|
||
)
|
||
|
||
app.add_middleware(
|
||
CORSMiddleware,
|
||
allow_origins=["*"],
|
||
allow_credentials=True,
|
||
allow_methods=["*"],
|
||
allow_headers=["*"],
|
||
)
|
||
|
||
pipeline: Optional[BirdCascadePipeline] = None
|
||
startup_error: Optional[str] = None
|
||
|
||
|
||
@app.on_event("startup")
|
||
def load_models():
|
||
global pipeline, startup_error
|
||
try:
|
||
pipeline = BirdCascadePipeline()
|
||
startup_error = None
|
||
except Exception as e:
|
||
pipeline = None
|
||
startup_error = str(e)
|
||
|
||
|
||
@app.get("/health")
|
||
def health():
|
||
return {
|
||
"status": "ok" if pipeline is not None else "error",
|
||
"det_model_loaded": pipeline is not None,
|
||
"cls_model_loaded": pipeline is not None,
|
||
"error": startup_error,
|
||
}
|
||
|
||
|
||
@app.post("/predict", response_model=PredictResponse)
|
||
async def predict(
|
||
file: UploadFile = File(...),
|
||
return_image: bool = Query(False, description="是否在返回 JSON 中包含 base64 图片"),
|
||
det_conf: Optional[float] = Query(None, ge=0.0, le=1.0, description="YOLO 检测置信度阈值"),
|
||
iou_thres: Optional[float] = Query(None, ge=0.0, le=1.0, description="NMS IoU 阈值"),
|
||
):
|
||
"""
|
||
图片识别接口。
|
||
|
||
无论 return_image 是否为 true，都会自动保存：
|
||
1. 打框图片；
|
||
2. JSON；
|
||
3. XML。
|
||
|
||
保存目录：
|
||
runs/images/
|
||
"""
|
||
if pipeline is None:
|
||
raise HTTPException(status_code=500, detail=f"模型未加载成功: {startup_error}")
|
||
|
||
try:
|
||
data = await file.read()
|
||
img_bgr = decode_image_from_bytes(data)
|
||
|
||
output = pipeline.predict(
|
||
img_bgr,
|
||
return_image=return_image,
|
||
det_conf=det_conf,
|
||
iou_thres=iou_thres,
|
||
)
|
||
|
||
saved_outputs = save_image_outputs(
|
||
img_bgr=img_bgr,
|
||
original_filename=file.filename or "input.jpg",
|
||
predict_output=output,
|
||
)
|
||
|
||
output["saved_outputs"] = saved_outputs
|
||
return output
|
||
|
||
except Exception as e:
|
||
raise HTTPException(status_code=400, detail=str(e))
|
||
|
||
|
||
@app.post("/predict_video")
|
||
async def predict_video(
|
||
file: UploadFile = File(...),
|
||
sample_interval: Optional[int] = Query(None, ge=1, description="视频抽帧间隔，例如 2 表示每 2 帧处理 1 帧"),
|
||
save_annotated_video: bool = Query(True, description="是否保存画框后视频"),
|
||
):
|
||
"""
|
||
视频识别接口。
|
||
|
||
自动保存：
|
||
1. 打框视频；
|
||
2. JSON；
|
||
3. XML；
|
||
4. 关键帧截图；
|
||
5. 原始视频副本。
|
||
|
||
保存目录：
|
||
runs/video/
|
||
"""
|
||
if pipeline is None:
|
||
raise HTTPException(status_code=500, detail=f"模型未加载成功: {startup_error}")
|
||
|
||
try:
|
||
suffix = Path(file.filename or "input.mp4").suffix or ".mp4"
|
||
|
||
with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
|
||
tmp.write(await file.read())
|
||
tmp_path = tmp.name
|
||
|
||
summary = process_video_file(
|
||
video_path=tmp_path,
|
||
pipeline=pipeline,
|
||
original_filename=file.filename or "input.mp4",
|
||
sample_interval=sample_interval,
|
||
save_annotated_video=save_annotated_video,
|
||
)
|
||
|
||
return summary
|
||
|
||
except Exception as e:
|
||
raise HTTPException(status_code=400, detail=str(e)) |