initial commit

app/main.py

@@ -0,0 +1,139 @@
+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))