wangguangyuan
/
yanshou
1
0
Fork 0
Code Issues 0 Pull Requests 0 Projects 0 Releases 0 Wiki Activity
延寿水库
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
3 Commits
1 Branch
29 MiB
 
 
Branch: master
Branches Tags
${ item.name }
Create branch ${ searchTerm }
from 'master'
${ noResults }
yanshou/app/main.py

866 lines
33 KiB
Raw Permalink Blame History

import json
import os
import tempfile
from typing import List, Optional
import cv2
import numpy as np
from fastapi import FastAPI, File, HTTPException, Path, UploadFile, Form
from pydantic import BaseModel, Field
from celery.result import AsyncResult
from starlette.responses import FileResponse
from app.model_registry import MODEL_MAP, SCENE_INFO, iter_models, load_model
from app.tasks import celery_app, analyze_video_stream, analyze_video_file
# 共享上传目录（API 和 Worker 容器通过 volume 共享）
UPLOAD_DIR = os.getenv("UPLOAD_DIR", "/tmp")
UPLOAD_DIR = os.getenv("UPLOAD_DIR", "/tmp")
from app.tasks import celery_app, analyze_video_stream, analyze_video_file
AUTO_SCENE_TASK_TYPE = "auto_scene"
SCENE_MARKDOWN = "\n".join(
f"- `{model_type}`: {scene['name']}，{scene['description']}"
for model_type, scene in SCENE_INFO.items()
)
SUPPORTED_MODEL_TYPES = ", ".join(MODEL_MAP.keys())
WEBHOOK_PAYLOAD_EXAMPLE = {
"event": "RESERVOIR_ALARM",
"scene": "floating",
"stream_url": "rtsp://example.com/live",
"details": [
{
"label": "floating_object",
"confidence": 0.93,
"bbox": [102.4, 55.2, 260.8, 188.6],
}
],
"msg": "detected floating event",
}
WEBHOOK_CALLBACK_SCHEMA = {
"type": "object",
"required": ["event", "scene", "stream_url", "details", "msg"],
"properties": {
"event": {
"type": "string",
"description": "事件类型，当前固定为 `RESERVOIR_ALARM`。",
"example": "RESERVOIR_ALARM",
},
"scene": {
"type": "string",
"description": "触发识别的场景类型，例如 `floating`、`gate`、`shore_garbage`、`water_gauge`。",
"example": "floating",
},
"stream_url": {
"type": "string",
"description": "原始视频流地址。",
"example": "rtsp://example.com/live",
},
"details": {
"type": "array",
"description": "识别结果列表。",
"items": {
"type": "object",
"required": ["label", "confidence", "bbox"],
"properties": {
"label": {
"type": "string",
"description": "识别类别名称。",
"example": "floating_object",
},
"confidence": {
"type": "number",
"format": "float",
"description": "识别置信度。",
"example": 0.93,
},
"bbox": {
"type": "array",
"description": "目标框坐标数组，格式为 `[x1, y1, x2, y2]`。",
"items": {"type": "number", "format": "float"},
"example": [102.4, 55.2, 260.8, 188.6],
},
},
},
},
"msg": {
"type": "string",
"description": "告警说明。",
"example": "detected floating event",
},
},
}
WEBHOOK_CALLBACK_OPENAPI = {
"onWebhookResult": {
"{$request.body#/webhook_url}": {
"post": {
"summary": "识别结果 webhook 回调",
"description": "当视频流识别到目标后，算法服务会向请求体中的 `webhook_url` 发送 HTTP POST 回调。",
"requestBody": {
"required": True,
"content": {
"application/json": {
"schema": WEBHOOK_CALLBACK_SCHEMA,
"example": WEBHOOK_PAYLOAD_EXAMPLE,
}
},
},
"responses": {
"200": {
"description": "业务方成功接收回调。"
}
},
}
}
}
}
LEGACY_STREAM_REQUEST_EXAMPLE = {
"source_url": "rtsp://admin:password@192.168.1.100:554/h264/ch1/main/av_stream",
"webhook_url": "https://example.com/api/reservoir/webhook",
"scene_type": "floating",
"roi": [[0, 0], [1920, 0], [1920, 1080], [0, 1080]],
}
STREAM_REQUEST_EXAMPLE = {
"source_url": "rtsp://admin:password@192.168.1.100:554/h264/ch1/main/av_stream",
"webhook_url": "https://example.com/api/reservoir/webhook",
"roi": [[0, 0], [1920, 0], [1920, 1080], [0, 1080]],
}
AUTO_STREAM_RESPONSE_EXAMPLE = {
"code": 200,
"msg": "task submitted",
"task_type": AUTO_SCENE_TASK_TYPE,
}
STREAM_RESPONSE_EXAMPLE = {
"code": 200,
"msg": "task submitted",
"task_type": "floating",
}
IMAGE_RESPONSE_EXAMPLE = {
"results": [
{
"name": "floating_object",
"class": 0,
"confidence": 0.93,
"box": {
"x1": 102.4,
"y1": 55.2,
"x2": 260.8,
"y2": 188.6,
},
}
]
}
AUTO_IMAGE_RESPONSE_EXAMPLE = {
"task_type": AUTO_SCENE_TASK_TYPE,
"results": [
{
"scene": "floating",
"name": "floating_object",
"class": 0,
"confidence": 0.93,
"box": {
"x1": 102.4,
"y1": 55.2,
"x2": 260.8,
"y2": 188.6,
},
}
],
}
WEBHOOK_FIELD_DESCRIPTION = (
"识别结果回调地址，算法服务会以 HTTP POST 方式回调该地址。"
"回调 JSON 参数说明：`event` 为事件类型，`scene` 为场景类型，"
"`stream_url` 为视频流地址，`details` 为识别结果列表，`msg` 为告警说明。"
)
WEBHOOK_CALLBACK_PARAMETER_DESCRIPTION = (
"回调参数中文说明：\n"
"- `event`：事件类型，当前固定为告警事件 `RESERVOIR_ALARM`。\n"
"- `scene`：本次识别命中的场景类型，例如 `floating`、`gate`、`shore_garbage`、`water_gauge`。\n"
"- `stream_url`：触发识别的视频流地址，即调用接口时传入的原始流地址。\n"
"- `details`：识别结果列表。\n"
"- `details[].label`：识别出的目标类别名称。\n"
"- `details[].confidence`：识别结果的置信度。\n"
"- `details[].bbox`：目标框坐标，格式为 `[x1, y1, x2, y2]`。\n"
"- `msg`：告警说明文本，用于描述当前识别事件。"
)
APP_DESCRIPTION = f"""
延寿水库智能识别服务，基于 YOLO 模型提供图像识别和视频流识别能力。
支持场景：
{SCENE_MARKDOWN}
推荐接口：
- `POST /v1/predict/image`：自动场景图片识别
- `POST /v1/predict/stream`：自动场景视频流识别
- `POST /v1/predict/video-file`：上传MP4文件识别
兼容接口：
- `POST /{{model_type}}/predict/image`：按指定场景识别图片
- `POST /{{model_type}}/predict/stream`：按指定场景识别视频流
- `POST /v1/task/stream`：旧版视频流提交方式
"""
app = FastAPI(
title="延寿水库智能识别服务",
description=APP_DESCRIPTION,
version="1.2.0",
openapi_tags=[
{"name": "图像识别", "description": "支持自动场景识别，也支持按指定场景模型执行图片识别。推荐优先使用 `/v1/predict/image`。"},
{"name": "视频流识别", "description": "支持自动场景识别，也支持按指定场景模型提交视频流异步识别任务。推荐优先使用 `/v1/predict/stream`。"},
{"name": "视频文件识别", "description": "支持上传MP4文件进行离线识别分析。推荐优先使用 `/v1/predict/video-file`。"},
{"name": "任务管理", "description": "查询异步任务状态。"},
{"name": "兼容接口", "description": "兼容旧调用方式的接口，建议逐步迁移到自动场景识别接口。"},
],
)
class TaskRequest(BaseModel):
source_url: str = Field(
...,
description="视频流地址，支持 RTSP、本地视频文件路径等 OpenCV 可读取的地址。",
example="rtsp://admin:password@192.168.1.100:554/h264/ch1/main/av_stream",
)
webhook_url: str = Field(
...,
description=WEBHOOK_FIELD_DESCRIPTION,
example="https://example.com/api/reservoir/webhook",
)
scene_type: Optional[str] = Field(
default=None,
description=f"兼容旧字段。传值时按指定场景识别；不传时自动遍历全部场景模型。可选值：{SUPPORTED_MODEL_TYPES}",
example="floating",
)
roi: Optional[List[List[int]]] = Field(
default=None,
description="可选电子围栏，多边形坐标数组，格式为 [[x1, y1], [x2, y2], ...]；不传时默认检测整幅画面。",
example=[[0, 0], [1920, 0], [1920, 1080], [0, 1080]],
)
class StreamPredictRequest(BaseModel):
source_url: str = Field(
...,
description="视频流地址，支持 RTSP、本地视频文件路径等 OpenCV 可读取的地址。",
example="rtsp://admin:password@192.168.1.100:554/h264/ch1/main/av_stream",
)
webhook_url: str = Field(
...,
description=WEBHOOK_FIELD_DESCRIPTION,
example="https://example.com/api/reservoir/webhook",
)
roi: Optional[List[List[int]]] = Field(
default=None,
description="可选电子围栏，多边形坐标数组，格式为 [[x1, y1], [x2, y2], ...]；不传时默认检测整幅画面。",
example=[[0, 0], [1920, 0], [1920, 1080], [0, 1080]],
)
class VideoFileRequest(BaseModel):
webhook_url: str = Field(
...,
description="识别结果回调地址",
example="https://example.com/api/reservoir/webhook",
)
roi: Optional[List[List[int]]] = Field(
default=None,
description="可选电子围栏",
example=[[0, 0], [1920, 0], [1920, 1080], [0, 1080]],
)
scene_type: Optional[str] = Field(
default=None,
description=f"指定场景类型，可选值：{SUPPORTED_MODEL_TYPES}",
)
class StreamTaskResponse(BaseModel):
code: int = Field(..., description="业务状态码，`200` 表示任务提交成功。", example=200)
msg: str = Field(..., description="接口返回说明信息。", example="task submitted")
task_type: str = Field(
...,
description=f"实际提交的任务类型。自动识别时为 `{AUTO_SCENE_TASK_TYPE}`，指定场景时为对应模型类型。",
example=AUTO_SCENE_TASK_TYPE,
)
class VideoFileTaskResponse(BaseModel):
code: int = Field(..., description="业务状态码，`200` 表示任务提交成功。", example=200)
msg: str = Field(..., description="接口返回说明信息。", example="video file task submitted")
task_id: str = Field(..., description="任务ID，用于查询任务状态。", example="abc123def456")
task_type: str = Field(
...,
description=f"实际提交的任务类型。自动识别时为 `{AUTO_SCENE_TASK_TYPE}`，指定场景时为对应模型类型。",
example=AUTO_SCENE_TASK_TYPE,
)
class DetectionBox(BaseModel):
x1: float = Field(..., description="检测框左上角 X 坐标。", example=102.4)
y1: float = Field(..., description="检测框左上角 Y 坐标。", example=55.2)
x2: float = Field(..., description="检测框右下角 X 坐标。", example=260.8)
y2: float = Field(..., description="检测框右下角 Y 坐标。", example=188.6)
class ImageDetectionItem(BaseModel):
name: str = Field(..., description="检测类别名称。", example="floating_object")
class_id: int = Field(..., alias="class", description="检测类别 ID。", example=0)
confidence: float = Field(..., description="检测置信度。", example=0.93)
box: DetectionBox = Field(..., description="目标框坐标。")
class AutoImageDetectionItem(ImageDetectionItem):
scene: str = Field(..., description="命中的场景模型类型。", example="floating")
class ImagePredictResponse(BaseModel):
results: List[ImageDetectionItem] = Field(
...,
description="识别结果列表。每个元素为模型输出的目标信息，通常包含类别名、类别 ID、置信度和检测框坐标。",
example=IMAGE_RESPONSE_EXAMPLE["results"],
)
class AutoImagePredictResponse(BaseModel):
task_type: str = Field(
...,
description=f"任务类型，固定为 `{AUTO_SCENE_TASK_TYPE}`。",
example=AUTO_SCENE_TASK_TYPE,
)
results: List[AutoImageDetectionItem] = Field(
...,
description="自动场景识别结果列表。每个元素额外包含 `scene` 字段，表示命中的场景模型。",
example=AUTO_IMAGE_RESPONSE_EXAMPLE["results"],
)
class WebhookDetectionItem(BaseModel):
label: str = Field(..., description="识别类别名称。", example="floating_object")
confidence: float = Field(..., description="识别置信度。", example=0.93)
bbox: List[float] = Field(..., description="目标框坐标数组，格式为 `[x1, y1, x2, y2]`。", example=[102.4, 55.2, 260.8, 188.6])
class WebhookPayload(BaseModel):
event: str = Field(..., description="事件类型。", example="RESERVOIR_ALARM")
scene: str = Field(..., description="触发识别的场景类型。", example="floating")
stream_url: str = Field(..., description="视频流地址。", example="rtsp://example.com/live")
details: List[WebhookDetectionItem] = Field(..., description="识别结果列表。")
msg: str = Field(..., description="告警说明。", example="detected floating event")
video_url: Optional[str] = Field(None, description="处理后的视频下载链接") # 新增
def validate_model_type(model_type: str) -> None:
if model_type not in MODEL_MAP:
raise HTTPException(
status_code=400,
detail=f"Unsupported model type: {model_type}. Supported types: {SUPPORTED_MODEL_TYPES}",
)
def decode_image(contents: bytes):
nparr = np.frombuffer(contents, np.uint8)
img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
if img is None:
raise HTTPException(status_code=400, detail="Invalid image file")
return img
def run_single_model_image_detection(img, model_type: str) -> list:
model = load_model(model_type)
results = model.predict(img, conf=0.6, verbose=False)
return json.loads(results[0].to_json())
def run_auto_image_detection(img) -> dict:
detections = []
for current_model_type, model in iter_models():
results = model.predict(img, conf=0.6, verbose=False)
for item in json.loads(results[0].to_json()):
item["scene"] = current_model_type
detections.append(item)
return {"task_type": AUTO_SCENE_TASK_TYPE, "results": detections}
def submit_stream_task(
source_url: str,
webhook_url: str,
roi: Optional[List[List[int]]],
model_type: Optional[str] = None,
):
if model_type is not None:
validate_model_type(model_type)
task = analyze_video_stream.delay(source_url, webhook_url, model_type, roi)
return {"code": 200, "msg": "task submitted", "task_type": model_type or AUTO_SCENE_TASK_TYPE, "task_id": task.id}
def submit_video_file_task(
video_path: str,
webhook_url: str,
roi: Optional[List[List[int]]],
model_type: Optional[str] = None,
original_filename: str = None,
):
"""提交视频文件识别任务"""
if model_type is not None:
validate_model_type(model_type)
task = analyze_video_file.delay(video_path, webhook_url, model_type, roi, original_filename)
return {
"code": 200,
"msg": "video file task submitted",
"task_id": task.id,
"task_type": model_type or AUTO_SCENE_TASK_TYPE
}
@app.post(
"/v1/task/stream",
tags=["兼容接口", "视频流识别"],
summary="兼容模式提交视频流识别任务",
description=(
"兼容旧调用方式提交视频流识别任务。\n\n"
"入参说明：\n"
"- `source_url`：视频流地址。\n"
"- `webhook_url`：识别结果回调地址。\n"
f"- `scene_type`：兼容旧字段，传值时按指定场景识别，不传时自动识别全部场景；可选值：{SUPPORTED_MODEL_TYPES}。\n"
"- `roi`：可选电子围栏。\n\n"
"模拟入参示例：\n"
f"```json\n{json.dumps(LEGACY_STREAM_REQUEST_EXAMPLE, ensure_ascii=False, indent=2)}\n```\n\n"
"出参说明：\n"
"- `code`：业务状态码，200 表示提交成功。\n"
"- `msg`：返回说明信息。\n"
f"- `task_type`：实际提交的任务类型，自动识别时为 `{AUTO_SCENE_TASK_TYPE}`。\n\n"
"模拟出参示例：\n"
f"```json\n{json.dumps(AUTO_STREAM_RESPONSE_EXAMPLE, ensure_ascii=False, indent=2)}\n```\n\n"
"说明：\n"
"- 建议新接入统一改用 `POST /v1/predict/stream`。\n"
"- 自动识别模式下，worker 会遍历全部已注册场景模型。\n"
"- Swagger 文档中的 Callbacks 区域可查看 `webhook_url` 的回调参数结构。\n\n"
f"{WEBHOOK_CALLBACK_PARAMETER_DESCRIPTION}\n"
),
response_model=StreamTaskResponse,
responses={
200: {
"description": "任务提交成功",
"content": {"application/json": {"example": AUTO_STREAM_RESPONSE_EXAMPLE}},
},
400: {"description": "请求参数错误或场景类型不支持。"},
},
deprecated=True,
openapi_extra={
"requestBody": {
"content": {
"application/json": {
"example": LEGACY_STREAM_REQUEST_EXAMPLE
}
}
},
"callbacks": WEBHOOK_CALLBACK_OPENAPI,
},
)
async def start_stream_task(req: TaskRequest):
return submit_stream_task(req.source_url, req.webhook_url, req.roi, req.scene_type)
@app.post(
"/v1/predict/stream",
tags=["视频流识别"],
summary="自动场景视频流识别任务",
description=(
"自动遍历所有已注册场景模型，对视频流执行异步识别任务。\n\n"
"入参说明：\n"
"- `source_url`：视频流地址。\n"
"- `webhook_url`：识别结果回调地址。\n"
"- `roi`：可选电子围栏。\n\n"
"模拟入参示例：\n"
f"```json\n{json.dumps(STREAM_REQUEST_EXAMPLE, ensure_ascii=False, indent=2)}\n```\n\n"
"出参说明：\n"
"- `code`：业务状态码，200 表示提交成功。\n"
"- `msg`：返回说明信息。\n"
f"- `task_type`：固定返回 `{AUTO_SCENE_TASK_TYPE}`。\n\n"
"模拟出参示例：\n"
f"```json\n{json.dumps(AUTO_STREAM_RESPONSE_EXAMPLE, ensure_ascii=False, indent=2)}\n```\n\n"
"检测到目标后会向 `webhook_url` 发送 POST 回调，回调数据示例：\n"
f"```json\n{json.dumps(WEBHOOK_PAYLOAD_EXAMPLE, ensure_ascii=False, indent=2)}\n```\n\n"
"回调说明：\n"
"- 自动识别模式下，如果同一帧命中多个场景，会按场景分别回调。\n"
"- `scene` 字段表示当前这条 webhook 对应的场景模型。\n"
"- Swagger 文档中的 Callbacks 区域可查看完整回调参数结构。\n\n"
f"{WEBHOOK_CALLBACK_PARAMETER_DESCRIPTION}\n"
),
response_model=StreamTaskResponse,
responses={
200: {
"description": "任务提交成功",
"content": {"application/json": {"example": AUTO_STREAM_RESPONSE_EXAMPLE}},
},
400: {"description": "请求参数错误。"},
},
openapi_extra={
"requestBody": {
"content": {
"application/json": {
"example": STREAM_REQUEST_EXAMPLE
}
}
},
"callbacks": WEBHOOK_CALLBACK_OPENAPI,
},
)
async def predict_stream_v1(req: StreamPredictRequest):
return submit_stream_task(req.source_url, req.webhook_url, req.roi, None)
@app.post(
"/{model_type}/predict/stream",
tags=["视频流识别"],
summary="按场景提交视频流识别任务",
description=(
"根据路径中的 `model_type` 选择对应场景模型，异步提交视频流识别任务。\n\n"
"入参说明：\n"
"- `model_type`：路径参数，场景模型类型。\n"
"- `source_url`：视频流地址。\n"
"- `webhook_url`：识别结果回调地址。\n"
"- `roi`：可选电子围栏。\n\n"
"模拟入参示例：\n"
f"```json\n{json.dumps(STREAM_REQUEST_EXAMPLE, ensure_ascii=False, indent=2)}\n```\n\n"
"出参说明：\n"
"- `code`：业务状态码，200 表示提交成功。\n"
"- `msg`：返回说明信息。\n"
"- `task_type`：实际提交的场景类型。\n\n"
"模拟出参示例：\n"
f"```json\n{json.dumps(STREAM_RESPONSE_EXAMPLE, ensure_ascii=False, indent=2)}\n```\n\n"
"检测到目标后会向 `webhook_url` 发送 POST 回调，回调数据示例：\n"
f"```json\n{json.dumps(WEBHOOK_PAYLOAD_EXAMPLE, ensure_ascii=False, indent=2)}\n```\n\n"
"回调说明：\n"
"- 指定场景模式下，只会使用路径里的 `model_type` 对应模型。\n"
"- webhook 的 `scene` 字段与路径中的 `model_type` 一致。\n"
"- Swagger 文档中的 Callbacks 区域可查看完整回调参数结构。\n\n"
f"{WEBHOOK_CALLBACK_PARAMETER_DESCRIPTION}\n"
),
response_model=StreamTaskResponse,
responses={
200: {
"description": "任务提交成功",
"content": {"application/json": {"example": STREAM_RESPONSE_EXAMPLE}},
},
400: {"description": "场景类型不支持或请求参数错误。"},
},
openapi_extra={
"requestBody": {
"content": {
"application/json": {
"example": STREAM_REQUEST_EXAMPLE
}
}
},
"callbacks": WEBHOOK_CALLBACK_OPENAPI,
},
)
async def predict_stream(
model_type: str = Path(..., description=f"场景模型类型。可选值：{SUPPORTED_MODEL_TYPES}"),
req: StreamPredictRequest = ...,
):
return submit_stream_task(req.source_url, req.webhook_url, req.roi, model_type)
@app.post(
"/v1/predict/video-file",
tags=["视频文件识别"],
summary="上传MP4文件进行识别",
description=(
"支持上传本地MP4文件进行视频识别分析。\n\n"
"入参说明：\n"
"- `file`：上传的MP4视频文件\n"
"- `webhook_url`：识别结果回调地址\n"
"- `roi`：可选电子围栏（JSON字符串格式）\n"
"- `scene_type`：可选指定场景类型，不传则自动识别\n\n"
"模拟入参示例：\n"
"```bash\n"
"curl -X POST http://localhost:8000/v1/predict/video-file \\\n"
" -H \"accept: application/json\" \\\n"
" -F \"file=@test.mp4;type=video/mp4\" \\\n"
" -F \"webhook_url=https://example.com/api/webhook\" \\\n"
" -F \"roi=[[0,0],[1920,0],[1920,1080],[0,1080]]\" \\\n"
" -F \"scene_type=floating\"\n"
"```\n\n"
"出参说明：\n"
"- `code`：业务状态码，200表示提交成功\n"
"- `msg`：返回说明信息\n"
"- `task_id`：任务ID，用于查询任务状态\n"
"- `task_type`：实际提交的任务类型\n\n"
"模拟出参示例：\n"
"```json\n"
"{\n"
" \"code\": 200,\n"
" \"msg\": \"video file task submitted\",\n"
" \"task_id\": \"abc123def456\",\n"
" \"task_type\": \"auto_scene\"\n"
"}\n"
"```\n\n"
"检测完成或检测到目标后会向 `webhook_url` 发送回调。\n"
"文件大小限制：500MB\n"
),
response_model=VideoFileTaskResponse,
responses={
200: {"description": "任务提交成功"},
400: {"description": "文件格式错误、大小超限或参数错误"},
500: {"description": "服务器内部错误"},
},
)
async def predict_video_file(
file: UploadFile = File(..., description="MP4视频文件", media_type="video/mp4"),
webhook_url: str = Form(..., description="识别结果回调地址"),
roi: Optional[str] = Form(default="[[0,0],[1920,0],[1920,1080],[0,1080]]", description="电子围栏JSON字符串"),
scene_type: Optional[str] = Form(default=None, description=f"场景类型，可选值：{SUPPORTED_MODEL_TYPES}"),
):
# 验证文件格式
if not file.filename.endswith(('.mp4', '.MP4')):
raise HTTPException(status_code=400, detail="Only MP4 files are supported")
# 验证文件大小（限制500MB）
file_size = 0
temp_file_path = None
try:
# 创建临时文件保存上传的视频（存入共享目录，确保 Worker 容器也能读取）
os.makedirs(UPLOAD_DIR, exist_ok=True)
with tempfile.NamedTemporaryFile(delete=False, suffix='.mp4', dir=UPLOAD_DIR) as temp_file:
# 分块读取并写入
chunk_size = 1024 * 1024 # 1MB chunks
while True:
chunk = await file.read(chunk_size)
if not chunk:
break
temp_file.write(chunk)
file_size += len(chunk)
# 限制文件大小为500MB
if file_size > 500 * 1024 * 1024:
raise HTTPException(status_code=400, detail="File size exceeds 500MB limit")
temp_file_path = temp_file.name
# 解析ROI参数
roi_list = None
if roi:
try:
roi_list = json.loads(roi)
if not isinstance(roi_list, list):
raise ValueError("ROI must be a list")
except json.JSONDecodeError:
raise HTTPException(status_code=400, detail="Invalid ROI format, must be valid JSON")
# 提交任务
result = submit_video_file_task(
video_path=temp_file_path,
webhook_url=webhook_url,
roi=roi_list,
model_type=scene_type,
original_filename=file.filename
)
return result
except HTTPException:
# 清理临时文件
if temp_file_path and os.path.exists(temp_file_path):
os.unlink(temp_file_path)
raise
except Exception as exc:
if temp_file_path and os.path.exists(temp_file_path):
os.unlink(temp_file_path)
raise HTTPException(status_code=500, detail=f"Failed to process video file: {str(exc)}") from exc
@app.get(
"/v1/task/status/{task_id}",
tags=["任务管理"],
summary="查询异步任务状态",
description="通过任务ID查询视频识别任务的执行状态和进度",
)
async def get_task_status(
task_id: str = Path(..., description="任务ID，从提交接口返回的task_id")
):
"""查询任务状态"""
try:
task = AsyncResult(task_id, app=celery_app)
if task.state == 'PENDING':
response = {
"state": "PENDING",
"progress": 0,
"message": "Task is waiting to be processed"
}
elif task.state == 'PROCESSING':
response = {
"state": "PROCESSING",
"progress": task.info.get('progress', 0) if task.info else 0,
"message": task.info.get('message', 'Processing video') if task.info else 'Processing',
"frame": task.info.get('frame', 0) if task.info else 0,
"total_frames": task.info.get('total_frames', 0) if task.info else 0
}
elif task.state == 'SUCCESS':
response = {
"state": "SUCCESS",
"result": task.result,
"message": "Task completed successfully"
}
elif task.state == 'FAILURE':
response = {
"state": "FAILURE",
"error": str(task.info),
"message": "Task failed"
}
else:
response = {
"state": task.state,
"message": "Unknown task state"
}
return response
except Exception as exc:
raise HTTPException(
status_code=500,
detail=f"Failed to query task status: {str(exc)}",
) from exc
@app.post(
"/v1/predict/image",
tags=["图像识别"],
summary="自动场景图片识别",
description=(
"不再需要指定 `model_type`，接口会自动遍历所有已注册场景模型，对上传图片执行场景识别。\n\n"
"入参说明：\n"
"- `file`：上传的待识别图片文件。\n\n"
"模拟入参示例：\n"
"```bash\n"
"curl -X POST http://localhost:8000/v1/predict/image \\\n"
" -H \"accept: application/json\" \\\n"
" -H \"Content-Type: multipart/form-data\" \\\n"
" -F \"file=@test.jpg;type=image/jpeg\"\n"
"```\n\n"
"出参说明：\n"
f"- `task_type`：固定返回 `{AUTO_SCENE_TASK_TYPE}`。\n"
"- `results`：识别结果列表。\n"
"- `results[].scene`：命中的场景模型类型。\n"
"- `results[].name`：类别名称。\n"
"- `results[].class`：类别 ID。\n"
"- `results[].confidence`：置信度。\n"
"- `results[].box`：目标框坐标。\n\n"
"模拟出参示例：\n"
f"```json\n{json.dumps(AUTO_IMAGE_RESPONSE_EXAMPLE, ensure_ascii=False, indent=2)}\n```\n\n"
"说明：\n"
"- 该接口会遍历全部已注册模型并合并命中结果。\n"
"- 若没有检测到目标，`results` 返回空数组。\n"
),
response_model=AutoImagePredictResponse,
responses={
200: {
"description": "图片自动场景识别成功",
"content": {"application/json": {"example": AUTO_IMAGE_RESPONSE_EXAMPLE}},
},
400: {"description": "上传文件不是有效图片。"},
},
)
async def predict_image_v1(
file: UploadFile = File(..., description="待识别图片文件，示例：file=@test.jpg"),
):
try:
contents = await file.read()
img = decode_image(contents)
return run_auto_image_detection(img)
except HTTPException:
raise
except Exception as exc:
raise HTTPException(status_code=500, detail=str(exc)) from exc
@app.post(
"/{model_type}/predict/image",
tags=["图像识别"],
summary="按场景执行图片识别",
description=(
"根据路径中的 `model_type` 选择对应场景模型，对上传图片执行识别。\n\n"
"入参说明：\n"
"- `model_type`：路径参数，场景模型类型。\n"
"- `file`：上传的待识别图片文件。\n\n"
"模拟入参示例：\n"
"```bash\n"
"curl -X POST http://localhost:8000/floating/predict/image \\\n"
" -H \"accept: application/json\" \\\n"
" -H \"Content-Type: multipart/form-data\" \\\n"
" -F \"file=@test.jpg;type=image/jpeg\"\n"
"```\n\n"
"出参说明：\n"
"- `results`：识别结果列表。\n"
"- `results[].name`：类别名称。\n"
"- `results[].class`：类别 ID。\n"
"- `results[].confidence`：置信度。\n"
"- `results[].box`：目标框坐标。\n\n"
"模拟出参示例：\n"
f"```json\n{json.dumps(IMAGE_RESPONSE_EXAMPLE, ensure_ascii=False, indent=2)}\n```\n\n"
"说明：\n"
"- 该接口只使用路径中的 `model_type` 对应模型。\n"
"- 若没有检测到目标，`results` 返回空数组。\n"
),
response_model=ImagePredictResponse,
responses={
200: {
"description": "图片识别成功",
"content": {"application/json": {"example": IMAGE_RESPONSE_EXAMPLE}},
},
400: {"description": "场景类型不支持，或上传文件不是有效图片。"},
},
)
async def predict_image(
model_type: str = Path(..., description=f"场景模型类型。可选值：{SUPPORTED_MODEL_TYPES}"),
file: UploadFile = File(..., description="待识别图片文件，示例：file=@test.jpg"),
):
validate_model_type(model_type)
try:
contents = await file.read()
img = decode_image(contents)
return {"results": run_single_model_image_detection(img, model_type)}
except HTTPException:
raise
except Exception as exc:
raise HTTPException(status_code=500, detail=str(exc)) from exc
@app.get("/v1/download/{filename}")
async def download_video(filename: str):
"""下载处理后的视频文件"""
video_dir = os.path.join(UPLOAD_DIR, "processed_videos")
print("video addr1",video_dir);
file_path = os.path.join(video_dir, filename)
if not os.path.exists(file_path):
raise HTTPException(status_code=404, detail="Video not found")
# 安全加固：防止路径穿越
if not os.path.realpath(file_path).startswith(os.path.realpath(video_dir)):
raise HTTPException(status_code=403, detail="Forbidden")
return FileResponse(file_path, media_type="video/mp4", filename=filename)
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=8000)