commit 0bafbee31f30b8b8c4800b28088dbae9c71471f8 Author: Ankle Date: Fri Aug 21 13:29:29 2026 +0800 initial commit diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 0000000..8e8ebcd --- /dev/null +++ b/.dockerignore @@ -0,0 +1,15 @@ +__pycache__/ +*.py[cod] +.pytest_cache/ +.git/ +.gitignore + +models/* +!models/.gitkeep +runs/ +input/* +!input/.gitkeep + +font_tmp/ +fonts-noto-cjk_*.deb +README？.md diff --git a/.git_backup_20260821 b/.git_backup_20260821 new file mode 160000 index 0000000..f07b49e --- /dev/null +++ b/.git_backup_20260821 @@ -0,0 +1 @@ +Subproject commit f07b49e0bb6f67bca805e2615d359cfb9024f862 diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000..0bb75f7 --- /dev/null +++ b/.gitattributes @@ -0,0 +1 @@ +*.onnx filter=lfs diff=lfs merge=lfs -text diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..2b16844 --- /dev/null +++ b/.gitignore @@ -0,0 +1,14 @@ +runs/ + +__pycache__/ +*.pyc +*.pyo +*.pyd +.pytest_cache/ + +.vscode/ +.idea/ + +*.log +*.tmp +.DS_Store diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..11f5919 --- /dev/null +++ b/Dockerfile @@ -0,0 +1,25 @@ +FROM python:3.10-slim + +ENV PYTHONDONTWRITEBYTECODE=1 +ENV PYTHONUNBUFFERED=1 + +WORKDIR /app + +# OpenCV 运行时依赖 +RUN apt-get update && apt-get install -y --no-install-recommends \ + libglib2.0-0 \ + libgl1 \ + libgomp1 \ + && rm -rf /var/lib/apt/lists/* + +COPY requirements.txt /app/requirements.txt +RUN pip install --no-cache-dir -r /app/requirements.txt + +COPY app /app/app +COPY assets/fonts /app/assets/fonts +COPY run_batch.sh /app/run_batch.sh +RUN mkdir -p /app/models /app/input /app/runs \ + && chmod +x /app/run_batch.sh + +ENTRYPOINT ["bash", "/app/run_batch.sh"] +CMD ["--input", "/app/input", "--output", "/app/runs"] diff --git a/README.md b/README.md new file mode 100644 index 0000000..23f2c01 --- /dev/null +++ b/README.md @@ -0,0 +1,300 @@ +# Bird Detect Cascade Batch + +这是一个 **YOLOv5 + ConvNeXt + 贝叶斯密度估计** 级联鸟类识别批处理程序。它通过 bash 命令处理图片、视频或文件夹，并输出标注结果、JSON、XML。 + +## 1. 最常用运行方式 + +处理单张图片、单个视频或一个文件夹： + +```bash +bash run_batch.sh --input /path/to/input +``` + +指定输出目录，并关闭密度估计： + +```bash +bash run_batch.sh --input /path/to/input --output /path/to/runs --disable-density +``` + +输入路径可以是图片、视频，也可以是包含图片/视频的文件夹。默认会递归处理子文件夹。 + +三个模型都可以单独关闭，例如只运行密度估计： + +```bash +bash run_batch.sh --input /path/to/input \ + --disable-detection \ + --disable-classification +``` + +## 2. 哪些参数可以直接写在 bash 指令里 + +下面这些参数都可以直接写在 `bash run_batch.sh ...` 后面，不需要改 Python 代码。 + +| bash 参数 | 默认值 | 含义 | +|---|---:|---| +| `--input PATH` | 无，必填 | 输入路径，可以是图片、视频或文件夹 | +| `--output PATH` | `./runs` | 输出根目录 | +| `--sample-interval INT` | `2` | 视频抽帧间隔，`2` 表示每 2 帧处理 1 帧 | +| `--recursive` | 开启 | 输入是文件夹时递归处理子目录 | +| `--no-recursive` | 关闭递归 | 只处理输入文件夹第一层文件 | +| `--enable-detection` | 开启 | 启用 YOLO 目标检测 | +| `--disable-detection` | 关闭目标检测 | 不加载、不运行 YOLO；检测结果为空 | +| `--enable-classification` | 开启 | 启用 ConvNeXt 细粒度识别 | +| `--disable-classification` | 关闭细粒度识别 | 不加载、不运行 ConvNeXt；保留 YOLO 检测框 | +| `--enable-density` | 开启 | 启用贝叶斯密度估计 | +| `--disable-density` | 关闭密度估计 | 不运行密度估计，只运行 YOLO + ConvNeXt | +| `--det-conf-thres FLOAT` | `0.25` | YOLO 第一层检测阈值，低于该值的框直接丢弃 | +| `--det-iou-thres FLOAT` | `0.45` | NMS IoU 阈值；越低，重叠框去得越狠 | +| `--det-second-conf-thres FLOAT` | `0.35` | YOLO 第二层过滤阈值，低于该值不会裁剪，也不会进入 ConvNeXt | +| `--min-box-area-ratio FLOAT` | `0.00002` | 过滤面积占比过小的噪声框 | +| `--max-box-area-ratio FLOAT` | `1.0` | 最大框面积占比；默认不按面积过滤大目标 | +| `--max-box-aspect-ratio FLOAT` | `6.0` | 过滤长宽比过大的细长框 | +| `--cls-show-conf-thres FLOAT` | `0.50` | 普通目标分类置信度低于该值时，不在结果图显示物种名 | +| `--large-box-area-ratio FLOAT` | `0.002` | 框面积占比超过该值时，认为是近景大目标 | +| `--large-box-cls-show-conf-thres FLOAT` | `0.25` | 大目标显示物种名的较低分类阈值 | +| `--show-large-box-name-always BOOL` | `true` | 大目标是否只要分类名非空就显示 | +| `--dense-scene-count-thres INT` | `20` | 密度估计数量达到该值时，认为是密集场景 | +| `--hide-cls-name-in-dense-scene BOOL` | `true` | 密集场景下是否隐藏所有物种名 | +| `--crop-padding-ratio FLOAT` | `0.08` | 送入 ConvNeXt 前，检测框向外扩的比例 | +| `--max-detections INT` | `50` | 单张图或单帧最多保留的检测框数量 | +| `--uniform-species-color BOOL` | `false` | `true` 时所有物种的检测框和标签统一使用同一种颜色 | + +示例：降低检测阈值、提高进入分类前的二次过滤阈值： + +```bash +bash run_batch.sh \ + --input /path/to/input \ + --det-conf-thres 0.15 \ + --det-second-conf-thres 0.45 +``` + +示例：密集场景仍然显示物种名： + +```bash +bash run_batch.sh \ + --input /path/to/input \ + --hide-cls-name-in-dense-scene false +``` + +示例：让所有物种的检测框和标签统一使用绿色： + +```bash +bash run_batch.sh \ + --input /path/to/input \ + --uniform-species-color true +``` + +## 3. 阈值怎么调 + +最关键的是这两个参数： + +- `--det-conf-thres`: YOLO 第一层阈值。低于这个值的框会在 YOLO 后处理阶段直接丢弃。 +- `--det-second-conf-thres`: YOLO 第二层阈值。低于这个值的框不会裁剪，也不会传给 ConvNeXt 分类。 + +常见调法： + +- 漏检多：降低 `--det-conf-thres`，例如 `0.25 -> 0.15`。 +- 误检多：提高 `--det-second-conf-thres`，例如 `0.35 -> 0.45`。 +- 小噪声框多：提高 `--min-box-area-ratio`。 +- 异常大框多：降低 `--max-box-area-ratio`；不要为了排除异常框把它设得过低，否则近景大鸟会被误删。 +- 细长误检框多：降低 `--max-box-aspect-ratio`。 +- 检测到了但不显示物种名：降低 `--cls-show-conf-thres`，或检查 `--hide-cls-name-in-dense-scene`。 +- 密集场景过早隐藏物种名：提高 `--dense-scene-count-thres`，或设置 `--hide-cls-name-in-dense-scene false`。 + +这些阈值现在都可以写在 bash 指令里，不需要去代码里调。 + +## 4. 三个模型的开关 + +三个模型默认都开启： + +```bash +bash run_batch.sh --input /path/to/input +``` + +关闭目标检测： + +```bash +bash run_batch.sh --input /path/to/input --disable-detection +``` + +关闭细粒度识别，只输出 YOLO 检测框： + +```bash +bash run_batch.sh --input /path/to/input --disable-classification +``` + +关闭密度估计： + +```bash +bash run_batch.sh --input /path/to/input --disable-density +``` + +也可以任意组合，例如只运行目标检测： + +```bash +bash run_batch.sh --input /path/to/input \ + --disable-classification \ + --disable-density +``` + +只运行密度估计： + +```bash +bash run_batch.sh --input /path/to/input \ + --disable-detection \ + --disable-classification +``` + +注意：细粒度识别依赖目标检测提供裁剪框，因此关闭目标检测后，ConvNeXt 不会加载或运行，即使显式传入 `--enable-classification`。密度估计与另外两个模型相互独立。 + +关闭某个模型后： + +- 对应 ONNX 模型不会加载，也不会执行推理； +- 关闭目标检测时，`results` 为空； +- 关闭细粒度识别时，每个检测结果的 `cls_class_id` 为 `-1`、`cls_class_name` 为空、`cls_confidence` 为 `0.0`； +- 关闭密度估计时，JSON 里的 `density` 为 `null`，结果图不显示密度热力图和密度估计数量。 + +颜色开关不影响模型推理和 JSON/XML，只影响结果图片或视频的可视化： + +- `--uniform-species-color false`：默认行为，每个物种使用各自的稳定颜色； +- `--uniform-species-color true`：所有物种的检测框和标签统一使用绿色。 + +## 5. 哪些配置用环境变量，不直接写 bash 参数 + +模型路径、输入尺寸、设备这类参数通常不频繁改，所以保留为环境变量。它们可以写在 bash 命令前面，也不需要改代码。 + +| 环境变量 | 默认值 | 含义 | +|---|---:|---| +| `ENABLE_DETECTION` | `true` | 是否启用 YOLO 目标检测 | +| `ENABLE_CLASSIFICATION` | `true` | 是否启用 ConvNeXt 细粒度识别 | +| `ENABLE_DENSITY` | `true` | 是否启用密度估计 | +| `DET_MODEL_PATH` | `models/YNDC_20251010_yolov5_optimize.onnx` | YOLOv5 检测模型 | +| `DET_CLASSES_PATH` | `models/YNDC_20251010_yolov5_optimize.txt` | YOLO 类别文件 | +| `CLS_MODEL_PATH` | `models/HHK_20260525_mmcls_optimize.onnx` | ConvNeXt 分类模型 | +| `CLS_CLASSES_PATH` | `models/HHK_20260525_mmcls_optimize.txt` | ConvNeXt 类别文件 | +| `DENSITY_MODEL_PATH` | `models/bird_density.onnx` | 密度估计模型 | +| `DET_IMG_SIZE` | `640` | YOLO 输入尺寸 | +| `CLS_IMG_SIZE` | `224` | ConvNeXt 输入尺寸 | +| `DENSITY_IMG_SIZE` | `0` | 密度模型输入尺寸；`0` 表示自动读取模型输入 | +| `TARGET_DET_CLASS_IDS` | `0` | 只保留哪些 YOLO 检测类别 | +| `DEVICE` | `cpu` | ONNX Runtime 运行设备，支持 `cpu` 或 `cuda`；CUDA 不可用时直接报错，不会静默回退 CPU | +| `UNIFORM_SPECIES_COLOR` | `false` | 是否让所有物种的检测框和标签使用同一种颜色 | + +示例： + +```bash +DEVICE=cpu TARGET_DET_CLASS_IDS=0 bash run_batch.sh --input /path/to/input +``` + +指定 GPU 0 运行： + +```bash +CUDA_VISIBLE_DEVICES=0 DEVICE=cuda bash run_batch.sh --input /path/to/input +``` + +使用 CUDA 前必须安装与服务器 CUDA/cuDNN 匹配的 `onnxruntime-gpu`。启动时程序会分别打印 YOLO、ConvNeXt、密度模型的实际 Provider；使用 GPU 时应看到 `CUDAExecutionProvider`。 + +## 6. 哪些情况才需要改代码 + +一般调阈值、开关三个模型、切换统一颜色、改模型路径都不需要改代码。 + +只有下面这些情况通常需要改代码： + +- YOLO ONNX 输出格式不是 `[1, num_boxes, 5 + num_classes]` 或 `[num_boxes, 6]`，需要改 `app/inference/yolo_onnx.py`。 +- ConvNeXt 训练时没有使用 ImageNet mean/std，需要改 `app/inference/convnext_onnx.py` 的归一化参数。 +- 密度估计模型的前处理或输出含义和当前实现不一致，需要改 `app/inference/density_onnx.py`。 +- 想改变结果 JSON/XML 字段结构，需要改 `app/batch_runner.py` 或 `app/utils.py`。 + +## 7. 输出结果 + +每次运行会生成一个任务目录： + +```text +runs/<时间戳>_<输入名>/ +├── result/ # 标注后的图片或视频 +├── json/ # 结构化结果 +├── xml/ # XML 标注结果 +└── index.json # 本次批处理汇总 +``` + +图片输出： + +- `result/<原文件名>.jpg` +- `json/<原文件名>.json` +- `xml/<原文件名>.xml` + +视频输出： + +- `result/<原文件名>.mp4` +- `json/<原文件名>.json` +- `xml/<原文件名>.xml` + +JSON 中每个目标包含： + +```json +{ + "bbox_xyxy": [100.0, 120.0, 240.0, 310.0], + "det_confidence": 0.91, + "det_class_id": 0, + "det_class_name": "bird", + "cls_class_id": 12, + "cls_class_name": "绿头鸭", + "cls_confidence": 0.87, + "box_area_ratio": 0.0032, + "is_large_box": true, + "is_dense_scene": false, + "show_cls_name": true +} +``` + +## 8. Docker 备选运行方式 + +如果不想在宿主机安装 Python 依赖，可以用 Docker。Docker 只是备选，默认仍然执行同一个 `run_batch.sh`。 + +构建镜像： + +```bash +docker build -t bird-cascade-batch:latest . +``` + +使用 Compose 处理项目根目录的 `input/`： + +```bash +docker compose run --rm bird-cascade-batch +``` + +临时覆盖参数： + +```bash +docker compose run --rm bird-cascade-batch \ + --input /app/input \ + --output /app/runs \ + --sample-interval 5 \ + --disable-density +``` + +`docker-compose.yml` 默认挂载： + +```yaml +volumes: + - ./models:/app/models:ro + - ./input:/app/input:ro + - ./runs:/app/runs +``` + +## 9. 程序结构 + +```text +app/ +├── batch_runner.py # 批处理入口 +├── config.py # 环境变量配置 +├── pipeline.py # 主级联流程 +├── utils.py # 图像、JSON、XML 工具 +├── video_pipeline.py # 视频平滑工具 +└── inference/ + ├── yolo_onnx.py # YOLOv5 ONNX 推理和后处理 + ├── convnext_onnx.py # ConvNeXt ONNX 推理 + └── density_onnx.py # 密度估计 ONNX 推理 +``` + +历史文件 `pipeline_raw.py`、`pipeline_add_density.py`、`config_raw.py`、`batch_runner_raw.py` 不参与默认 bash 批处理主流程。 diff --git a/README？.md b/README？.md new file mode 100644 index 0000000..8c5d4e5 --- /dev/null +++ b/README？.md @@ -0,0 +1,455 @@ +# Bird Detect Cascade Service：YOLOv5 + ConvNeXt 鸟类级联识别服务 + +本项目用于将 **YOLOv5 目标检测模型** 和 **ConvNeXt 鸟类细粒度分类模型** 级联起来： + +1. 接收图片或视频； +2. 使用 YOLOv5 检测鸟类目标位置； +3. 对每个检测框进行目标裁剪； +4. 将裁剪图送入 ConvNeXt 分类模型； +5. 返回每个目标的位置、检测置信度、分类类别、分类置信度； +6. 可选返回画框后的图片或视频。 + +项目默认使用 **ONNX Runtime** 推理，适合用 Docker 部署成 HTTP API 服务，方便网站后端调用。 + +--- + +## 1. 项目结构 + +```text +bird_detect_cascade_service/ +├── app/ +│ ├── main.py # FastAPI 接口入口 +│ ├── config.py # 配置读取 +│ ├── schemas.py # 返回结构 +│ ├── pipeline.py # YOLO + 裁剪 + ConvNeXt 级联流程 +│ ├── utils.py # 图像工具函数 +│ └── inference/ +│ ├── yolo_onnx.py # YOLOv5 ONNX 推理与后处理 +│ └── convnext_onnx.py # ConvNeXt ONNX 推理 +├── models/ +│ ├── yolov5.onnx # 你的 YOLOv5 ONNX 模型，需自行放入 +│ ├── convnext.onnx # 你的 ConvNeXt ONNX 模型，需自行放入 +│ ├── det_classes.txt # YOLO 类别名 +│ └── cls_classes.txt # ConvNeXt 类别名 +├── test_client.py # 本地接口测试脚本 +├── requirements.txt +├── Dockerfile +├── docker-compose.yml +└── README.md +``` + +--- + +## 2. 模型准备 + +### 2.1 YOLOv5 导出 ONNX + +如果你的 YOLOv5 是 `.pt` 权重，一般可在 YOLOv5 项目目录下执行： + +```bash +python export.py \ + --weights /path/to/best.pt \ + --img 960 960 \ + --batch 1 \ + --device cpu \ + --include onnx \ + --simplify +``` + +导出后把 ONNX 文件复制为： + +```bash +models/yolov5.onnx +``` + +### 2.2 ConvNeXt 导出 ONNX + +如果你已有 ConvNeXt 的 `.onnx`，直接复制为： + +```bash +models/convnext.onnx +``` + +如果你只有 PyTorch 权重，需要用你训练时的模型定义导出，输入通常为： + +```text +1 × 3 × 224 × 224 +``` + +如果你的 ConvNeXt 输入尺寸不是 224，需要修改 `.env` 或 `docker-compose.yml` 中的： + +```bash +CLS_IMG_SIZE=224 +``` + +--- + +## 3. 类别文件格式 + +### 3.1 YOLO 检测类别 + +`models/det_classes.txt` + +如果 YOLO 只检测鸟，可以写： + +```text +bird +``` + +如果 YOLO 有多个检测类别，例如： + +```text +bird +person +cat +dog +``` + +每一行对应一个类别，行号从 0 开始。 + +### 3.2 ConvNeXt 分类类别 + +`models/cls_classes.txt` + +每一行一个鸟类类别，顺序必须和训练时类别索引一致，例如： + +```text +Aix_galericulata +Anas_crecca +Anser_cygnoides +Cygnus_columbianus +``` + +--- + +## 4. 本地直接运行 + +安装依赖： + +```bash +pip install -r requirements.txt +``` + +启动服务： + +```bash +uvicorn app.main:app --host 0.0.0.0 --port 8000 +``` + +访问接口文档： + +```text +http://127.0.0.1:8000/docs +``` + +--- + +## 5. Docker 构建与运行 + +### 5.1 构建镜像 + +```bash +docker build -t bird-detect-cascade-service:latest . +``` + +### 5.2 启动容器 + +```bash +docker run -d \ + --name bird-detect-cascade-service \ + -p 8000:8000 \ + -v $(pwd)/models:/app/models \ + bird-detect-cascade-service:latest +``` + +### 5.3 使用 docker compose + +```bash +docker compose up -d --build +``` + +查看日志： + +```bash +docker logs -f bird-detect-cascade-service +``` + +停止服务： + +```bash +docker compose down +``` + +--- + +## 6. API 说明 + +### 6.1 健康检查 + +```bash +curl http://127.0.0.1:8000/health +``` + +返回示例： + +```json +{ + "status": "ok", + "det_model_loaded": true, + "cls_model_loaded": true +} +``` + +### 6.2 图片识别接口 + +接口： + +```text +POST /predict +``` + +参数： + +| 参数 | 类型 | 默认值 | 说明 | +|---|---|---:|---| +| file | file | 必填 | 上传图片 | +| return_image | bool | false | 是否返回画框后的 base64 图片 | +| det_conf | float | 配置文件默认值 | YOLO 检测置信度阈值 | +| iou_thres | float | 配置文件默认值 | NMS 的 IoU 阈值 | + +curl 示例： + +```bash +curl -X POST "http://127.0.0.1:8000/predict?return_image=true" \ + -F "file=@/path/to/test.jpg" +``` + +返回示例： + +```json +{ + "success": true, + "image_width": 1920, + "image_height": 1080, + "num_objects": 1, + "results": [ + { + "bbox_xyxy": [500.2, 130.7, 820.9, 600.4], + "det_confidence": 0.91, + "det_class_id": 0, + "det_class_name": "bird", + "cls_class_id": 12, + "cls_class_name": "Aix_galericulata", + "cls_confidence": 0.87 + } + ], + "annotated_image_base64": "..." +} +``` + +--- + +## 7. Python 测试调用 + +```bash +python test_client.py \ + --url http://127.0.0.1:8000/predict \ + --image /path/to/test.jpg \ + --save result.jpg +``` + +--- + +## 8. 关键环境变量 + +这些配置可以在 `docker-compose.yml` 中修改。 + +| 变量 | 默认值 | 说明 | +|---|---:|---| +| DET_MODEL_PATH | /app/models/yolov5.onnx | YOLOv5 ONNX 路径 | +| CLS_MODEL_PATH | /app/models/convnext.onnx | ConvNeXt ONNX 路径 | +| DET_CLASSES_PATH | /app/models/det_classes.txt | 检测类别文件 | +| CLS_CLASSES_PATH | /app/models/cls_classes.txt | 分类类别文件 | +| DET_IMG_SIZE | 960 | YOLO 输入尺寸 | +| CLS_IMG_SIZE | 224 | ConvNeXt 输入尺寸 | +| DET_CONF_THRES | 0.25 | YOLO 默认置信度阈值 | +| DET_IOU_THRES | 0.45 | NMS 默认 IoU 阈值 | +| CROP_PADDING_RATIO | 0.08 | 裁剪框外扩比例 | +| MAX_DETECTIONS | 50 | 单张图片最多返回目标数 | +| DEVICE | cpu | onnxruntime 设备，默认 cpu | + +--- + +## 9. 给网站同事的集成方式 + +网站后端只需要以 `multipart/form-data` 方式上传图片到： + +```text +POST http://服务器IP:8000/predict +``` + +前端或后端拿到 JSON 后，根据 `results` 字段展示识别结果即可。 + +如果网站只需要结果，不需要画框图： + +```text +return_image=false +``` + +如果网站需要展示画框图： + +```text +return_image=true +``` + +返回的 `annotated_image_base64` 是 JPEG 的 base64 字符串，前端可拼接为： + +```javascript +const imgSrc = "data:image/jpeg;base64," + annotated_image_base64; +``` + +--- + +## 10. 注意事项 + +1. YOLO 的输入尺寸必须和导出 ONNX 时一致，例如导出时是 `--img 960 960`，配置里也建议用 `DET_IMG_SIZE=960`。 +2. ConvNeXt 的类别文件顺序必须和训练时完全一致。 +3. 如果 YOLO 输出不是标准 YOLOv5 ONNX 格式 `[1, num_boxes, 5 + num_classes]`，需要根据实际输出改 `app/inference/yolo_onnx.py`。 +4. 如果 ConvNeXt 训练时使用了非 ImageNet 的 mean/std，需要修改 `app/inference/convnext_onnx.py` 中的归一化参数。 +5. 如果部署机器没有 GPU，默认 CPU 推理即可；如果需要 GPU 版本，需要替换 `onnxruntime` 为 `onnxruntime-gpu`，并使用支持 CUDA 的基础镜像。 + +--- + +## 11. 常见问题 + +### 11.1 为什么要用 HTTP API？ + +因为最终要集成到网站中，HTTP API 是最直接的对接方式。网站后端不需要关心模型细节，只需要上传图片并接收 JSON。 + +### 11.2 为什么先检测再分类？ + +YOLOv5 负责解决“鸟在哪里”，ConvNeXt 负责解决“这只鸟是什么种类”。先裁剪目标区域可以减少背景干扰，更适合细粒度分类。 + +### 11.3 多只鸟怎么处理？ + +YOLO 检测到几个目标，就裁剪几个目标，并分别送入 ConvNeXt 分类，最终返回一个结果列表。 + +### 11.4 没检测到鸟怎么办？ + +返回： + +```json +{ + "success": true, + "num_objects": 0, + "results": [] +} +``` + + + +--- + +## 12. 视频识别接口 + +当前版本新增视频接口： + +```text +POST /predict_video +``` + +### 12.1 后台测试视频上传 + +启动服务后，在项目根目录执行： + +```bash +bash backend_test_video.sh /path/to/test.mp4 +``` + +等价于： + +```bash +curl -X POST "http://127.0.0.1:8000/predict_video?sample_interval=25&save_annotated_video=true" \ + -F "file=@/path/to/test.mp4" +``` + +其中： + +- `sample_interval=25` 表示每 25 帧处理 1 帧； +- `save_annotated_video=true` 表示保存画框后视频； +- 如果视频很长，可以把 `sample_interval` 调大，例如 50、100，减少计算量。 + +### 12.2 视频运行结果保存位置 + +如果使用 `docker compose up -d --build` 启动，`docker-compose.yml` 已经挂载： + +```yaml +volumes: + - ./models:/app/models + - ./runs:/app/runs +``` + +所以结果会同时出现在： + +容器内： + +```text +/app/runs/video_时间_随机ID/ +``` + +宿主机项目目录下： + +```text +./runs/video_时间_随机ID/ +``` + +每次处理一个视频，会生成一个新的任务目录，例如： + +```text +runs/ +└── video_20260518_123000_a1b2c3d4/ + ├── source_video.mp4 # 上传的原始视频副本 + ├── annotated.mp4 # 画框后的结果视频 + ├── summary.json # 本次视频任务摘要 + ├── frame_results.json # 每个抽帧的详细识别结果 + └── frames/ + ├── frame_00000000.jpg # 有识别结果的关键帧截图 + ├── frame_00000025.jpg + └── ... +``` + +### 12.3 返回结果示例 + +```json +{ + "success": true, + "task_dir": "/app/runs/video_20260518_123000_a1b2c3d4", + "source_video": "/app/runs/video_20260518_123000_a1b2c3d4/source_video.mp4", + "annotated_video": "/app/runs/video_20260518_123000_a1b2c3d4/annotated.mp4", + "frames_dir": "/app/runs/video_20260518_123000_a1b2c3d4/frames", + "frame_results_json": "/app/runs/video_20260518_123000_a1b2c3d4/frame_results.json", + "summary_json": "/app/runs/video_20260518_123000_a1b2c3d4/summary.json", + "sample_interval": 25, + "processed_frames": 120, + "detected_frames": 87, + "elapsed_sec": 35.6 +} +``` + +### 12.4 网站集成建议 + +网站后端上传视频到： + +```text +POST http://服务器IP:8000/predict_video +``` + +如果网站需要展示结果，可以让网站后端读取返回的： + +```text +annotated_video +frame_results_json +summary_json +``` + +正式部署时更推荐把 `runs/` 目录映射到网站后端可访问的静态文件目录，或者在本服务里再增加 `/files/...` 静态下载接口。 diff --git a/app/__init__.py b/app/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/app/batch_runner.py b/app/batch_runner.py new file mode 100644 index 0000000..718725a --- /dev/null +++ b/app/batch_runner.py @@ -0,0 +1,924 @@ +from __future__ import annotations + +import argparse +import json +import time +from pathlib import Path +from typing import Dict, List, Sequence + +import cv2 + +from app.config import settings +from app.pipeline import BirdCascadePipeline +from app.utils import ( + decode_image_from_bytes, + safe_stem, + save_json, + write_image_xml, + write_video_xml, +) +from app.video_pipeline import VideoResultSmoother + + +IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".bmp", ".tif", ".tiff", ".webp"} +VIDEO_EXTS = {".mp4", ".avi", ".mov", ".mkv", ".flv", ".wmv", ".m4v"} + + +def parse_bool(value: str | bool) -> bool: + if isinstance(value, bool): + return value + + normalized = value.strip().lower() + if normalized in {"1", "true", "yes", "y", "on"}: + return True + if normalized in {"0", "false", "no", "n", "off"}: + return False + + raise argparse.ArgumentTypeError( + f"布尔值必须是 true/false、yes/no、1/0 之一，当前为: {value}" + ) + + +def now_timestamp_minute() -> str: + return time.strftime("%Y%m%d%H%M") + + +def is_image(path: Path) -> bool: + return path.suffix.lower() in IMAGE_EXTS + + +def is_video(path: Path) -> bool: + return path.suffix.lower() in VIDEO_EXTS + + +def collect_inputs(input_path: Path, recursive: bool = False) -> List[Path]: + if input_path.is_file(): + if is_image(input_path) or is_video(input_path): + return [input_path] + raise ValueError(f"不支持的文件类型: {input_path}") + + if not input_path.is_dir(): + raise FileNotFoundError(f"输入路径不存在: {input_path}") + + iterator = input_path.rglob("*") if recursive else input_path.glob("*") + + files = [] + for p in iterator: + if p.is_file() and (is_image(p) or is_video(p)): + files.append(p) + + files.sort() + return files + + +def make_output_dirs(input_path: Path, output_base: Path) -> Dict[str, Path]: + timestamp = now_timestamp_minute() + input_name = safe_stem(input_path.name) + + output_root = output_base / f"{timestamp}_{input_name}" + + if output_root.exists(): + idx = 1 + while True: + candidate = output_base / f"{timestamp}_{input_name}_{idx:02d}" + if not candidate.exists(): + output_root = candidate + break + idx += 1 + + result_dir = output_root / "result" + xml_dir = output_root / "xml" + json_dir = output_root / "json" + + result_dir.mkdir(parents=True, exist_ok=True) + xml_dir.mkdir(parents=True, exist_ok=True) + json_dir.mkdir(parents=True, exist_ok=True) + + return { + "root": output_root, + "result": result_dir, + "xml": xml_dir, + "json": json_dir, + } + + +def unique_path(path: Path) -> Path: + if not path.exists(): + return path + + stem = path.stem + suffix = path.suffix + parent = path.parent + + idx = 1 + while True: + candidate = parent / f"{stem}_{idx:02d}{suffix}" + if not candidate.exists(): + return candidate + idx += 1 + + +def process_image( + image_path: Path, + pipeline: BirdCascadePipeline, + dirs: Dict[str, Path], +) -> Dict: + img_bgr = decode_image_from_bytes(image_path.read_bytes()) + + output = pipeline.predict(img_bgr, return_image=False) + results = output["results"] + density_json = output.get("density") + density_result = output.get("density_result") + + result_img_path = unique_path(dirs["result"] / f"{image_path.stem}.jpg") + json_path = unique_path(dirs["json"] / f"{image_path.stem}.json") + xml_path = unique_path(dirs["xml"] / f"{image_path.stem}.xml") + + annotated = pipeline.draw_full_result( + img_bgr=img_bgr, + results=results, + density_result=density_result, + ) + cv2.imwrite(str(result_img_path), annotated) + + json_data = { + "success": True, + "type": "image", + "source_path": str(image_path), + "result_path": str(result_img_path), + "xml_path": str(xml_path), + "json_path": str(json_path), + "image_width": output["image_width"], + "image_height": output["image_height"], + "num_objects": output["num_objects"], + "density": density_json, + "results": results, + } + + save_json(json_path, json_data) + + write_image_xml( + xml_path=xml_path, + image_filename=result_img_path.name, + image_path=str(result_img_path), + width=int(output["image_width"]), + height=int(output["image_height"]), + results=results, + ) + + return json_data + + +def process_video( + video_path: Path, + pipeline: BirdCascadePipeline, + dirs: Dict[str, Path], + sample_interval: int, +) -> Dict: + result_video_path = unique_path(dirs["result"] / f"{video_path.stem}.mp4") + json_path = unique_path(dirs["json"] / f"{video_path.stem}.json") + xml_path = unique_path(dirs["xml"] / f"{video_path.stem}.xml") + + cap = cv2.VideoCapture(str(video_path)) + if not cap.isOpened(): + raise ValueError(f"视频打开失败: {video_path}") + + fps = cap.get(cv2.CAP_PROP_FPS) + if fps is None or fps <= 0: + fps = 25.0 + + width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH) or 0) + height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT) or 0) + total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT) or 0) + + fourcc = cv2.VideoWriter_fourcc(*"mp4v") + writer = cv2.VideoWriter(str(result_video_path), fourcc, fps, (width, height)) + + smoother = VideoResultSmoother( + hold_frames=settings.smooth_hold_frames, + alpha=settings.box_smooth_alpha, + ) + + frame_results = [] + processed_frames = 0 + detected_frames = 0 + frame_index = 0 + start_time = time.time() + + last_density_result = None + last_density_json = None + + try: + while True: + ok, frame = cap.read() + if not ok: + break + + should_process = frame_index % sample_interval == 0 + + if should_process: + output = pipeline.predict(frame, return_image=False) + + raw_results = output["results"] + density_result = output.get("density_result") + density_json = output.get("density") + + last_density_result = density_result + last_density_json = density_json + + processed_frames += 1 + + if settings.smooth_video_boxes: + results = smoother.update(raw_results, frame_index) + else: + results = raw_results + + if len(results) > 0: + detected_frames += 1 + + frame_results.append( + { + "frame_index": frame_index, + "time_sec": round(frame_index / fps, 3), + "num_objects": len(results), + "density": density_json, + "results": results, + } + ) + + else: + if settings.smooth_video_boxes: + results = smoother.get_for_skipped_frame(frame_index) + else: + results = [] + + density_result = last_density_result + density_json = last_density_json + + annotated_frame = pipeline.draw_full_result( + img_bgr=frame, + results=results, + density_result=density_result, + ) + writer.write(annotated_frame) + + frame_index += 1 + + finally: + cap.release() + writer.release() + + elapsed = round(time.time() - start_time, 3) + + summary = { + "success": True, + "type": "video", + "source_path": str(video_path), + "result_path": str(result_video_path), + "xml_path": str(xml_path), + "json_path": str(json_path), + "video_info": { + "width": width, + "height": height, + "fps": fps, + "total_frames": total_frames, + }, + "sample_interval": sample_interval, + "smooth_video_boxes": settings.smooth_video_boxes, + "smooth_hold_frames": settings.smooth_hold_frames, + "box_smooth_alpha": settings.box_smooth_alpha, + "processed_frames": processed_frames, + "detected_frames": detected_frames, + "elapsed_sec": elapsed, + } + + json_data = { + "summary": summary, + "frames": frame_results, + } + + save_json(json_path, json_data) + + write_video_xml( + xml_path=xml_path, + video_filename=result_video_path.name, + video_path=str(result_video_path), + width=width, + height=height, + fps=fps, + total_frames=total_frames, + sample_interval=sample_interval, + frame_results=frame_results, + ) + + return summary + + +def run_batch( + input_path: str, + output_base: str, + recursive: bool, + sample_interval: int, +) -> Dict: + input_path_obj = Path(input_path).resolve() + output_base_obj = Path(output_base).resolve() + output_base_obj.mkdir(parents=True, exist_ok=True) + + dirs = make_output_dirs(input_path_obj, output_base_obj) + files = collect_inputs(input_path_obj, recursive=recursive) + + if len(files) == 0: + raise ValueError(f"没有找到图片或视频文件: {input_path_obj}") + + pipeline = BirdCascadePipeline() + + success_results = [] + failed_results = [] + + for idx, file_path in enumerate(files, start=1): + print(f"[{idx}/{len(files)}] processing: {file_path}") + + try: + if is_image(file_path): + result = process_image( + image_path=file_path, + pipeline=pipeline, + dirs=dirs, + ) + elif is_video(file_path): + result = process_video( + video_path=file_path, + pipeline=pipeline, + dirs=dirs, + sample_interval=sample_interval, + ) + else: + continue + + success_results.append(result) + + except Exception as e: + failed = { + "source_path": str(file_path), + "error": str(e), + } + failed_results.append(failed) + print(f"[ERROR] {file_path}: {e}") + + index_json = { + "success": len(failed_results) == 0, + "input_path": str(input_path_obj), + "output_root": str(dirs["root"]), + "result_dir": str(dirs["result"]), + "xml_dir": str(dirs["xml"]), + "json_dir": str(dirs["json"]), + "total_files": len(files), + "success_files": len(success_results), + "failed_files": len(failed_results), + "failed": failed_results, + } + + index_path = dirs["root"] / "index.json" + index_path.write_text( + json.dumps(index_json, ensure_ascii=False, indent=2), + encoding="utf-8", + ) + + print("\nDone.") + print(f"Output root: {dirs['root']}") + print(f"Result dir : {dirs['result']}") + print(f"XML dir : {dirs['xml']}") + print(f"JSON dir : {dirs['json']}") + print(f"Index JSON : {index_path}") + + return index_json + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description="Bird cascade batch inference") + parser.add_argument( + "--input", + required=True, + help="输入路径，可以是单张图片、单个视频，或包含图片/视频的文件夹", + ) + parser.add_argument( + "--output", + default=settings.output_dir, + help="输出根目录，默认输出到项目 runs 下", + ) + + recursive_group = parser.add_mutually_exclusive_group() + recursive_group.add_argument( + "--recursive", + dest="recursive", + action="store_true", + help="输入为文件夹时递归处理子文件夹（默认）", + ) + recursive_group.add_argument( + "--no-recursive", + dest="recursive", + action="store_false", + help="输入为文件夹时只处理第一层文件", + ) + parser.set_defaults(recursive=True) + + parser.add_argument( + "--sample-interval", + type=int, + default=settings.video_sample_interval, + help="视频抽帧间隔", + ) + + detection_group = parser.add_mutually_exclusive_group() + detection_group.add_argument( + "--enable-detection", + dest="enable_detection", + action="store_true", + help="启用 YOLO 目标检测（默认）", + ) + detection_group.add_argument( + "--disable-detection", + dest="enable_detection", + action="store_false", + help="关闭 YOLO 目标检测", + ) + parser.set_defaults(enable_detection=settings.enable_detection) + + classification_group = parser.add_mutually_exclusive_group() + classification_group.add_argument( + "--enable-classification", + dest="enable_classification", + action="store_true", + help="启用 ConvNeXt 细粒度识别（默认）", + ) + classification_group.add_argument( + "--disable-classification", + dest="enable_classification", + action="store_false", + help="关闭 ConvNeXt 细粒度识别，仅保留 YOLO 检测框", + ) + parser.set_defaults(enable_classification=settings.enable_classification) + + density_group = parser.add_mutually_exclusive_group() + density_group.add_argument( + "--enable-density", + dest="enable_density", + action="store_true", + help="启用贝叶斯密度估计（默认）", + ) + density_group.add_argument( + "--disable-density", + dest="enable_density", + action="store_false", + help="关闭贝叶斯密度估计", + ) + parser.set_defaults(enable_density=settings.enable_density) + + parser.add_argument( + "--det-conf-thres", + type=float, + default=settings.det_conf_thres, + help="YOLO 第一层检测置信度阈值", + ) + parser.add_argument( + "--det-iou-thres", + type=float, + default=settings.det_iou_thres, + help="NMS IoU 阈值", + ) + parser.add_argument( + "--det-second-conf-thres", + type=float, + default=settings.det_second_conf_thres, + help="进入 ConvNeXt 前的第二层检测置信度阈值", + ) + parser.add_argument( + "--min-box-area-ratio", + type=float, + default=settings.min_box_area_ratio, + help="检测框最小面积占比", + ) + parser.add_argument( + "--max-box-area-ratio", + type=float, + default=settings.max_box_area_ratio, + help="检测框最大面积占比", + ) + parser.add_argument( + "--max-box-aspect-ratio", + type=float, + default=settings.max_box_aspect_ratio, + help="检测框最大长宽比", + ) + parser.add_argument( + "--detection-roi-y-min-ratio", + type=float, + default=settings.detection_roi_y_min_ratio, + help="只保留中心点位于该纵向比例以下区域之外的框，默认 0", + ) + parser.add_argument( + "--detection-roi-y-max-ratio", + type=float, + default=settings.detection_roi_y_max_ratio, + help="只保留中心点位于该纵向比例以上区域之前的框，默认 1", + ) + parser.add_argument( + "--cls-show-conf-thres", + type=float, + default=settings.cls_show_conf_thres, + help="普通目标显示物种名的分类置信度阈值", + ) + parser.add_argument( + "--large-box-area-ratio", + type=float, + default=settings.large_box_area_ratio, + help="近景大目标的框面积占比阈值", + ) + parser.add_argument( + "--large-box-cls-show-conf-thres", + type=float, + default=settings.large_box_cls_show_conf_thres, + help="大目标显示物种名的分类置信度阈值", + ) + parser.add_argument( + "--show-large-box-name-always", + type=parse_bool, + default=settings.show_large_box_name_always, + metavar="BOOL", + help="大目标是否始终显示非空物种名", + ) + parser.add_argument( + "--dense-scene-count-thres", + type=int, + default=settings.dense_scene_count_thres, + help="判定密集场景的密度估计数量阈值", + ) + parser.add_argument( + "--hide-cls-name-in-dense-scene", + type=parse_bool, + default=settings.hide_cls_name_in_dense_scene, + metavar="BOOL", + help="密集场景下是否隐藏物种名", + ) + parser.add_argument( + "--crop-padding-ratio", + type=float, + default=settings.crop_padding_ratio, + help="送入 ConvNeXt 前检测框向外扩的比例", + ) + parser.add_argument( + "--max-detections", + type=int, + default=settings.max_detections, + help="单张图或单帧最多保留的检测框数量", + ) + parser.add_argument( + "--uniform-species-color", + type=parse_bool, + default=settings.uniform_species_color, + metavar="BOOL", + help="所有物种的检测框和标签是否统一使用同一种颜色", + ) + parser.add_argument( + "--force-species-name", + dest="forced_species_name", + default=settings.forced_species_name, + help="人工确认整段视频为单一物种时，强制覆盖所有细粒度标签；留空不覆盖", + ) + parser.add_argument( + "--force-species-class-id", + dest="forced_species_class_id", + type=int, + default=settings.forced_species_class_id, + help="强制物种对应的类别 ID；默认 -1 时按标签文件自动查找", + ) + parser.add_argument( + "--force-species-min-box-area-ratio", + dest="forced_species_min_box_area_ratio", + type=float, + default=settings.forced_species_min_box_area_ratio, + help="只覆盖面积占比达到该值的目标；默认 0 表示覆盖全部检测目标", + ) + parser.add_argument( + "--species-name-remap", + default=settings.species_name_remap, + help="按人工复核纠正指定标签，格式：东方白鹳=白鹤,白鹭=白鹤", + ) + parser.add_argument( + "--species-name-remap-min-confidence", + type=float, + default=settings.species_name_remap_min_confidence, + help="只有原始分类置信度达到该值时才应用物种标签映射", + ) + parser.add_argument( + "--allowed-species-names", + default=settings.allowed_species_names, + help="只允许这些物种显示标签，英文逗号分隔；留空表示不限制", + ) + parser.add_argument( + "--binary-dark-species-name", + default=settings.binary_dark_species_name, + help="双物种明暗规则中的深色物种名称；留空表示关闭", + ) + parser.add_argument( + "--binary-light-species-name", + default=settings.binary_light_species_name, + help="双物种明暗规则中的浅色物种名称；留空表示关闭", + ) + parser.add_argument( + "--binary-dark-pixel-threshold", + type=int, + default=settings.binary_dark_pixel_threshold, + help="灰度低于该值的像素视为暗色像素，范围 0~255", + ) + parser.add_argument( + "--binary-dark-ratio-threshold", + type=float, + default=settings.binary_dark_ratio_threshold, + help="暗色像素比例达到该值时使用深色物种标签", + ) + parser.add_argument( + "--dark-override-species-name", + default=settings.dark_override_species_name, + help="暗色目标人工覆盖标签；留空表示关闭", + ) + parser.add_argument( + "--dark-override-species-class-id", + type=int, + default=settings.dark_override_species_class_id, + help="暗色目标人工覆盖标签的 class id；默认 -1 时按类别文件查找", + ) + parser.add_argument( + "--dark-override-pixel-threshold", + type=int, + default=settings.dark_override_pixel_threshold, + help="灰度低于该值的像素视为暗色像素，范围 0~255", + ) + parser.add_argument( + "--dark-override-ratio-threshold", + type=float, + default=settings.dark_override_ratio_threshold, + help="暗色像素比例达到该值时覆盖为指定标签", + ) + parser.add_argument( + "--dark-override-min-box-area-ratio", + type=float, + default=settings.dark_override_min_box_area_ratio, + help="只覆盖面积占比达到该值的暗色目标", + ) + parser.add_argument( + "--dark-override-max-box-area-ratio", + type=float, + default=settings.dark_override_max_box_area_ratio, + help="只覆盖面积占比不超过该值的暗色目标", + ) + parser.add_argument( + "--dark-override-roi-y-min-ratio", + type=float, + default=settings.dark_override_roi_y_min_ratio, + help="暗色覆盖只作用于中心点纵向比例不小于该值的目标", + ) + parser.add_argument( + "--dark-override-roi-y-max-ratio", + type=float, + default=settings.dark_override_roi_y_max_ratio, + help="暗色覆盖只作用于中心点纵向比例不大于该值的目标", + ) + parser.add_argument( + "--region-override-species-name", + default=settings.region_override_species_name, + help="区域人工覆盖标签；留空表示关闭", + ) + parser.add_argument( + "--region-override-species-class-id", + type=int, + default=settings.region_override_species_class_id, + help="区域人工覆盖标签的 class id；默认 -1 时按类别文件查找", + ) + parser.add_argument( + "--region-override-source-species-names", + default=settings.region_override_source_species_names, + help="只覆盖这些原标签，英文逗号分隔；留空表示覆盖区域内全部非空标签", + ) + parser.add_argument( + "--region-override-roi-y-min-ratio", + type=float, + default=settings.region_override_roi_y_min_ratio, + help="区域覆盖只作用于中心点纵向比例不小于该值的目标", + ) + parser.add_argument( + "--region-override-roi-y-max-ratio", + type=float, + default=settings.region_override_roi_y_max_ratio, + help="区域覆盖只作用于中心点纵向比例不大于该值的目标", + ) + parser.add_argument( + "--enable-crane-appearance-correction", + action="store_true", + default=settings.enable_crane_appearance_correction, + help="启用白鹤/白头鹤/白枕鹤外观二次纠正", + ) + parser.add_argument( + "--crane-white-ratio-threshold", + type=float, + default=settings.crane_white_ratio_threshold, + help="鹤类纠正中判定白鹤的亮白像素比例阈值", + ) + parser.add_argument( + "--crane-dark-ratio-threshold", + type=float, + default=settings.crane_dark_ratio_threshold, + help="鹤类纠正中判定白枕鹤/白头鹤的暗色像素比例阈值", + ) + parser.add_argument( + "--crane-tail-white-ratio-threshold", + type=float, + default=settings.crane_tail_white_ratio_threshold, + help="鹤类纠正中判定白枕鹤白尾/白斑的局部亮白比例阈值", + ) + parser.add_argument( + "--crane-uncertain-hide-label", + type=parse_bool, + default=settings.crane_uncertain_hide_label, + metavar="BOOL", + help="鹤类纠正无法确认时是否隐藏标签", + ) + + return parser + + +def validate_args( + parser: argparse.ArgumentParser, + args: argparse.Namespace, +) -> argparse.Namespace: + if args.sample_interval < 1: + parser.error("--sample-interval 必须大于等于 1") + + probability_fields = { + "--det-conf-thres": args.det_conf_thres, + "--det-iou-thres": args.det_iou_thres, + "--det-second-conf-thres": args.det_second_conf_thres, + "--cls-show-conf-thres": args.cls_show_conf_thres, + "--large-box-cls-show-conf-thres": args.large_box_cls_show_conf_thres, + "--species-name-remap-min-confidence": ( + args.species_name_remap_min_confidence + ), + "--detection-roi-y-min-ratio": args.detection_roi_y_min_ratio, + "--detection-roi-y-max-ratio": args.detection_roi_y_max_ratio, + "--binary-dark-ratio-threshold": args.binary_dark_ratio_threshold, + "--dark-override-ratio-threshold": args.dark_override_ratio_threshold, + "--dark-override-roi-y-min-ratio": args.dark_override_roi_y_min_ratio, + "--dark-override-roi-y-max-ratio": args.dark_override_roi_y_max_ratio, + "--region-override-roi-y-min-ratio": ( + args.region_override_roi_y_min_ratio + ), + "--region-override-roi-y-max-ratio": ( + args.region_override_roi_y_max_ratio + ), + "--crane-white-ratio-threshold": args.crane_white_ratio_threshold, + "--crane-dark-ratio-threshold": args.crane_dark_ratio_threshold, + "--crane-tail-white-ratio-threshold": ( + args.crane_tail_white_ratio_threshold + ), + } + for option, value in probability_fields.items(): + if not 0.0 <= value <= 1.0: + parser.error(f"{option} 必须在 0 到 1 之间") + + ratio_fields = { + "--min-box-area-ratio": args.min_box_area_ratio, + "--max-box-area-ratio": args.max_box_area_ratio, + "--large-box-area-ratio": args.large_box_area_ratio, + "--crop-padding-ratio": args.crop_padding_ratio, + "--force-species-min-box-area-ratio": ( + args.forced_species_min_box_area_ratio + ), + "--dark-override-min-box-area-ratio": ( + args.dark_override_min_box_area_ratio + ), + "--dark-override-max-box-area-ratio": ( + args.dark_override_max_box_area_ratio + ), + } + for option, value in ratio_fields.items(): + if value < 0.0: + parser.error(f"{option} 不能小于 0") + + if args.min_box_area_ratio > args.max_box_area_ratio: + parser.error("--min-box-area-ratio 不能大于 --max-box-area-ratio") + if args.detection_roi_y_min_ratio > args.detection_roi_y_max_ratio: + parser.error( + "--detection-roi-y-min-ratio 不能大于 " + "--detection-roi-y-max-ratio" + ) + if args.max_box_aspect_ratio <= 0.0: + parser.error("--max-box-aspect-ratio 必须大于 0") + if args.dense_scene_count_thres < 0: + parser.error("--dense-scene-count-thres 不能小于 0") + if args.max_detections < 1: + parser.error("--max-detections 必须大于等于 1") + if not 0 <= args.binary_dark_pixel_threshold <= 255: + parser.error("--binary-dark-pixel-threshold 必须在 0 到 255 之间") + if not 0 <= args.dark_override_pixel_threshold <= 255: + parser.error("--dark-override-pixel-threshold 必须在 0 到 255 之间") + if args.dark_override_min_box_area_ratio > args.dark_override_max_box_area_ratio: + parser.error( + "--dark-override-min-box-area-ratio 不能大于 " + "--dark-override-max-box-area-ratio" + ) + if args.dark_override_roi_y_min_ratio > args.dark_override_roi_y_max_ratio: + parser.error( + "--dark-override-roi-y-min-ratio 不能大于 " + "--dark-override-roi-y-max-ratio" + ) + if args.region_override_roi_y_min_ratio > args.region_override_roi_y_max_ratio: + parser.error( + "--region-override-roi-y-min-ratio 不能大于 " + "--region-override-roi-y-max-ratio" + ) + if bool(args.binary_dark_species_name) != bool( + args.binary_light_species_name + ): + parser.error( + "--binary-dark-species-name 和 " + "--binary-light-species-name 必须同时设置" + ) + + return args + + +def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: + parser = build_parser() + return validate_args(parser, parser.parse_args(argv)) + + +def apply_cli_settings(args: argparse.Namespace) -> None: + setting_names = ( + "enable_detection", + "enable_classification", + "enable_density", + "det_conf_thres", + "det_iou_thres", + "det_second_conf_thres", + "min_box_area_ratio", + "max_box_area_ratio", + "max_box_aspect_ratio", + "detection_roi_y_min_ratio", + "detection_roi_y_max_ratio", + "cls_show_conf_thres", + "large_box_area_ratio", + "large_box_cls_show_conf_thres", + "show_large_box_name_always", + "dense_scene_count_thres", + "hide_cls_name_in_dense_scene", + "crop_padding_ratio", + "max_detections", + "uniform_species_color", + "forced_species_name", + "forced_species_class_id", + "forced_species_min_box_area_ratio", + "species_name_remap", + "species_name_remap_min_confidence", + "allowed_species_names", + "binary_dark_species_name", + "binary_light_species_name", + "binary_dark_pixel_threshold", + "binary_dark_ratio_threshold", + "dark_override_species_name", + "dark_override_species_class_id", + "dark_override_pixel_threshold", + "dark_override_ratio_threshold", + "dark_override_min_box_area_ratio", + "dark_override_max_box_area_ratio", + "dark_override_roi_y_min_ratio", + "dark_override_roi_y_max_ratio", + "region_override_species_name", + "region_override_species_class_id", + "region_override_source_species_names", + "region_override_roi_y_min_ratio", + "region_override_roi_y_max_ratio", + "enable_crane_appearance_correction", + "crane_white_ratio_threshold", + "crane_dark_ratio_threshold", + "crane_tail_white_ratio_threshold", + "crane_uncertain_hide_label", + ) + + for name in setting_names: + object.__setattr__(settings, name, getattr(args, name)) + + +def main(argv: Sequence[str] | None = None): + args = parse_args(argv) + apply_cli_settings(args) + + return run_batch( + input_path=args.input, + output_base=args.output, + recursive=args.recursive, + sample_interval=args.sample_interval, + ) + + +if __name__ == "__main__": + main() diff --git a/app/batch_runner_raw.py b/app/batch_runner_raw.py new file mode 100644 index 0000000..82b79b2 --- /dev/null +++ b/app/batch_runner_raw.py @@ -0,0 +1,396 @@ +from __future__ import annotations + +import argparse +import json +import time +from pathlib import Path +from typing import Dict, List + +import cv2 + +from app.config import settings +from app.pipeline import BirdCascadePipeline +from app.utils import ( + decode_image_from_bytes, + draw_results, + safe_stem, + save_json, + write_image_xml, + write_video_xml, +) +from app.video_pipeline import VideoResultSmoother + + +IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".bmp", ".tif", ".tiff", ".webp"} +VIDEO_EXTS = {".mp4", ".avi", ".mov", ".mkv", ".flv", ".wmv", ".m4v"} + + +def now_timestamp_minute() -> str: + return time.strftime("%Y%m%d%H%M") + + +def is_image(path: Path) -> bool: + return path.suffix.lower() in IMAGE_EXTS + + +def is_video(path: Path) -> bool: + return path.suffix.lower() in VIDEO_EXTS + + +def collect_inputs(input_path: Path, recursive: bool = False) -> List[Path]: + if input_path.is_file(): + if is_image(input_path) or is_video(input_path): + return [input_path] + raise ValueError(f"不支持的文件类型: {input_path}") + + if not input_path.is_dir(): + raise FileNotFoundError(f"输入路径不存在: {input_path}") + + iterator = input_path.rglob("*") if recursive else input_path.glob("*") + + files = [] + for p in iterator: + if p.is_file() and (is_image(p) or is_video(p)): + files.append(p) + + files.sort() + return files + + +def make_output_dirs(input_path: Path, output_base: Path) -> Dict[str, Path]: + timestamp = now_timestamp_minute() + input_name = safe_stem(input_path.name) + + output_root = output_base / f"{timestamp}_{input_name}" + + if output_root.exists(): + idx = 1 + while True: + candidate = output_base / f"{timestamp}_{input_name}_{idx:02d}" + if not candidate.exists(): + output_root = candidate + break + idx += 1 + + result_dir = output_root / "result" + xml_dir = output_root / "xml" + json_dir = output_root / "json" + + result_dir.mkdir(parents=True, exist_ok=True) + xml_dir.mkdir(parents=True, exist_ok=True) + json_dir.mkdir(parents=True, exist_ok=True) + + return { + "root": output_root, + "result": result_dir, + "xml": xml_dir, + "json": json_dir, + } + + +def unique_path(path: Path) -> Path: + """ + 如果不同子目录里存在同名文件，避免覆盖。 + 例如 result/1.jpg 已存在，则自动变成 result/1_01.jpg。 + """ + if not path.exists(): + return path + + stem = path.stem + suffix = path.suffix + parent = path.parent + + idx = 1 + while True: + candidate = parent / f"{stem}_{idx:02d}{suffix}" + if not candidate.exists(): + return candidate + idx += 1 + + +def process_image( + image_path: Path, + pipeline: BirdCascadePipeline, + dirs: Dict[str, Path], +) -> Dict: + img_bgr = decode_image_from_bytes(image_path.read_bytes()) + + output = pipeline.predict(img_bgr, return_image=False) + results = output["results"] + + result_img_path = unique_path(dirs["result"] / f"{image_path.stem}.jpg") + json_path = unique_path(dirs["json"] / f"{image_path.stem}.json") + xml_path = unique_path(dirs["xml"] / f"{image_path.stem}.xml") + + annotated = draw_results(img_bgr, results) + cv2.imwrite(str(result_img_path), annotated) + + json_data = { + "success": True, + "type": "image", + "source_path": str(image_path), + "result_path": str(result_img_path), + "xml_path": str(xml_path), + "json_path": str(json_path), + "image_width": output["image_width"], + "image_height": output["image_height"], + "num_objects": output["num_objects"], + "results": results, + } + + save_json(json_path, json_data) + + write_image_xml( + xml_path=xml_path, + image_filename=result_img_path.name, + image_path=str(result_img_path), + width=int(output["image_width"]), + height=int(output["image_height"]), + results=results, + ) + + return json_data + + +def process_video( + video_path: Path, + pipeline: BirdCascadePipeline, + dirs: Dict[str, Path], + sample_interval: int, +) -> Dict: + result_video_path = unique_path(dirs["result"] / f"{video_path.stem}.mp4") + json_path = unique_path(dirs["json"] / f"{video_path.stem}.json") + xml_path = unique_path(dirs["xml"] / f"{video_path.stem}.xml") + + cap = cv2.VideoCapture(str(video_path)) + if not cap.isOpened(): + raise ValueError(f"视频打开失败: {video_path}") + + fps = cap.get(cv2.CAP_PROP_FPS) + if fps is None or fps <= 0: + fps = 25.0 + + width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH) or 0) + height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT) or 0) + total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT) or 0) + + fourcc = cv2.VideoWriter_fourcc(*"mp4v") + writer = cv2.VideoWriter(str(result_video_path), fourcc, fps, (width, height)) + + smoother = VideoResultSmoother( + hold_frames=settings.smooth_hold_frames, + alpha=settings.box_smooth_alpha, + ) + + frame_results = [] + processed_frames = 0 + detected_frames = 0 + frame_index = 0 + start_time = time.time() + + try: + while True: + ok, frame = cap.read() + if not ok: + break + + should_process = frame_index % sample_interval == 0 + + if should_process: + output = pipeline.predict(frame, return_image=False) + raw_results = output["results"] + processed_frames += 1 + + if settings.smooth_video_boxes: + results = smoother.update(raw_results, frame_index) + else: + results = raw_results + + if len(results) > 0: + detected_frames += 1 + + frame_results.append( + { + "frame_index": frame_index, + "time_sec": round(frame_index / fps, 3), + "num_objects": len(results), + "results": results, + } + ) + else: + if settings.smooth_video_boxes: + results = smoother.get_for_skipped_frame(frame_index) + else: + results = [] + + annotated_frame = draw_results(frame, results) + writer.write(annotated_frame) + + frame_index += 1 + + finally: + cap.release() + writer.release() + + elapsed = round(time.time() - start_time, 3) + + summary = { + "success": True, + "type": "video", + "source_path": str(video_path), + "result_path": str(result_video_path), + "xml_path": str(xml_path), + "json_path": str(json_path), + "video_info": { + "width": width, + "height": height, + "fps": fps, + "total_frames": total_frames, + }, + "sample_interval": sample_interval, + "smooth_video_boxes": settings.smooth_video_boxes, + "smooth_hold_frames": settings.smooth_hold_frames, + "box_smooth_alpha": settings.box_smooth_alpha, + "processed_frames": processed_frames, + "detected_frames": detected_frames, + "elapsed_sec": elapsed, + } + + json_data = { + "summary": summary, + "frames": frame_results, + } + + save_json(json_path, json_data) + + write_video_xml( + xml_path=xml_path, + video_filename=result_video_path.name, + video_path=str(result_video_path), + width=width, + height=height, + fps=fps, + total_frames=total_frames, + sample_interval=sample_interval, + frame_results=frame_results, + ) + + return summary + + +def run_batch( + input_path: str, + output_base: str, + recursive: bool, + sample_interval: int, +) -> Dict: + input_path_obj = Path(input_path).resolve() + output_base_obj = Path(output_base).resolve() + output_base_obj.mkdir(parents=True, exist_ok=True) + + dirs = make_output_dirs(input_path_obj, output_base_obj) + files = collect_inputs(input_path_obj, recursive=recursive) + + if len(files) == 0: + raise ValueError(f"没有找到图片或视频文件: {input_path_obj}") + + pipeline = BirdCascadePipeline() + + success_results = [] + failed_results = [] + + for idx, file_path in enumerate(files, start=1): + print(f"[{idx}/{len(files)}] processing: {file_path}") + + try: + if is_image(file_path): + result = process_image( + image_path=file_path, + pipeline=pipeline, + dirs=dirs, + ) + elif is_video(file_path): + result = process_video( + video_path=file_path, + pipeline=pipeline, + dirs=dirs, + sample_interval=sample_interval, + ) + else: + continue + + success_results.append(result) + + except Exception as e: + failed = { + "source_path": str(file_path), + "error": str(e), + } + failed_results.append(failed) + print(f"[ERROR] {file_path}: {e}") + + index_json = { + "success": len(failed_results) == 0, + "input_path": str(input_path_obj), + "output_root": str(dirs["root"]), + "result_dir": str(dirs["result"]), + "xml_dir": str(dirs["xml"]), + "json_dir": str(dirs["json"]), + "total_files": len(files), + "success_files": len(success_results), + "failed_files": len(failed_results), + "failed": failed_results, + } + + index_path = dirs["root"] / "index.json" + index_path.write_text( + json.dumps(index_json, ensure_ascii=False, indent=2), + encoding="utf-8", + ) + + print("\nDone.") + print(f"Output root: {dirs['root']}") + print(f"Result dir : {dirs['result']}") + print(f"XML dir : {dirs['xml']}") + print(f"JSON dir : {dirs['json']}") + print(f"Index JSON : {index_path}") + + return index_json + + +def main(): + parser = argparse.ArgumentParser(description="Bird cascade batch inference") + parser.add_argument( + "--input", + required=True, + help="输入路径，可以是单张图片、单个视频，或包含图片/视频的文件夹", + ) + parser.add_argument( + "--output", + default="/data2/zhangw_data/Bird_detect_cascade_service/runs", + help="输出根目录，默认直接输出到项目 runs 下", + ) + parser.add_argument( + "--recursive", + action="store_true", + help="输入为文件夹时，是否递归处理子文件夹", + ) + parser.add_argument( + "--sample-interval", + type=int, + default=settings.video_sample_interval, + help="视频抽帧间隔", + ) + + args = parser.parse_args() + + run_batch( + input_path=args.input, + output_base=args.output, + recursive=args.recursive, + sample_interval=args.sample_interval, + ) + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/app/config.py b/app/config.py new file mode 100644 index 0000000..a93eb25 --- /dev/null +++ b/app/config.py @@ -0,0 +1,494 @@ +import os +from dataclasses import dataclass + + +# ========================= +# 项目根目录 +# ========================= +# 例如： +# 宿主机： +# /data2/zhangw_data/Bird_detect_cascade_service +# +# Docker： +# /Bird_detect_cascade_service +# +# 会自动适配，不需要写死绝对路径 +# ========================= +BASE_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) + +MODELS_DIR = os.path.join(BASE_DIR, "models") +RUNS_DIR = os.path.join(BASE_DIR, "runs") + + +def _get_str(name: str, default: str) -> str: + return os.getenv(name, default) + + +def _get_int(name: str, default: int) -> int: + value = os.getenv(name) + if value is None or value == "": + return default + return int(value) + + +def _get_float(name: str, default: float) -> float: + value = os.getenv(name) + if value is None or value == "": + return default + return float(value) + + +def _get_bool(name: str, default: bool) -> bool: + value = os.getenv(name) + if value is None or value == "": + return default + return value.lower() in ["1", "true", "yes", "y"] + + +def _get_int_list(name: str, default: str) -> list[int]: + value = os.getenv(name, default) + + if value is None or value.strip() == "": + return [] + + return [ + int(x.strip()) + for x in value.split(",") + if x.strip() != "" + ] + + +@dataclass(frozen=True) +class Settings: + + # ========================= + # 模型开关 + # ========================= + enable_detection: bool = _get_bool( + "ENABLE_DETECTION", + True, + ) + + enable_classification: bool = _get_bool( + "ENABLE_CLASSIFICATION", + True, + ) + + # ========================= + # YOLOv5 模型 + # ========================= + det_model_path: str = _get_str( + "DET_MODEL_PATH", + os.path.join( + MODELS_DIR, + "YNDC_20251010_yolov5_optimize.onnx", + ), + ) + + det_classes_path: str = _get_str( + "DET_CLASSES_PATH", + os.path.join( + MODELS_DIR, + "YNDC_20251010_yolov5_optimize.txt", + ), + ) + + # ========================= + # ConvNeXt 分类模型 + # ========================= + cls_model_path: str = _get_str( + "CLS_MODEL_PATH", + os.path.join( + MODELS_DIR, + "JXPYH_20260204_mmcls_optimize.onnx", + ), + ) + + cls_classes_path: str = _get_str( + "CLS_CLASSES_PATH", + os.path.join( + MODELS_DIR, + "JXPYH_20260204_mmcls_optimize.txt", + ), + ) + + # ========================= + # 密度估计模型 + # ========================= + enable_density: bool = _get_bool( + "ENABLE_DENSITY", + True, + ) + + density_model_path: str = _get_str( + "DENSITY_MODEL_PATH", + os.path.join( + MODELS_DIR, + "bird_density.onnx", + ), + ) + + # 如果模型输入固定，例如 400×400 + # 可以设置 DENSITY_IMG_SIZE=400 + # 默认 0 表示自动读取 ONNX 输入尺寸 + density_img_size: int = _get_int( + "DENSITY_IMG_SIZE", + 0, + ) + + # ========================= + # 输入尺寸 + # ========================= + det_img_size: int = _get_int( + "DET_IMG_SIZE", + 640, + ) + + cls_img_size: int = _get_int( + "CLS_IMG_SIZE", + 224, + ) + + # ========================= + # YOLO 阈值 + # ========================= + det_conf_thres: float = _get_float( + "DET_CONF_THRES", + 0.25, + ) + + det_iou_thres: float = _get_float( + "DET_IOU_THRES", + 0.45, + ) + + # 二次过滤 YOLO 检测框，避免把明显错误的框传给 ConvNeXt + det_second_conf_thres: float = _get_float( + "DET_SECOND_CONF_THRES", + 0.35, + ) + + # 最小框面积占整图比例，太小的噪声框过滤 + min_box_area_ratio: float = _get_float( + "MIN_BOX_AREA_RATIO", + 0.00002, + ) + + # 最大框面积占整图比例，太大的误检框过滤 + max_box_area_ratio: float = _get_float( + "MAX_BOX_AREA_RATIO", + 1.0, + ) + + # 最大长宽比，过细过长的框过滤 + max_box_aspect_ratio: float = _get_float( + "MAX_BOX_ASPECT_RATIO", + 6.0, + ) + + # 可选检测区域：按检测框中心点的纵向位置过滤。 + # 默认 0~1 表示整幅画面；可排除固定字幕、时间戳或岸边噪声区域。 + detection_roi_y_min_ratio: float = _get_float( + "DETECTION_ROI_Y_MIN_RATIO", + 0.0, + ) + + detection_roi_y_max_ratio: float = _get_float( + "DETECTION_ROI_Y_MAX_RATIO", + 1.0, + ) + + # 分类置信度低于该值时，不显示物种名称 + cls_show_conf_thres: float = _get_float( + "CLS_SHOW_CONF_THRES", + 0.50, + ) + + # 大目标判断阈值：框面积占整图比例超过该值，认为是近景大目标 + large_box_area_ratio: float = _get_float( + "LARGE_BOX_AREA_RATIO", + 0.002, + ) + + # 大目标分类显示阈值，近景大目标即使分类置信度稍低，也显示物种名 + large_box_cls_show_conf_thres: float = _get_float( + "LARGE_BOX_CLS_SHOW_CONF_THRES", + 0.25, + ) + + # 数量达到该阈值时，认为是密集场景，不显示物种名 + dense_scene_count_thres: int = _get_int( + "DENSE_SCENE_COUNT_THRES", + 20, + ) + + # 大目标强制显示物种名，只要分类结果非空即可 + show_large_box_name_always: bool = _get_bool( + "SHOW_LARGE_BOX_NAME_ALWAYS", + True, + ) + + # 是否在密集场景隐藏所有物种名 + hide_cls_name_in_dense_scene: bool = _get_bool( + "HIDE_CLS_NAME_IN_DENSE_SCENE", + True, + ) + + # 所有物种是否统一使用同一种标签和检测框颜色 + uniform_species_color: bool = _get_bool( + "UNIFORM_SPECIES_COLOR", + False, + ) + + # 可选：整段素材已人工确认只有一个物种时，覆盖细粒度模型标签。 + # 留空表示使用模型输出。 + forced_species_name: str = _get_str( + "FORCED_SPECIES_NAME", + "", + ).strip() + + forced_species_class_id: int = _get_int( + "FORCED_SPECIES_CLASS_ID", + -1, + ) + + # 只对达到该框面积占比的目标应用人工物种覆盖。 + # 可用来统一主要大型鸟类，同时保留画面中的小型伴生鸟类。 + forced_species_min_box_area_ratio: float = _get_float( + "FORCED_SPECIES_MIN_BOX_AREA_RATIO", + 0.0, + ) + + # 视频级人工复核映射，格式示例： + # 东方白鹳=白鹤,白鹭=白鹤 + species_name_remap: str = _get_str( + "SPECIES_NAME_REMAP", + "", + ).strip() + + # 只有原始细粒度分类置信度达到该阈值时才应用标签映射。 + # 默认 0 保持原有行为；混合物种场景可提高该值以避免低置信度误映射。 + species_name_remap_min_confidence: float = _get_float( + "SPECIES_NAME_REMAP_MIN_CONFIDENCE", + 0.0, + ) + + # 留空表示允许全部；非空时，只有白名单中的物种显示名称和物种色。 + allowed_species_names: str = _get_str( + "ALLOWED_SPECIES_NAMES", + "", + ).strip() + + # 可选双物种明暗规则。两个名称都非空时，不运行细粒度分类模型， + # 而是根据检测框内暗色像素比例区分两个外观差异明显的物种。 + binary_dark_species_name: str = _get_str( + "BINARY_DARK_SPECIES_NAME", + "", + ).strip() + + binary_light_species_name: str = _get_str( + "BINARY_LIGHT_SPECIES_NAME", + "", + ).strip() + + binary_dark_pixel_threshold: int = _get_int( + "BINARY_DARK_PIXEL_THRESHOLD", + 85, + ) + + binary_dark_ratio_threshold: float = _get_float( + "BINARY_DARK_RATIO_THRESHOLD", + 0.12, + ) + + # 可选：对画面中的深色目标做人工标签覆盖。 + # 适合模型把黑色雁鸭类误判成白骨顶等细粒度类别的场景。 + dark_override_species_name: str = _get_str( + "DARK_OVERRIDE_SPECIES_NAME", + "", + ).strip() + + dark_override_species_class_id: int = _get_int( + "DARK_OVERRIDE_SPECIES_CLASS_ID", + -1, + ) + + dark_override_pixel_threshold: int = _get_int( + "DARK_OVERRIDE_PIXEL_THRESHOLD", + 80, + ) + + dark_override_ratio_threshold: float = _get_float( + "DARK_OVERRIDE_RATIO_THRESHOLD", + 0.35, + ) + + dark_override_min_box_area_ratio: float = _get_float( + "DARK_OVERRIDE_MIN_BOX_AREA_RATIO", + 0.0, + ) + + dark_override_max_box_area_ratio: float = _get_float( + "DARK_OVERRIDE_MAX_BOX_AREA_RATIO", + 1.0, + ) + + dark_override_roi_y_min_ratio: float = _get_float( + "DARK_OVERRIDE_ROI_Y_MIN_RATIO", + 0.0, + ) + + dark_override_roi_y_max_ratio: float = _get_float( + "DARK_OVERRIDE_ROI_Y_MAX_RATIO", + 1.0, + ) + + # 可选：按画面纵向区域做人工标签覆盖。 + # 用于同一视频中上/中/下区域物种分布稳定的场景。 + region_override_species_name: str = _get_str( + "REGION_OVERRIDE_SPECIES_NAME", + "", + ).strip() + + region_override_species_class_id: int = _get_int( + "REGION_OVERRIDE_SPECIES_CLASS_ID", + -1, + ) + + region_override_source_species_names: str = _get_str( + "REGION_OVERRIDE_SOURCE_SPECIES_NAMES", + "", + ).strip() + + region_override_roi_y_min_ratio: float = _get_float( + "REGION_OVERRIDE_ROI_Y_MIN_RATIO", + 0.0, + ) + + region_override_roi_y_max_ratio: float = _get_float( + "REGION_OVERRIDE_ROI_Y_MAX_RATIO", + 1.0, + ) + + # 可选：鹤类混淆纠正。针对白鹤/白头鹤/白枕鹤相互误判， + # 根据检测框内亮白像素和暗色像素比例做保守二次校验。 + enable_crane_appearance_correction: bool = _get_bool( + "ENABLE_CRANE_APPEARANCE_CORRECTION", + False, + ) + + crane_white_ratio_threshold: float = _get_float( + "CRANE_WHITE_RATIO_THRESHOLD", + 0.30, + ) + + crane_dark_ratio_threshold: float = _get_float( + "CRANE_DARK_RATIO_THRESHOLD", + 0.28, + ) + + crane_tail_white_ratio_threshold: float = _get_float( + "CRANE_TAIL_WHITE_RATIO_THRESHOLD", + 0.16, + ) + + crane_uncertain_hide_label: bool = _get_bool( + "CRANE_UNCERTAIN_HIDE_LABEL", + True, + ) + + # ========================= + # 裁剪外扩比例 + # ========================= + crop_padding_ratio: float = _get_float( + "CROP_PADDING_RATIO", + 0.08, + ) + + # ========================= + # 最多检测目标数 + # ========================= + max_detections: int = _get_int( + "MAX_DETECTIONS", + 50, + ) + + # ========================= + # 推理设备 + # cpu / cuda + # ========================= + device: str = _get_str( + "DEVICE", + "cpu", + ) + + # ========================= + # 输出目录 + # ========================= + output_dir: str = _get_str( + "OUTPUT_DIR", + RUNS_DIR, + ) + + image_output_subdir: str = _get_str( + "IMAGE_OUTPUT_SUBDIR", + "images", + ) + + video_output_subdir: str = _get_str( + "VIDEO_OUTPUT_SUBDIR", + "video", + ) + + # ========================= + # 视频抽帧间隔 + # ========================= + video_sample_interval: int = _get_int( + "VIDEO_SAMPLE_INTERVAL", + 2, + ) + + # ========================= + # 是否保存画框视频 + # ========================= + save_annotated_video: bool = _get_bool( + "SAVE_ANNOTATED_VIDEO", + True, + ) + + # ========================= + # 只保留 YOLO 类别 + # 默认只保留 bird + # ========================= + target_det_class_ids: list[int] = None + + # ========================= + # 视频框平滑 + # ========================= + smooth_video_boxes: bool = _get_bool( + "SMOOTH_VIDEO_BOXES", + True, + ) + + smooth_hold_frames: int = _get_int( + "SMOOTH_HOLD_FRAMES", + 3, + ) + + box_smooth_alpha: float = _get_float( + "BOX_SMOOTH_ALPHA", + 0.25, + ) + + def __post_init__(self): + object.__setattr__( + self, + "target_det_class_ids", + _get_int_list( + "TARGET_DET_CLASS_IDS", + "0", + ), + ) + + +settings = Settings() diff --git a/app/config_raw.py b/app/config_raw.py new file mode 100644 index 0000000..0c238a8 --- /dev/null +++ b/app/config_raw.py @@ -0,0 +1,105 @@ +import os +from dataclasses import dataclass + + +def _get_str(name: str, default: str) -> str: + return os.getenv(name, default) + + +def _get_int(name: str, default: int) -> int: + value = os.getenv(name) + if value is None or value == "": + return default + return int(value) + + +def _get_float(name: str, default: float) -> float: + value = os.getenv(name) + if value is None or value == "": + return default + return float(value) + + +def _get_bool(name: str, default: bool) -> bool: + value = os.getenv(name) + if value is None or value == "": + return default + return value.lower() in ["1", "true", "yes", "y"] + + +def _get_int_list(name: str, default: str) -> list[int]: + value = os.getenv(name, default) + if value is None or value.strip() == "": + return [] + return [int(x.strip()) for x in value.split(",") if x.strip() != ""] + + +@dataclass(frozen=True) +class Settings: + # Docker 内默认模型路径 + det_model_path: str = _get_str( + "DET_MODEL_PATH", + "/app/models/YNDC_20251010_yolov5_optimize.onnx", + ) + cls_model_path: str = _get_str( + "CLS_MODEL_PATH", + "/app/models/HHK_20260525_mmcls_optimize.onnx", + ) + det_classes_path: str = _get_str( + "DET_CLASSES_PATH", + "/app/models/YNDC_20251010_yolov5_optimize.txt", + ) + cls_classes_path: str = _get_str( + "CLS_CLASSES_PATH", + "/app/models/HHK_20260525_mmcls_optimize.txt", + ) + + # 输入尺寸 + det_img_size: int = _get_int("DET_IMG_SIZE", 640) + cls_img_size: int = _get_int("CLS_IMG_SIZE", 224) + + # YOLO 检测阈值 + det_conf_thres: float = _get_float("DET_CONF_THRES", 0.25) + det_iou_thres: float = _get_float("DET_IOU_THRES", 0.45) + + # 目标裁剪外扩比例 + crop_padding_ratio: float = _get_float("CROP_PADDING_RATIO", 0.08) + + # 最多保留目标数量 + max_detections: int = _get_int("MAX_DETECTIONS", 50) + + # 当前环境没有 onnxruntime-gpu 时，建议用 cpu + device: str = _get_str("DEVICE", "cpu") + + # 输出目录 + output_dir: str = _get_str("OUTPUT_DIR", "/app/runs") + image_output_subdir: str = _get_str("IMAGE_OUTPUT_SUBDIR", "images") + video_output_subdir: str = _get_str("VIDEO_OUTPUT_SUBDIR", "video") + + # 视频抽帧间隔 + video_sample_interval: int = _get_int("VIDEO_SAMPLE_INTERVAL", 2) + + # 是否保存画框后视频 + save_annotated_video: bool = _get_bool("SAVE_ANNOTATED_VIDEO", True) + + # 只保留 YOLO 类别 0，也就是“鸟” + target_det_class_ids: list[int] = None + + # 视频平滑参数 + smooth_video_boxes: bool = _get_bool("SMOOTH_VIDEO_BOXES", True) + + # 未处理帧或短暂漏检时，最多复用多少帧上一结果 + smooth_hold_frames: int = _get_int("SMOOTH_HOLD_FRAMES", 3) + + # 框坐标平滑系数，越大越稳定，但响应越慢 + box_smooth_alpha: float = _get_float("BOX_SMOOTH_ALPHA", 0.25) + + def __post_init__(self): + object.__setattr__( + self, + "target_det_class_ids", + _get_int_list("TARGET_DET_CLASS_IDS", "0"), + ) + + +settings = Settings() \ No newline at end of file diff --git a/app/image_pipeline.py b/app/image_pipeline.py new file mode 100644 index 0000000..eeec26c --- /dev/null +++ b/app/image_pipeline.py @@ -0,0 +1,121 @@ +from pathlib import Path +from typing import Dict + +import cv2 +import numpy as np + +from app.config import settings +from app.utils import ( + draw_results, + now_timestamp_minute, + safe_stem, + save_json, + write_image_xml, +) + + +def make_unique_result_dir(parent_dir: Path, timestamp: str, original_filename: str) -> tuple[Path, str, str]: + """ + 创建单次图片识别结果目录。 + + 目录名： + 时间戳_原文件名stem + + 例如： + runs/images/202605181431_DJI_20250820151801_0006_Z/ + + 返回： + result_dir, result_base_name, original_stem + """ + original_stem = safe_stem(original_filename) + base_name = f"{timestamp}_{original_stem}" + + result_dir = parent_dir / base_name + + if not result_dir.exists(): + result_dir.mkdir(parents=True, exist_ok=True) + return result_dir, base_name, original_stem + + idx = 1 + while True: + candidate_base = f"{base_name}_{idx:02d}" + candidate_dir = parent_dir / candidate_base + if not candidate_dir.exists(): + candidate_dir.mkdir(parents=True, exist_ok=True) + return candidate_dir, candidate_base, original_stem + idx += 1 + + +def save_image_outputs( + img_bgr: np.ndarray, + original_filename: str, + predict_output: Dict, +) -> Dict: + """ + 图片结果自动保存到： + runs/images/时间戳_原图片名/ + + 目录内部生成： + 原图片名.jpg + 时间戳_原图片名.json + 时间戳_原图片名.xml + + 示例： + runs/images/202605181431_DJI_20250820151801_0006_Z/ + ├── DJI_20250820151801_0006_Z.jpg + ├── 202605181431_DJI_20250820151801_0006_Z.json + └── 202605181431_DJI_20250820151801_0006_Z.xml + """ + images_root = Path(settings.output_dir) / settings.image_output_subdir + images_root.mkdir(parents=True, exist_ok=True) + + timestamp = now_timestamp_minute() + result_dir, result_base_name, original_stem = make_unique_result_dir( + parent_dir=images_root, + timestamp=timestamp, + original_filename=original_filename, + ) + + # 图片文件名不带时间戳，按你的要求保留原图片名 stem + image_path = result_dir / f"{original_stem}.jpg" + + # JSON 和 XML 带时间戳 + json_path = result_dir / f"{result_base_name}.json" + xml_path = result_dir / f"{result_base_name}.xml" + + results = predict_output.get("results", []) + + annotated = draw_results(img_bgr, results) + cv2.imwrite(str(image_path), annotated) + + json_data = { + "success": True, + "type": "image", + "original_filename": original_filename, + "result_dir": str(result_dir), + "saved_image": str(image_path), + "saved_json": str(json_path), + "saved_xml": str(xml_path), + "image_width": predict_output.get("image_width"), + "image_height": predict_output.get("image_height"), + "num_objects": predict_output.get("num_objects"), + "results": results, + } + + save_json(json_path, json_data) + + write_image_xml( + xml_path=xml_path, + image_filename=image_path.name, + image_path=str(image_path), + width=int(predict_output.get("image_width", img_bgr.shape[1])), + height=int(predict_output.get("image_height", img_bgr.shape[0])), + results=results, + ) + + return { + "result_dir": str(result_dir), + "image": str(image_path), + "json": str(json_path), + "xml": str(xml_path), + } \ No newline at end of file diff --git a/app/inference/__init__.py b/app/inference/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/app/inference/convnext_onnx.py b/app/inference/convnext_onnx.py new file mode 100644 index 0000000..2f4d172 --- /dev/null +++ b/app/inference/convnext_onnx.py @@ -0,0 +1,101 @@ +from typing import List, Dict +import cv2 +import numpy as np +import onnxruntime as ort + + +def softmax(x: np.ndarray) -> np.ndarray: + x = x.astype(np.float32) + x = x - np.max(x) + exp = np.exp(x) + return exp / (np.sum(exp) + 1e-12) + + +def to_probabilities(output: np.ndarray) -> np.ndarray: + """ + 兼容两种常见 ONNX 导出： + 1. 模型直接输出 logits，需要 softmax； + 2. 模型图中已经包含 softmax，直接输出概率。 + + 对概率再次做 softmax 会把原本很高的置信度压到接近均匀分布。 + """ + values = np.asarray(output, dtype=np.float32).reshape(-1) + total = float(np.sum(values)) + is_probability_output = ( + values.size > 0 + and np.all(np.isfinite(values)) + and float(np.min(values)) >= -1e-5 + and float(np.max(values)) <= 1.0 + 1e-5 + and abs(total - 1.0) <= 1e-3 + ) + if is_probability_output: + return values / max(total, 1e-12) + return softmax(values) + + +class ConvNeXtONNXClassifier: + def __init__( + self, + model_path: str, + class_names: List[str], + img_size: int = 224, + device: str = "cpu", + ): + self.model_path = model_path + self.class_names = class_names + self.img_size = img_size + + providers = ["CPUExecutionProvider"] + if device.lower() in ["cuda", "gpu"]: + providers = ["CUDAExecutionProvider", "CPUExecutionProvider"] + + self.session = ort.InferenceSession(model_path, providers=providers) + active_providers = self.session.get_providers() + print("[ConvNeXt] providers:", active_providers) + if ( + device.lower() in ["cuda", "gpu"] + and "CUDAExecutionProvider" not in active_providers + ): + raise RuntimeError( + "请求使用 CUDA 推理，但 ConvNeXt 会话未启用 " + "CUDAExecutionProvider。请安装匹配 CUDA/cuDNN 的 " + "onnxruntime-gpu。" + ) + self.input_name = self.session.get_inputs()[0].name + self.output_names = [o.name for o in self.session.get_outputs()] + + # 常见 ImageNet 归一化参数。若训练 ConvNeXt 时使用了其他 mean/std，需要保持一致。 + self.mean = np.array([0.485, 0.456, 0.406], dtype=np.float32) + self.std = np.array([0.229, 0.224, 0.225], dtype=np.float32) + + def preprocess(self, img_bgr: np.ndarray) -> np.ndarray: + img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB) + img_rgb = cv2.resize(img_rgb, (self.img_size, self.img_size), interpolation=cv2.INTER_LINEAR) + img = img_rgb.astype(np.float32) / 255.0 + img = (img - self.mean) / self.std + img = img.transpose(2, 0, 1) + img = np.expand_dims(img, axis=0).astype(np.float32) + return img + + def infer(self, img_bgr: np.ndarray) -> Dict: + input_tensor = self.preprocess(img_bgr) + outputs = self.session.run(self.output_names, {self.input_name: input_tensor}) + logits = np.asarray(outputs[0]) + + if logits.ndim == 2: + logits = logits[0] + elif logits.ndim == 1: + pass + else: + logits = logits.reshape(-1) + + probs = to_probabilities(logits) + class_id = int(np.argmax(probs)) + confidence = float(probs[class_id]) + class_name = self.class_names[class_id] if 0 <= class_id < len(self.class_names) else str(class_id) + + return { + "cls_class_id": class_id, + "cls_class_name": class_name, + "cls_confidence": confidence, + } diff --git a/app/inference/density_onnx.py b/app/inference/density_onnx.py new file mode 100644 index 0000000..ff3eee0 --- /dev/null +++ b/app/inference/density_onnx.py @@ -0,0 +1,90 @@ +import cv2 +import numpy as np +import onnxruntime as ort + + +class DensityEstimatorONNX: + def __init__(self, model_path, input_size=None, device="cpu", providers=None): + self.model_path = model_path + self.input_size = input_size + + if providers is None: + providers = ["CPUExecutionProvider"] + if device.lower() in ["cuda", "gpu"]: + providers = ["CUDAExecutionProvider", "CPUExecutionProvider"] + + self.session = ort.InferenceSession(model_path, providers=providers) + active_providers = self.session.get_providers() + print("[Density] providers:", active_providers) + if ( + device.lower() in ["cuda", "gpu"] + and "CUDAExecutionProvider" not in active_providers + ): + raise RuntimeError( + "请求使用 CUDA 推理，但密度模型会话未启用 " + "CUDAExecutionProvider。请安装匹配 CUDA/cuDNN 的 " + "onnxruntime-gpu。" + ) + + self.input_name = self.session.get_inputs()[0].name + self.input_shape = self.session.get_inputs()[0].shape + self.output_names = [o.name for o in self.session.get_outputs()] + + print("[Density] input_name:", self.input_name) + print("[Density] input_shape:", self.input_shape) + print("[Density] output_names:", self.output_names) + + if self.input_size is None: + h = self.input_shape[2] + w = self.input_shape[3] + if isinstance(h, int) and isinstance(w, int): + self.input_size = (w, h) + else: + self.input_size = (400, 400) + + def preprocess(self, image_bgr): + original_h, original_w = image_bgr.shape[:2] + + image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB) + resized = cv2.resize(image_rgb, self.input_size) + + img = resized.astype(np.float32) / 255.0 + + mean = np.array([0.485, 0.456, 0.406], dtype=np.float32) + std = np.array([0.229, 0.224, 0.225], dtype=np.float32) + img = (img - mean) / std + + img = img.transpose(2, 0, 1) + img = np.expand_dims(img, axis=0).astype(np.float32) + + return img, original_w, original_h + + def infer(self, image_bgr): + input_tensor, original_w, original_h = self.preprocess(image_bgr) + + outputs = self.session.run(self.output_names, {self.input_name: input_tensor}) + density = outputs[0] + + density = np.squeeze(density).astype(np.float32) + + if density.ndim != 2: + density = density.reshape(density.shape[-2], density.shape[-1]) + + count = float(density.sum()) + + density_map = cv2.resize(density, (original_w, original_h)) + + return { + "count": count, + "density_map": density_map + } + + @staticmethod + def colorize_density_map(density_map): + d = density_map.copy() + d = d - d.min() + if d.max() > 0: + d = d / d.max() + d = (d * 255).astype(np.uint8) + heatmap = cv2.applyColorMap(d, cv2.COLORMAP_JET) + return heatmap diff --git a/app/inference/yolo_onnx.py b/app/inference/yolo_onnx.py new file mode 100644 index 0000000..ddadae9 --- /dev/null +++ b/app/inference/yolo_onnx.py @@ -0,0 +1,216 @@ +from typing import List, Dict +import cv2 +import numpy as np +import onnxruntime as ort + + +def letterbox( + img: np.ndarray, + new_shape=(960, 960), + color=(114, 114, 114), + auto=False, + scale_fill=False, + scale_up=True, + stride=32, +): + shape = img.shape[:2] # h, w + + if isinstance(new_shape, int): + new_shape = (new_shape, new_shape) + + r = min(new_shape[0] / shape[0], new_shape[1] / shape[1]) + if not scale_up: + r = min(r, 1.0) + + ratio = (r, r) + new_unpad = (int(round(shape[1] * r)), int(round(shape[0] * r))) + + dw = new_shape[1] - new_unpad[0] + dh = new_shape[0] - new_unpad[1] + + if auto: + dw, dh = np.mod(dw, stride), np.mod(dh, stride) + elif scale_fill: + dw, dh = 0.0, 0.0 + new_unpad = (new_shape[1], new_shape[0]) + ratio = (new_shape[1] / shape[1], new_shape[0] / shape[0]) + + dw /= 2 + dh /= 2 + + if shape[::-1] != new_unpad: + img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR) + + top = int(round(dh - 0.1)) + bottom = int(round(dh + 0.1)) + left = int(round(dw - 0.1)) + right = int(round(dw + 0.1)) + + img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color) + return img, ratio, (dw, dh) + + +def xywh2xyxy(x: np.ndarray) -> np.ndarray: + y = np.copy(x) + y[:, 0] = x[:, 0] - x[:, 2] / 2 + y[:, 1] = x[:, 1] - x[:, 3] / 2 + y[:, 2] = x[:, 0] + x[:, 2] / 2 + y[:, 3] = x[:, 1] + x[:, 3] / 2 + return y + + +def box_iou(box: np.ndarray, boxes: np.ndarray) -> np.ndarray: + x1 = np.maximum(box[0], boxes[:, 0]) + y1 = np.maximum(box[1], boxes[:, 1]) + x2 = np.minimum(box[2], boxes[:, 2]) + y2 = np.minimum(box[3], boxes[:, 3]) + + inter = np.maximum(0, x2 - x1) * np.maximum(0, y2 - y1) + area1 = max(0, box[2] - box[0]) * max(0, box[3] - box[1]) + area2 = np.maximum(0, boxes[:, 2] - boxes[:, 0]) * np.maximum(0, boxes[:, 3] - boxes[:, 1]) + union = area1 + area2 - inter + 1e-7 + return inter / union + + +def nms_numpy(boxes: np.ndarray, scores: np.ndarray, iou_thres: float) -> List[int]: + if len(boxes) == 0: + return [] + + order = scores.argsort()[::-1] + keep = [] + + while order.size > 0: + i = order[0] + keep.append(int(i)) + if order.size == 1: + break + ious = box_iou(boxes[i], boxes[order[1:]]) + inds = np.where(ious <= iou_thres)[0] + order = order[inds + 1] + + return keep + + +def scale_boxes_from_letterbox(boxes: np.ndarray, ratio, pad, original_shape) -> np.ndarray: + # boxes: xyxy in letterboxed image + gain = ratio[0] + pad_x, pad_y = pad + + boxes[:, [0, 2]] -= pad_x + boxes[:, [1, 3]] -= pad_y + boxes[:, :4] /= gain + + h, w = original_shape[:2] + boxes[:, [0, 2]] = boxes[:, [0, 2]].clip(0, w - 1) + boxes[:, [1, 3]] = boxes[:, [1, 3]].clip(0, h - 1) + return boxes + + +class YOLOv5ONNXDetector: + def __init__( + self, + model_path: str, + class_names: List[str], + img_size: int = 960, + device: str = "cpu", + ): + self.model_path = model_path + self.class_names = class_names + self.img_size = img_size + + providers = ["CPUExecutionProvider"] + if device.lower() in ["cuda", "gpu"]: + providers = ["CUDAExecutionProvider", "CPUExecutionProvider"] + + self.session = ort.InferenceSession(model_path, providers=providers) + active_providers = self.session.get_providers() + print("[YOLO] providers:", active_providers) + if ( + device.lower() in ["cuda", "gpu"] + and "CUDAExecutionProvider" not in active_providers + ): + raise RuntimeError( + "请求使用 CUDA 推理，但 YOLO 会话未启用 " + "CUDAExecutionProvider。请安装匹配 CUDA/cuDNN 的 " + "onnxruntime-gpu。" + ) + self.input_name = self.session.get_inputs()[0].name + self.output_names = [o.name for o in self.session.get_outputs()] + + def preprocess(self, img_bgr: np.ndarray): + img_lb, ratio, pad = letterbox(img_bgr, new_shape=(self.img_size, self.img_size), auto=False) + img_rgb = cv2.cvtColor(img_lb, cv2.COLOR_BGR2RGB) + img = img_rgb.transpose(2, 0, 1).astype(np.float32) / 255.0 + img = np.expand_dims(img, axis=0) + return img, ratio, pad + + def infer( + self, + img_bgr: np.ndarray, + conf_thres: float = 0.25, + iou_thres: float = 0.45, + max_det: int = 50, + ) -> List[Dict]: + input_tensor, ratio, pad = self.preprocess(img_bgr) + outputs = self.session.run(self.output_names, {self.input_name: input_tensor}) + + pred = outputs[0] + + # 常见 YOLOv5 ONNX 输出: [1, num_boxes, 5 + num_classes] + if isinstance(pred, list): + pred = pred[0] + pred = np.asarray(pred) + + if pred.ndim == 3: + pred = pred[0] + elif pred.ndim == 2: + pass + else: + raise ValueError(f"不支持的 YOLO 输出维度: {pred.shape}") + + if pred.shape[1] < 6: + raise ValueError(f"YOLO 输出最后一维应至少为 6，当前为: {pred.shape}") + + boxes_xywh = pred[:, :4] + obj_conf = pred[:, 4] + + if pred.shape[1] == 6: + # 某些导出格式可能已经是 x,y,w,h,conf,class_id + cls_ids = pred[:, 5].astype(np.int32) + scores = obj_conf + else: + cls_scores = pred[:, 5:] + cls_ids = np.argmax(cls_scores, axis=1) + cls_conf = cls_scores[np.arange(len(cls_scores)), cls_ids] + scores = obj_conf * cls_conf + + mask = scores >= conf_thres + if not np.any(mask): + return [] + + boxes_xywh = boxes_xywh[mask] + scores = scores[mask] + cls_ids = cls_ids[mask] + + boxes = xywh2xyxy(boxes_xywh) + boxes = scale_boxes_from_letterbox(boxes, ratio, pad, img_bgr.shape) + + final_results = [] + + for class_id in np.unique(cls_ids): + idxs = np.where(cls_ids == class_id)[0] + keep = nms_numpy(boxes[idxs], scores[idxs], iou_thres) + + for k in keep: + original_idx = idxs[k] + cid = int(class_id) + cname = self.class_names[cid] if 0 <= cid < len(self.class_names) else str(cid) + final_results.append({ + "bbox_xyxy": boxes[original_idx].astype(float).tolist(), + "det_confidence": float(scores[original_idx]), + "det_class_id": cid, + "det_class_name": cname, + }) + + final_results.sort(key=lambda x: x["det_confidence"], reverse=True) + return final_results[:max_det] diff --git a/app/main.py b/app/main.py new file mode 100644 index 0000000..b44613d --- /dev/null +++ b/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)) \ No newline at end of file diff --git a/app/pipeline.py b/app/pipeline.py new file mode 100644 index 0000000..cf39efa --- /dev/null +++ b/app/pipeline.py @@ -0,0 +1,794 @@ +from typing import Dict, List +import os +import cv2 +import numpy as np +from PIL import Image, ImageDraw, ImageFont + +from app.config import settings, BASE_DIR +from app.utils import read_classes, crop_with_padding, encode_image_to_base64 +from app.inference.yolo_onnx import YOLOv5ONNXDetector +from app.inference.convnext_onnx import ConvNeXtONNXClassifier +from app.inference.density_onnx import DensityEstimatorONNX + + +class BirdCascadePipeline: + def __init__(self): + self.det_classes = [] + self.cls_classes = [] + self.detector = None + self.classifier = None + self.binary_dark_species_name = getattr( + settings, + "binary_dark_species_name", + "", + ).strip() + self.binary_light_species_name = getattr( + settings, + "binary_light_species_name", + "", + ).strip() + self.use_binary_species_rule = bool( + self.binary_dark_species_name + and self.binary_light_species_name + ) + + if settings.enable_detection: + self.det_classes = read_classes( + settings.det_classes_path, + keep_empty=False, + ) + + if not os.path.exists(settings.det_model_path): + raise FileNotFoundError( + f"未找到 YOLOv5 模型: {settings.det_model_path}" + ) + + self.detector = YOLOv5ONNXDetector( + model_path=settings.det_model_path, + class_names=self.det_classes, + img_size=settings.det_img_size, + device=settings.device, + ) + + if settings.enable_classification or self.use_binary_species_rule: + self.cls_classes = read_classes( + settings.cls_classes_path, + keep_empty=True, + ) + + if settings.enable_classification and not self.use_binary_species_rule: + if not os.path.exists(settings.cls_model_path): + raise FileNotFoundError( + f"未找到 ConvNeXt 模型: {settings.cls_model_path}" + ) + + self.classifier = ConvNeXtONNXClassifier( + model_path=settings.cls_model_path, + class_names=self.cls_classes, + img_size=settings.cls_img_size, + device=settings.device, + ) + + self.density_estimator = None + if settings.enable_density: + if not os.path.exists(settings.density_model_path): + raise FileNotFoundError(f"未找到密度估计模型: {settings.density_model_path}") + + density_input_size = None + if settings.density_img_size and settings.density_img_size > 0: + density_input_size = (settings.density_img_size, settings.density_img_size) + + self.density_estimator = DensityEstimatorONNX( + model_path=settings.density_model_path, + input_size=density_input_size, + device=settings.device, + ) + + self.target_det_class_ids = settings.target_det_class_ids + self.species_name_remap = {} + remap_text = getattr(settings, "species_name_remap", "").replace("，", ",") + for rule in remap_text.split(","): + rule = rule.strip() + if not rule: + continue + if "=" not in rule: + raise ValueError( + "SPECIES_NAME_REMAP/--species-name-remap 格式错误: " + f"{rule!r}，应为 原标签=新标签" + ) + source_name, target_name = [ + value.strip() + for value in rule.split("=", 1) + ] + if not source_name or not target_name: + raise ValueError( + "物种标签映射的原标签和新标签都不能为空" + ) + self.species_name_remap[source_name] = target_name + allowed_text = getattr( + settings, + "allowed_species_names", + "", + ).replace("，", ",") + self.allowed_species_names = { + name.strip() + for name in allowed_text.split(",") + if name.strip() + } + self.font_path = self._find_chinese_font() + + def _find_chinese_font(self) -> str: + candidates = [ + os.path.join(BASE_DIR, "assets", "fonts", "NotoSansCJK-Regular.ttc"), + os.path.join(BASE_DIR, "assets", "fonts", "NotoSansSC-Regular.otf"), + os.path.join(BASE_DIR, "assets", "fonts", "SimHei.ttf"), + os.path.join(BASE_DIR, "assets", "fonts", "msyh.ttc"), + "/usr/share/fonts/opentype/noto/NotoSansCJK-Regular.ttc", + "/usr/share/fonts/truetype/wqy/wqy-microhei.ttc", + "C:/Windows/Fonts/msyh.ttc", + "C:/Windows/Fonts/simhei.ttf", + ] + + for p in candidates: + if os.path.exists(p): + print(f"[Font] 使用中文字体: {p}") + return p + + raise FileNotFoundError( + f"未找到中文字体，请将 NotoSansCJK-Regular.ttc 放到 " + f"{os.path.join(BASE_DIR, 'assets', 'fonts')} 下。" + ) + + def _get_font_by_target_height(self, target_height: int) -> ImageFont.FreeTypeFont: + lo, hi = 8, 180 + best_size = 28 + sample_text = "绿头鸭" + + while lo <= hi: + mid = (lo + hi) // 2 + font = ImageFont.truetype(self.font_path, mid) + bbox = font.getbbox(sample_text) + text_h = bbox[3] - bbox[1] + + if text_h <= target_height: + best_size = mid + lo = mid + 1 + else: + hi = mid - 1 + + return ImageFont.truetype(self.font_path, best_size) + + def _get_color_by_class_id(self, class_id: int) -> tuple[int, int, int]: + if settings.uniform_species_color: + return (0, 190, 0) + + palette = [ + (0, 190, 0), + (220, 60, 60), + (70, 130, 255), + (255, 160, 0), + (180, 80, 220), + (0, 180, 180), + (255, 80, 180), + (120, 180, 40), + (160, 120, 60), + (80, 80, 220), + (255, 210, 0), + (0, 120, 255), + (255, 90, 90), + (90, 220, 160), + (200, 80, 140), + ] + + if class_id < 0: + return (155, 155, 155) + + return palette[class_id % len(palette)] + + def _filter_target_detections(self, detections: List[Dict], img_w: int, img_h: int) -> List[Dict]: + filtered = [] + img_area = float(img_w * img_h) + + det_second_conf_thres = getattr(settings, "det_second_conf_thres", settings.det_conf_thres) + min_box_area_ratio = getattr(settings, "min_box_area_ratio", 0.00002) + max_box_area_ratio = getattr(settings, "max_box_area_ratio", 1.0) + max_box_aspect_ratio = getattr(settings, "max_box_aspect_ratio", 6.0) + roi_y_min = float( + getattr(settings, "detection_roi_y_min_ratio", 0.0) + ) + roi_y_max = float( + getattr(settings, "detection_roi_y_max_ratio", 1.0) + ) + + for det in detections: + det_class_id = int(det.get("det_class_id", -1)) + det_conf = float(det.get("det_confidence", 0.0)) + + if self.target_det_class_ids and det_class_id not in self.target_det_class_ids: + continue + + if det_conf < det_second_conf_thres: + continue + + x1, y1, x2, y2 = [float(v) for v in det["bbox_xyxy"]] + box_w = max(0.0, x2 - x1) + box_h = max(0.0, y2 - y1) + + if box_w <= 1 or box_h <= 1: + continue + + area_ratio = (box_w * box_h) / img_area + + if area_ratio < min_box_area_ratio: + continue + + if area_ratio > max_box_area_ratio: + continue + + box_ratio = max(box_w / box_h, box_h / box_w) + if box_ratio > max_box_aspect_ratio: + continue + + center_y_ratio = ((y1 + y2) * 0.5) / float(img_h) + if center_y_ratio < roi_y_min or center_y_ratio > roi_y_max: + continue + + filtered.append(det) + + return filtered + + def _draw_count_text(self, img_bgr: np.ndarray, count: int) -> np.ndarray: + h, w = img_bgr.shape[:2] + short_side = min(h, w) + + target_text_height = int(short_side * 0.035) + target_text_height = max(30, min(target_text_height, 48)) + + font = self._get_font_by_target_height(target_text_height) + text = f"数量：{int(count)}" + + img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB) + pil_img = Image.fromarray(img_rgb) + draw = ImageDraw.Draw(pil_img) + + bbox = draw.textbbox((0, 0), text, font=font, stroke_width=2) + text_h = bbox[3] - bbox[1] + + margin_x = max(18, min(int(short_side * 0.018), 36)) + margin_y = max(18, min(int(short_side * 0.018), 36)) + + x = margin_x + y = h - margin_y - text_h + + draw.text( + (x, y), + text, + font=font, + fill=(255, 255, 255), + stroke_width=2, + stroke_fill=(0, 0, 0), + ) + + return cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR) + + def _draw_detection_results( + self, + img_bgr: np.ndarray, + results: List[Dict], + is_dense_scene: bool = False, + ) -> np.ndarray: + out = img_bgr.copy() + h, w = out.shape[:2] + short_side = min(h, w) + + box_thick = max(2, int(short_side * 0.0022)) + + # 密集场景中每只鸟仍显示自己的标签，因此字号保持紧凑， + # 避免标签底色遮住相邻目标。 + label_target_height = int(short_side * 0.014) + label_target_height = max(14, min(label_target_height, 24)) + label_font = self._get_font_by_target_height(label_target_height) + + img_rgb = cv2.cvtColor(out, cv2.COLOR_BGR2RGB) + pil_img = Image.fromarray(img_rgb) + draw = ImageDraw.Draw(pil_img) + + hide_in_dense = getattr(settings, "hide_cls_name_in_dense_scene", True) + for item in results: + x1, y1, x2, y2 = [int(round(v)) for v in item["bbox_xyxy"]] + + x1 = max(0, min(x1, w - 1)) + y1 = max(0, min(y1, h - 1)) + x2 = max(0, min(x2, w - 1)) + y2 = max(0, min(y2, h - 1)) + + cls_id = int(item.get("cls_class_id", -1)) + cls_name = str(item.get("cls_class_name", "") or "").strip() + + color = self._get_color_by_class_id(cls_id) + + draw.rectangle( + [(x1, y1), (x2, y2)], + outline=color, + width=box_thick, + ) + + if is_dense_scene and hide_in_dense: + continue + + if cls_name == "": + continue + + if not bool(item.get("show_cls_name", False)): + continue + + text = cls_name + + text_bbox = draw.textbbox((0, 0), text, font=label_font, stroke_width=0) + text_w = text_bbox[2] - text_bbox[0] + text_h = text_bbox[3] - text_bbox[1] + + pad_x = max(3, int(label_target_height * 0.18)) + pad_y = max(2, int(label_target_height * 0.10)) + + label_width = text_w + 2 * pad_x + label_height = text_h + 2 * pad_y + + tx = x1 + ty = y1 - label_height + + if ty < 0: + ty = y1 + + if tx + label_width > w: + tx = max(0, w - label_width) + + draw.rectangle( + [ + (tx, ty), + (tx + label_width, ty + label_height), + ], + fill=color, + ) + + text_x = tx + (label_width - text_w) // 2 - text_bbox[0] + text_y = ty + (label_height - text_h) // 2 - text_bbox[1] + + draw.text( + (text_x, text_y), + text, + font=label_font, + fill=(0, 0, 0), + ) + + return cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR) + + def _draw_density_panel(self, img_bgr: np.ndarray, density_result: Dict | None) -> np.ndarray: + if density_result is None or "density_map" not in density_result: + return img_bgr + + out = img_bgr.copy() + h, w = out.shape[:2] + short_side = min(h, w) + + panel_w = int(w * 0.28) + panel_h = int(h * 0.24) + + panel_w = max(300, min(panel_w, int(w * 0.32))) + panel_h = max(180, min(panel_h, int(h * 0.28))) + + panel_w = min(panel_w, w - 40) + panel_h = min(panel_h, h - 40) + + margin = max(18, min(int(short_side * 0.018), 36)) + + x1 = w - panel_w - margin + y1 = h - panel_h - margin + x2 = w - margin + y2 = h - margin + + density_map = density_result["density_map"] + heatmap = DensityEstimatorONNX.colorize_density_map(density_map) + heatmap_small = cv2.resize(heatmap, (panel_w, panel_h)) + + roi = out[y1:y2, x1:x2] + blended = cv2.addWeighted(roi, 0.25, heatmap_small, 0.75, 0) + out[y1:y2, x1:x2] = blended + + border_thick = max(2, int(short_side * 0.002)) + cv2.rectangle(out, (x1, y1), (x2, y2), (255, 255, 255), border_thick) + + count = int(round(float(density_result.get("count", 0.0)))) + out = self._draw_count_text(out, count) + + return out + + def draw_full_result( + self, + img_bgr: np.ndarray, + results: List[Dict], + density_result: Dict | None = None, + ) -> np.ndarray: + count = 0 + if density_result is not None: + count = int(round(float(density_result.get("count", 0.0)))) + + dense_scene_count_thres = getattr(settings, "dense_scene_count_thres", 50) + is_dense_scene = count >= dense_scene_count_thres + + annotated = self._draw_detection_results( + img_bgr, + results, + is_dense_scene=is_dense_scene, + ) + + annotated = self._draw_density_panel(annotated, density_result) + + return annotated + + def predict( + self, + img_bgr: np.ndarray, + return_image: bool = False, + det_conf: float | None = None, + iou_thres: float | None = None, + ) -> Dict: + h, w = img_bgr.shape[:2] + + det_conf = settings.det_conf_thres if det_conf is None else det_conf + iou_thres = settings.det_iou_thres if iou_thres is None else iou_thres + + density_result = None + density_count = 0 + + if self.density_estimator is not None: + density_result = self.density_estimator.infer(img_bgr) + density_count = int(round(float(density_result.get("count", 0.0)))) + + dense_scene_count_thres = getattr(settings, "dense_scene_count_thres", 50) + is_dense_scene = density_count >= dense_scene_count_thres + + detections = [] + if self.detector is not None: + detections = self.detector.infer( + img_bgr, + conf_thres=det_conf, + iou_thres=iou_thres, + max_det=settings.max_detections, + ) + + detections = self._filter_target_detections( + detections, + img_w=w, + img_h=h, + ) + + results: List[Dict] = [] + + for det in detections: + cls_id = -1 + cls_name = "" + cls_conf = 0.0 + + if self.use_binary_species_rule: + crop = crop_with_padding( + img_bgr, + det["bbox_xyxy"], + padding_ratio=0.0, + ) + gray = cv2.cvtColor(crop, cv2.COLOR_BGR2GRAY) + dark_pixel_threshold = int( + getattr(settings, "binary_dark_pixel_threshold", 85) + ) + dark_ratio_threshold = float( + getattr(settings, "binary_dark_ratio_threshold", 0.12) + ) + dark_ratio = float( + np.mean(gray < dark_pixel_threshold) + ) + if dark_ratio >= dark_ratio_threshold: + cls_name = self.binary_dark_species_name + else: + cls_name = self.binary_light_species_name + if cls_name in self.cls_classes: + cls_id = self.cls_classes.index(cls_name) + cls_conf = 1.0 + elif self.classifier is not None: + crop = crop_with_padding( + img_bgr, + det["bbox_xyxy"], + padding_ratio=settings.crop_padding_ratio, + ) + + cls_result = self.classifier.infer(crop) + + cls_id = int(cls_result["cls_class_id"]) + cls_name = str(cls_result["cls_class_name"] or "").strip() + cls_conf = float(cls_result["cls_confidence"]) + + remap_min_confidence = float( + getattr( + settings, + "species_name_remap_min_confidence", + 0.0, + ) + ) + if ( + cls_name in self.species_name_remap + and cls_conf >= remap_min_confidence + ): + cls_name = self.species_name_remap[cls_name] + if cls_name in self.cls_classes: + cls_id = self.cls_classes.index(cls_name) + cls_conf = 1.0 + + if ( + self.allowed_species_names + and cls_name not in self.allowed_species_names + ): + cls_id = -1 + cls_name = "" + cls_conf = 0.0 + + x1, y1, x2, y2 = [float(v) for v in det["bbox_xyxy"]] + box_area_ratio = ((x2 - x1) * (y2 - y1)) / float(w * h) + + dark_override_species_name = getattr( + settings, + "dark_override_species_name", + "", + ).strip() + if dark_override_species_name: + crop = crop_with_padding( + img_bgr, + det["bbox_xyxy"], + padding_ratio=0.0, + ) + gray = cv2.cvtColor(crop, cv2.COLOR_BGR2GRAY) + dark_pixel_threshold = int( + getattr(settings, "dark_override_pixel_threshold", 80) + ) + dark_ratio_threshold = float( + getattr(settings, "dark_override_ratio_threshold", 0.35) + ) + min_dark_area = float( + getattr(settings, "dark_override_min_box_area_ratio", 0.0) + ) + max_dark_area = float( + getattr(settings, "dark_override_max_box_area_ratio", 1.0) + ) + dark_roi_y_min = float( + getattr(settings, "dark_override_roi_y_min_ratio", 0.0) + ) + dark_roi_y_max = float( + getattr(settings, "dark_override_roi_y_max_ratio", 1.0) + ) + center_y_ratio = ((y1 + y2) * 0.5) / float(h) + dark_ratio = float(np.mean(gray < dark_pixel_threshold)) + if ( + dark_ratio >= dark_ratio_threshold + and box_area_ratio >= min_dark_area + and box_area_ratio <= max_dark_area + and center_y_ratio >= dark_roi_y_min + and center_y_ratio <= dark_roi_y_max + ): + dark_class_id = int( + getattr(settings, "dark_override_species_class_id", -1) + ) + if ( + dark_class_id < 0 + and dark_override_species_name in self.cls_classes + ): + dark_class_id = self.cls_classes.index( + dark_override_species_name + ) + cls_id = dark_class_id + cls_name = dark_override_species_name + cls_conf = 1.0 + + forced_species_name = getattr( + settings, + "forced_species_name", + "", + ).strip() + forced_min_area = float( + getattr( + settings, + "forced_species_min_box_area_ratio", + 0.0, + ) + ) + if ( + forced_species_name + and box_area_ratio >= forced_min_area + ): + forced_class_id = int( + getattr(settings, "forced_species_class_id", -1) + ) + if forced_class_id < 0 and forced_species_name in self.cls_classes: + forced_class_id = self.cls_classes.index(forced_species_name) + cls_id = forced_class_id + cls_name = forced_species_name + cls_conf = 1.0 + + region_override_species_name = getattr( + settings, + "region_override_species_name", + "", + ).strip() + if region_override_species_name and cls_name: + region_y_min = float( + getattr(settings, "region_override_roi_y_min_ratio", 0.0) + ) + region_y_max = float( + getattr(settings, "region_override_roi_y_max_ratio", 1.0) + ) + center_y_ratio = ((y1 + y2) * 0.5) / float(h) + source_text = getattr( + settings, + "region_override_source_species_names", + "", + ).replace("，", ",") + source_names = { + value.strip() + for value in source_text.split(",") + if value.strip() + } + if ( + center_y_ratio >= region_y_min + and center_y_ratio <= region_y_max + and (not source_names or cls_name in source_names) + ): + region_class_id = int( + getattr(settings, "region_override_species_class_id", -1) + ) + if ( + region_class_id < 0 + and region_override_species_name in self.cls_classes + ): + region_class_id = self.cls_classes.index( + region_override_species_name + ) + cls_id = region_class_id + cls_name = region_override_species_name + cls_conf = 1.0 + + crane_names = {"白鹤", "白头鹤", "白枕鹤"} + if ( + getattr(settings, "enable_crane_appearance_correction", False) + and cls_name in crane_names + ): + crop = crop_with_padding( + img_bgr, + det["bbox_xyxy"], + padding_ratio=0.0, + ) + gray = cv2.cvtColor(crop, cv2.COLOR_BGR2GRAY) + white_ratio = float(np.mean(gray > 168)) + dark_ratio = float(np.mean(gray < 78)) + crop_h, crop_w = gray.shape[:2] + lower_white_ratio = float( + np.mean(gray[crop_h // 2 :, :] > 168) + ) + left_white_ratio = float( + np.mean(gray[:, : max(1, crop_w // 3)] > 168) + ) + right_white_ratio = float( + np.mean(gray[:, max(0, crop_w * 2 // 3) :] > 168) + ) + tail_white_ratio = max( + lower_white_ratio, + left_white_ratio, + right_white_ratio, + ) + white_thres = float( + getattr(settings, "crane_white_ratio_threshold", 0.30) + ) + dark_thres = float( + getattr(settings, "crane_dark_ratio_threshold", 0.28) + ) + tail_white_thres = float( + getattr( + settings, + "crane_tail_white_ratio_threshold", + 0.16, + ) + ) + + corrected_name = "" + if white_ratio >= white_thres and dark_ratio < dark_thres: + corrected_name = "白鹤" + elif ( + white_ratio < white_thres + and tail_white_ratio >= tail_white_thres + ): + corrected_name = "白枕鹤" + elif dark_ratio >= dark_thres and white_ratio < white_thres: + corrected_name = "白头鹤" + elif not getattr(settings, "crane_uncertain_hide_label", True): + corrected_name = cls_name + + if corrected_name: + cls_name = corrected_name + if cls_name in self.cls_classes: + cls_id = self.cls_classes.index(cls_name) + cls_conf = max(cls_conf, 0.90) + else: + cls_id = -1 + cls_name = "" + cls_conf = 0.0 + + normal_show_thres = getattr(settings, "cls_show_conf_thres", 0.60) + large_box_area_ratio = getattr(settings, "large_box_area_ratio", 0.002) + large_box_show_thres = getattr(settings, "large_box_cls_show_conf_thres", 0.25) + show_large_always = getattr(settings, "show_large_box_name_always", True) + hide_in_dense = getattr(settings, "hide_cls_name_in_dense_scene", True) + generic_class_names = { + "鸟", + "bird", + "unknown", + "未知", + "未知鸟类", + } + + is_large_box = box_area_ratio >= large_box_area_ratio + + if is_dense_scene and hide_in_dense: + show_cls_name = False + elif cls_name == "" or cls_name.lower() in generic_class_names: + show_cls_name = False + elif is_large_box and show_large_always: + show_cls_name = True + else: + show_cls_name = ( + cls_conf >= normal_show_thres + or + (is_large_box and cls_conf >= large_box_show_thres) + ) + + merged = { + "bbox_xyxy": [round(float(v), 2) for v in det["bbox_xyxy"]], + "det_confidence": round(float(det["det_confidence"]), 6), + "det_class_id": int(det["det_class_id"]), + "det_class_name": str(det["det_class_name"]), + "cls_class_id": cls_id, + "cls_class_name": cls_name, + "cls_confidence": round(cls_conf, 6), + "box_area_ratio": round(float(box_area_ratio), 8), + "is_large_box": bool(is_large_box), + "is_dense_scene": bool(is_dense_scene), + "show_cls_name": bool(show_cls_name), + } + + results.append(merged) + + density_json = None + + if density_result is not None: + density_json = { + "count": int(round(float(density_result.get("count", 0.0)))), + "is_dense_scene": bool(is_dense_scene), + } + + annotated_image_base64 = None + + if return_image: + annotated = self.draw_full_result( + img_bgr, + results, + density_result, + ) + annotated_image_base64 = encode_image_to_base64(annotated) + + return { + "success": True, + "image_width": w, + "image_height": h, + "num_objects": len(results), + "target_det_class_ids": self.target_det_class_ids, + "density": density_json, + "density_result": density_result, + "results": results, + "annotated_image_base64": annotated_image_base64, + "message": "ok", + } diff --git a/app/pipeline_add_density.py b/app/pipeline_add_density.py new file mode 100644 index 0000000..02a8216 --- /dev/null +++ b/app/pipeline_add_density.py @@ -0,0 +1,309 @@ +from typing import Dict, List +import os +import cv2 +import numpy as np +from PIL import Image, ImageDraw, ImageFont + +from app.config import settings, BASE_DIR +from app.utils import read_classes, crop_with_padding, draw_results, encode_image_to_base64 +from app.inference.yolo_onnx import YOLOv5ONNXDetector +from app.inference.convnext_onnx import ConvNeXtONNXClassifier +from app.inference.density_onnx import DensityEstimatorONNX + + +class BirdCascadePipeline: + def __init__(self): + self.det_classes = read_classes(settings.det_classes_path, keep_empty=False) + self.cls_classes = read_classes(settings.cls_classes_path, keep_empty=True) + + if not os.path.exists(settings.det_model_path): + raise FileNotFoundError(f"未找到 YOLOv5 模型: {settings.det_model_path}") + + if not os.path.exists(settings.cls_model_path): + raise FileNotFoundError(f"未找到 ConvNeXt 模型: {settings.cls_model_path}") + + self.detector = YOLOv5ONNXDetector( + model_path=settings.det_model_path, + class_names=self.det_classes, + img_size=settings.det_img_size, + device=settings.device, + ) + + self.classifier = ConvNeXtONNXClassifier( + model_path=settings.cls_model_path, + class_names=self.cls_classes, + img_size=settings.cls_img_size, + device=settings.device, + ) + + self.density_estimator = None + if settings.enable_density: + if not os.path.exists(settings.density_model_path): + raise FileNotFoundError(f"未找到密度估计模型: {settings.density_model_path}") + + density_input_size = None + if settings.density_img_size and settings.density_img_size > 0: + density_input_size = (settings.density_img_size, settings.density_img_size) + + self.density_estimator = DensityEstimatorONNX( + model_path=settings.density_model_path, + input_size=density_input_size, + ) + + self.target_det_class_ids = settings.target_det_class_ids + self.font_path = self._find_chinese_font() + + def _find_chinese_font(self) -> str: + candidates = [ + os.path.join(BASE_DIR, "assets", "fonts", "NotoSansCJK-Regular.ttc"), + os.path.join(BASE_DIR, "assets", "fonts", "NotoSansSC-Regular.otf"), + os.path.join(BASE_DIR, "assets", "fonts", "SimHei.ttf"), + os.path.join(BASE_DIR, "assets", "fonts", "msyh.ttc"), + "/usr/share/fonts/opentype/noto/NotoSansCJK-Regular.ttc", + "/usr/share/fonts/opentype/noto/NotoSansCJK-Regular.otf", + "/usr/share/fonts/truetype/wqy/wqy-microhei.ttc", + "/usr/share/fonts/truetype/wqy/wqy-zenhei.ttc", + "C:/Windows/Fonts/msyh.ttc", + "C:/Windows/Fonts/simhei.ttf", + ] + + for p in candidates: + if os.path.exists(p): + print(f"[Font] 使用中文字体: {p}") + return p + + raise FileNotFoundError( + "未找到中文字体。请将 NotoSansCJK-Regular.ttc 放到 " + f"{os.path.join(BASE_DIR, 'assets', 'fonts')} 下。" + ) + + def _get_font_by_target_height(self, target_height: int) -> ImageFont.FreeTypeFont: + """ + 用目标像素高度反推字号，避免不同图片字体视觉大小不一致。 + """ + lo, hi = 8, 160 + best_size = 28 + sample_text = "数量：888" + + while lo <= hi: + mid = (lo + hi) // 2 + font = ImageFont.truetype(self.font_path, mid) + bbox = font.getbbox(sample_text) + text_h = bbox[3] - bbox[1] + + if text_h <= target_height: + best_size = mid + lo = mid + 1 + else: + hi = mid - 1 + + return ImageFont.truetype(self.font_path, best_size) + + def _draw_count_text(self, img_bgr: np.ndarray, count: int) -> np.ndarray: + h, w = img_bgr.shape[:2] + short_side = min(h, w) + + # 关键：不要直接用 font_size，而是控制最终“字高” + # 0.035 比 0.06 稳定很多；同时设置上下限，避免极端图过大或过小。 + target_text_height = int(short_side * 0.035) + target_text_height = max(30, min(target_text_height, 48)) + + font = self._get_font_by_target_height(target_text_height) + + text = f"数量：{int(count)}" + + img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB) + pil_img = Image.fromarray(img_rgb) + draw = ImageDraw.Draw(pil_img) + + bbox = draw.textbbox((0, 0), text, font=font, stroke_width=2) + text_h = bbox[3] - bbox[1] + + margin_x = int(short_side * 0.018) + margin_y = int(short_side * 0.018) + + margin_x = max(18, min(margin_x, 36)) + margin_y = max(18, min(margin_y, 36)) + + x = margin_x + y = h - margin_y - text_h + + draw.text( + (x, y), + text, + font=font, + fill=(255, 255, 255), + stroke_width=2, + stroke_fill=(0, 0, 0), + ) + + return cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR) + + def _filter_target_detections(self, detections: List[Dict], img_w: int, img_h: int) -> List[Dict]: + filtered = [] + img_area = float(img_w * img_h) + + det_second_conf_thres = getattr(settings, "det_second_conf_thres", settings.det_conf_thres) + min_box_area_ratio = getattr(settings, "min_box_area_ratio", 0.00002) + max_box_area_ratio = getattr(settings, "max_box_area_ratio", 0.08) + max_box_aspect_ratio = getattr(settings, "max_box_aspect_ratio", 6.0) + + for det in detections: + det_class_id = int(det.get("det_class_id", -1)) + det_conf = float(det.get("det_confidence", 0.0)) + + if self.target_det_class_ids and det_class_id not in self.target_det_class_ids: + continue + + if det_conf < det_second_conf_thres: + continue + + x1, y1, x2, y2 = [float(v) for v in det["bbox_xyxy"]] + box_w = max(0.0, x2 - x1) + box_h = max(0.0, y2 - y1) + + if box_w <= 1 or box_h <= 1: + continue + + box_area = box_w * box_h + area_ratio = box_area / img_area + + if area_ratio < min_box_area_ratio: + continue + + if area_ratio > max_box_area_ratio: + continue + + box_ratio = max(box_w / box_h, box_h / box_w) + if box_ratio > max_box_aspect_ratio: + continue + + filtered.append(det) + + return filtered + + def _draw_density_panel(self, img_bgr: np.ndarray, density_result: Dict | None) -> np.ndarray: + if density_result is None or "density_map" not in density_result: + return img_bgr + + out = img_bgr.copy() + h, w = out.shape[:2] + short_side = min(h, w) + + panel_w = int(w * 0.28) + panel_h = int(h * 0.24) + + panel_w = max(300, min(panel_w, int(w * 0.32))) + panel_h = max(180, min(panel_h, int(h * 0.28))) + + panel_w = min(panel_w, w - 40) + panel_h = min(panel_h, h - 40) + + margin = int(short_side * 0.018) + margin = max(18, min(margin, 36)) + + x1 = w - panel_w - margin + y1 = h - panel_h - margin + x2 = w - margin + y2 = h - margin + + density_map = density_result["density_map"] + heatmap = DensityEstimatorONNX.colorize_density_map(density_map) + heatmap_small = cv2.resize(heatmap, (panel_w, panel_h)) + + roi = out[y1:y2, x1:x2] + blended = cv2.addWeighted(roi, 0.25, heatmap_small, 0.75, 0) + out[y1:y2, x1:x2] = blended + + border_thick = max(2, int(short_side * 0.002)) + cv2.rectangle(out, (x1, y1), (x2, y2), (255, 255, 255), border_thick) + + count = int(round(float(density_result.get("count", 0.0)))) + out = self._draw_count_text(out, count) + + return out + + def draw_full_result( + self, + img_bgr: np.ndarray, + results: List[Dict], + density_result: Dict | None = None, + ) -> np.ndarray: + annotated = draw_results(img_bgr, results) + annotated = self._draw_density_panel(annotated, density_result) + return annotated + + def predict( + self, + img_bgr: np.ndarray, + return_image: bool = False, + det_conf: float | None = None, + iou_thres: float | None = None, + ) -> Dict: + h, w = img_bgr.shape[:2] + + det_conf = settings.det_conf_thres if det_conf is None else det_conf + iou_thres = settings.det_iou_thres if iou_thres is None else iou_thres + + detections = self.detector.infer( + img_bgr, + conf_thres=det_conf, + iou_thres=iou_thres, + max_det=settings.max_detections, + ) + + detections = self._filter_target_detections(detections, img_w=w, img_h=h) + + results: List[Dict] = [] + + for det in detections: + crop = crop_with_padding( + img_bgr, + det["bbox_xyxy"], + padding_ratio=settings.crop_padding_ratio, + ) + + cls_result = self.classifier.infer(crop) + + cls_id = int(cls_result["cls_class_id"]) + cls_name = str(cls_result["cls_class_name"] or "") + + merged = { + "bbox_xyxy": [round(float(v), 2) for v in det["bbox_xyxy"]], + "det_confidence": round(float(det["det_confidence"]), 6), + "det_class_id": int(det["det_class_id"]), + "det_class_name": str(det["det_class_name"]), + "cls_class_id": cls_id, + "cls_class_name": cls_name, + "cls_confidence": round(float(cls_result["cls_confidence"]), 6), + } + + results.append(merged) + + density_result = None + if self.density_estimator is not None: + density_result = self.density_estimator.infer(img_bgr) + + density_json = None + if density_result is not None: + density_json = { + "count": int(round(float(density_result.get("count", 0.0)))), + } + + annotated_image_base64 = None + if return_image: + annotated = self.draw_full_result(img_bgr, results, density_result) + annotated_image_base64 = encode_image_to_base64(annotated) + + return { + "success": True, + "image_width": w, + "image_height": h, + "num_objects": len(results), + "target_det_class_ids": self.target_det_class_ids, + "density": density_json, + "density_result": density_result, + "results": results, + "annotated_image_base64": annotated_image_base64, + "message": "ok", + } \ No newline at end of file diff --git a/app/pipeline_raw.py b/app/pipeline_raw.py new file mode 100644 index 0000000..4ccca97 --- /dev/null +++ b/app/pipeline_raw.py @@ -0,0 +1,115 @@ +from typing import Dict, List +import os +import numpy as np + +from app.config import settings +from app.utils import read_classes, crop_with_padding, draw_results, encode_image_to_base64 +from app.inference.yolo_onnx import YOLOv5ONNXDetector +from app.inference.convnext_onnx import ConvNeXtONNXClassifier + + +class BirdCascadePipeline: + def __init__(self): + self.det_classes = read_classes(settings.det_classes_path, keep_empty=False) + + # 分类类别文件必须保留空项，因为 ;; 表示该类别不显示 + self.cls_classes = read_classes(settings.cls_classes_path, keep_empty=True) + + if not os.path.exists(settings.det_model_path): + raise FileNotFoundError(f"未找到 YOLOv5 模型: {settings.det_model_path}") + + if not os.path.exists(settings.cls_model_path): + raise FileNotFoundError(f"未找到 ConvNeXt 模型: {settings.cls_model_path}") + + self.detector = YOLOv5ONNXDetector( + model_path=settings.det_model_path, + class_names=self.det_classes, + img_size=settings.det_img_size, + device=settings.device, + ) + + self.classifier = ConvNeXtONNXClassifier( + model_path=settings.cls_model_path, + class_names=self.cls_classes, + img_size=settings.cls_img_size, + device=settings.device, + ) + + self.target_det_class_ids = settings.target_det_class_ids + + def _filter_target_detections(self, detections: List[Dict]) -> List[Dict]: + """ + 只保留 YOLO 检测类别为 0 的目标，也就是“鸟”。 + """ + if self.target_det_class_ids is None or len(self.target_det_class_ids) == 0: + return detections + + filtered = [] + for det in detections: + det_class_id = int(det.get("det_class_id", -1)) + if det_class_id in self.target_det_class_ids: + filtered.append(det) + + return filtered + + def predict( + self, + img_bgr: np.ndarray, + return_image: bool = False, + det_conf: float | None = None, + iou_thres: float | None = None, + ) -> Dict: + h, w = img_bgr.shape[:2] + + det_conf = settings.det_conf_thres if det_conf is None else det_conf + iou_thres = settings.det_iou_thres if iou_thres is None else iou_thres + + detections = self.detector.infer( + img_bgr, + conf_thres=det_conf, + iou_thres=iou_thres, + max_det=settings.max_detections, + ) + + detections = self._filter_target_detections(detections) + + results: List[Dict] = [] + + for det in detections: + crop = crop_with_padding( + img_bgr, + det["bbox_xyxy"], + padding_ratio=settings.crop_padding_ratio, + ) + + cls_result = self.classifier.infer(crop) + + cls_id = int(cls_result["cls_class_id"]) + cls_name = str(cls_result["cls_class_name"] or "") + + merged = { + "bbox_xyxy": [round(float(v), 2) for v in det["bbox_xyxy"]], + "det_confidence": round(float(det["det_confidence"]), 6), + "det_class_id": int(det["det_class_id"]), + "det_class_name": str(det["det_class_name"]), + "cls_class_id": cls_id, + "cls_class_name": cls_name, + "cls_confidence": round(float(cls_result["cls_confidence"]), 6), + } + results.append(merged) + + annotated_image_base64 = None + if return_image: + annotated = draw_results(img_bgr, results) + annotated_image_base64 = encode_image_to_base64(annotated) + + return { + "success": True, + "image_width": w, + "image_height": h, + "num_objects": len(results), + "target_det_class_ids": self.target_det_class_ids, + "results": results, + "annotated_image_base64": annotated_image_base64, + "message": "ok", + } \ No newline at end of file diff --git a/app/schemas.py b/app/schemas.py new file mode 100644 index 0000000..7fd75ec --- /dev/null +++ b/app/schemas.py @@ -0,0 +1,23 @@ +from typing import List, Optional, Dict, Any +from pydantic import BaseModel + + +class ObjectResult(BaseModel): + bbox_xyxy: List[float] + det_confidence: float + det_class_id: int + det_class_name: str + cls_class_id: int + cls_class_name: str + cls_confidence: float + + +class PredictResponse(BaseModel): + success: bool + image_width: int + image_height: int + num_objects: int + results: List[ObjectResult] + annotated_image_base64: Optional[str] = None + saved_outputs: Optional[Dict[str, Any]] = None + message: str = "ok" \ No newline at end of file diff --git a/app/utils.py b/app/utils.py new file mode 100644 index 0000000..c088d43 --- /dev/null +++ b/app/utils.py @@ -0,0 +1,451 @@ +from pathlib import Path +from typing import List, Dict +import base64 +import hashlib +import json +import os +import time +import xml.etree.ElementTree as ET + +import cv2 +import numpy as np +from PIL import Image, ImageDraw, ImageFont + +from app.config import settings + + +def read_classes(path: str, keep_empty: bool = False) -> List[str]: + """ + 支持两种类别文件格式： + 1. 每行一个类别； + 2. 用英文分号 ; 分隔类别。 + + keep_empty=True 时，会保留空类别。 + 你的分类文件里有 ;;，表示该类别不显示，所以分类类别必须保留空项。 + """ + p = Path(path) + if not p.exists(): + return [] + + text = p.read_text(encoding="utf-8").strip() + + if ";" in text: + parts = text.split(";") + if keep_empty: + return [x.strip() for x in parts] + return [x.strip() for x in parts if x.strip()] + + lines = text.splitlines() + if keep_empty: + return [x.strip() for x in lines] + return [x.strip() for x in lines if x.strip()] + + +def decode_image_from_bytes(data: bytes) -> np.ndarray: + arr = np.frombuffer(data, dtype=np.uint8) + img = cv2.imdecode(arr, cv2.IMREAD_COLOR) + if img is None: + raise ValueError("图片解码失败，请确认上传的是有效图片文件") + return img + + +def encode_image_to_base64(img_bgr: np.ndarray) -> str: + ok, buf = cv2.imencode(".jpg", img_bgr, [int(cv2.IMWRITE_JPEG_QUALITY), 90]) + if not ok: + raise ValueError("图片编码失败") + return base64.b64encode(buf.tobytes()).decode("utf-8") + + +def clip_box(box, width: int, height: int): + x1, y1, x2, y2 = box + x1 = max(0, min(float(x1), width - 1)) + y1 = max(0, min(float(y1), height - 1)) + x2 = max(0, min(float(x2), width - 1)) + y2 = max(0, min(float(y2), height - 1)) + return [x1, y1, x2, y2] + + +def crop_with_padding(img_bgr: np.ndarray, box, padding_ratio: float = 0.08) -> np.ndarray: + h, w = img_bgr.shape[:2] + x1, y1, x2, y2 = clip_box(box, w, h) + + bw = x2 - x1 + bh = y2 - y1 + pad_x = bw * padding_ratio + pad_y = bh * padding_ratio + + nx1 = int(max(0, x1 - pad_x)) + ny1 = int(max(0, y1 - pad_y)) + nx2 = int(min(w, x2 + pad_x)) + ny2 = int(min(h, y2 + pad_y)) + + if nx2 <= nx1 or ny2 <= ny1: + return img_bgr.copy() + + return img_bgr[ny1:ny2, nx1:nx2].copy() + + +def stable_color_for_class(class_id: int, class_name: str = "") -> tuple[int, int, int]: + """ + 根据类别 ID 和类别名生成稳定颜色。 + 同一物种每次运行颜色一致。 + OpenCV 使用 BGR 顺序。 + """ + key = f"{class_id}_{class_name}".encode("utf-8") + digest = hashlib.md5(key).hexdigest() + + b = int(digest[0:2], 16) + g = int(digest[2:4], 16) + r = int(digest[4:6], 16) + + # 避免颜色太暗 + b = 80 + b % 176 + g = 80 + g % 176 + r = 80 + r % 176 + + return b, g, r + + +def bgr_to_rgb(color_bgr: tuple[int, int, int]) -> tuple[int, int, int]: + b, g, r = color_bgr + return r, g, b + + +def get_display_label(item: Dict) -> str: + """ + 可视化时只显示物种中文名。 + 如果分类文件中该类别是空字符串，则不显示文字。 + """ + name = str(item.get("cls_class_name", "") or "").strip() + return name + + +def find_chinese_font() -> str | None: + """ + 自动寻找中文字体。 + 可优先通过环境变量 FONT_PATH 指定。 + """ + env_font = os.getenv("FONT_PATH", "").strip() + if env_font and Path(env_font).exists(): + return env_font + + candidates = [ + "/data2/zhangw_data/Bird_detect_cascade_service/fonts/NotoSansCJK-Regular.ttc", + "/data2/zhangw_data/Bird_detect_cascade_service/fonts/NotoSansCJK-Regular.otf", + "/data2/zhangw_data/Bird_detect_cascade_service/fonts/SourceHanSansSC-Regular.otf", + "/usr/share/fonts/opentype/noto/NotoSansCJK-Regular.ttc", + "/usr/share/fonts/opentype/noto/NotoSansCJK-Regular.otf", + "/usr/share/fonts/truetype/noto/NotoSansCJK-Regular.ttc", + "/usr/share/fonts/truetype/noto/NotoSansCJK-Regular.otf", + "/usr/share/fonts/truetype/wqy/wqy-microhei.ttc", + "/usr/share/fonts/truetype/wqy/wqy-zenhei.ttc", + "/usr/share/fonts/truetype/arphic/uming.ttc", + "/usr/share/fonts/truetype/arphic/ukai.ttc", + "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", + ] + + for p in candidates: + if Path(p).exists(): + return p + + return None + + +_FONT_CACHE: dict[int, ImageFont.FreeTypeFont | ImageFont.ImageFont] = {} + + +def get_font(font_size: int = 26): + """ + 获取支持中文的字体。 + 如果找不到中文字体，会退回默认字体，但默认字体仍可能无法显示中文。 + """ + if font_size in _FONT_CACHE: + return _FONT_CACHE[font_size] + + font_path = find_chinese_font() + + if font_path: + font = ImageFont.truetype(font_path, font_size) + else: + font = ImageFont.load_default() + + _FONT_CACHE[font_size] = font + return font + + +def draw_chinese_text_with_background( + img_bgr: np.ndarray, + text: str, + left: int, + top: int, + color_bgr: tuple[int, int, int], + font_size: int = 26, +) -> np.ndarray: + """ + 使用 PIL 绘制中文标签。 + 输入输出都是 OpenCV BGR 图像。 + + 这个版本解决中文文字在色块中偏下的问题： + 1. 先计算文字真实边界； + 2. 再根据背景框中心点反推文字绘制位置； + 3. 保证中文视觉上垂直居中。 + """ + if not text: + return img_bgr + + img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB) + pil_img = Image.fromarray(img_rgb) + draw = ImageDraw.Draw(pil_img) + + font = get_font(font_size) + color_rgb = bgr_to_rgb(color_bgr) + + # 计算文字真实边界。注意：中文字体 bbox 的 y 方向通常不是从 0 开始。 + bbox = draw.textbbox((0, 0), text, font=font) + text_left = bbox[0] + text_top = bbox[1] + text_right = bbox[2] + text_bottom = bbox[3] + + text_w = text_right - text_left + text_h = text_bottom - text_top + + pad_x = 6 + pad_y = 4 + + img_w, img_h = pil_img.size + + x1 = max(0, min(int(left), img_w - 1)) + y1 = max(0, min(int(top), img_h - 1)) + + bg_w = text_w + 2 * pad_x + bg_h = text_h + 2 * pad_y + + # 默认把标签放在检测框上方；如果上方空间不够，就放在框内部顶端。 + bg_left = x1 + if y1 - bg_h >= 0: + bg_top = y1 - bg_h + else: + bg_top = y1 + + bg_right = min(img_w - 1, bg_left + bg_w) + bg_bottom = min(img_h - 1, bg_top + bg_h) + + # 画背景色块 + draw.rectangle( + [bg_left, bg_top, bg_right, bg_bottom], + fill=color_rgb, + ) + + # 关键修正： + # 让文字真实 bbox 的中心对齐背景框中心，而不是简单 bg_top + pad_y。 + bg_center_x = (bg_left + bg_right) / 2 + bg_center_y = (bg_top + bg_bottom) / 2 + + text_draw_x = bg_center_x - text_w / 2 - text_left + text_draw_y = bg_center_y - text_h / 2 - text_top + + draw.text( + (text_draw_x, text_draw_y), + text, + font=font, + fill=(0, 0, 0), + ) + + out_rgb = np.array(pil_img) + out_bgr = cv2.cvtColor(out_rgb, cv2.COLOR_RGB2BGR) + return out_bgr + +def draw_results(img_bgr: np.ndarray, results: List[dict]) -> np.ndarray: + """ + 只画类别名，不画置信度。 + 中文文字使用 PIL 绘制，避免 OpenCV 中文乱码。 + 空类别名只画框，不显示文字。 + """ + out = img_bgr.copy() + + for item in results: + x1, y1, x2, y2 = [int(round(v)) for v in item["bbox_xyxy"]] + + cls_id = int(item.get("cls_class_id", -1)) + cls_name = str(item.get("cls_class_name", "") or "") + label = get_display_label(item) + + if settings.uniform_species_color: + color = (0, 190, 0) + else: + color = stable_color_for_class(cls_id, cls_name) + + # 先用 OpenCV 画框 + cv2.rectangle(out, (x1, y1), (x2, y2), color, 2) + + # 再用 PIL 画中文 + if label: + # 根据图像大小自适应字体大小，避免无人机大图上文字太小 + h, w = out.shape[:2] + font_size = max(22, int(min(w, h) * 0.018)) + font_size = min(font_size, 42) + + out = draw_chinese_text_with_background( + img_bgr=out, + text=label, + left=x1, + top=y1, + color_bgr=color, + font_size=font_size, + ) + + return out + + +def now_timestamp_minute() -> str: + return time.strftime("%Y%m%d%H%M") + + +def safe_stem(filename: str) -> str: + name = Path(filename).stem + name = name.replace(" ", "_") + return name + + +def make_unique_base(output_dir: Path, timestamp: str, original_filename: str) -> str: + """ + 生成 时间戳_原文件名。 + 如果同一分钟内重复处理同一个文件，则自动加 _01、_02。 + """ + stem = safe_stem(original_filename) + base = f"{timestamp}_{stem}" + + candidate = base + idx = 1 + while ( + (output_dir / f"{candidate}.json").exists() + or (output_dir / f"{candidate}.xml").exists() + or (output_dir / f"{candidate}.jpg").exists() + or (output_dir / f"{candidate}.mp4").exists() + ): + candidate = f"{base}_{idx:02d}" + idx += 1 + + return candidate + + +def save_json(path: Path, data: Dict): + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(data, ensure_ascii=False, indent=2), encoding="utf-8") + + +def indent_xml(elem, level: int = 0): + i = "\n" + level * " " + if len(elem): + if not elem.text or not elem.text.strip(): + elem.text = i + " " + for child in elem: + indent_xml(child, level + 1) + if not child.tail or not child.tail.strip(): + child.tail = i + if level and (not elem.tail or not elem.tail.strip()): + elem.tail = i + + +def write_image_xml( + xml_path: Path, + image_filename: str, + image_path: str, + width: int, + height: int, + results: List[Dict], +): + """ + 图片 XML，接近 Pascal VOC 格式。 + 空类别名为了避免 XML name 为空，会写成 class_。 + """ + root = ET.Element("annotation") + + ET.SubElement(root, "filename").text = image_filename + ET.SubElement(root, "path").text = image_path + + size = ET.SubElement(root, "size") + ET.SubElement(size, "width").text = str(width) + ET.SubElement(size, "height").text = str(height) + ET.SubElement(size, "depth").text = "3" + + for item in results: + obj = ET.SubElement(root, "object") + + cls_id = int(item.get("cls_class_id", -1)) + cls_name = str(item.get("cls_class_name", "") or "").strip() + xml_name = cls_name if cls_name else f"class_{cls_id}" + + ET.SubElement(obj, "name").text = xml_name + ET.SubElement(obj, "class_id").text = str(cls_id) + + box = item["bbox_xyxy"] + bndbox = ET.SubElement(obj, "bndbox") + ET.SubElement(bndbox, "xmin").text = str(int(round(box[0]))) + ET.SubElement(bndbox, "ymin").text = str(int(round(box[1]))) + ET.SubElement(bndbox, "xmax").text = str(int(round(box[2]))) + ET.SubElement(bndbox, "ymax").text = str(int(round(box[3]))) + + indent_xml(root) + tree = ET.ElementTree(root) + xml_path.parent.mkdir(parents=True, exist_ok=True) + tree.write(xml_path, encoding="utf-8", xml_declaration=True) + + +def write_video_xml( + xml_path: Path, + video_filename: str, + video_path: str, + width: int, + height: int, + fps: float, + total_frames: int, + sample_interval: int, + frame_results: List[Dict], +): + """ + 视频 XML：记录每个处理帧上的目标框。 + """ + root = ET.Element("video_annotation") + + ET.SubElement(root, "filename").text = video_filename + ET.SubElement(root, "path").text = video_path + + info = ET.SubElement(root, "video_info") + ET.SubElement(info, "width").text = str(width) + ET.SubElement(info, "height").text = str(height) + ET.SubElement(info, "fps").text = str(fps) + ET.SubElement(info, "total_frames").text = str(total_frames) + ET.SubElement(info, "sample_interval").text = str(sample_interval) + + frames_node = ET.SubElement(root, "frames") + + for frame_item in frame_results: + frame_node = ET.SubElement(frames_node, "frame") + ET.SubElement(frame_node, "frame_index").text = str(frame_item["frame_index"]) + ET.SubElement(frame_node, "time_sec").text = str(frame_item["time_sec"]) + + objects_node = ET.SubElement(frame_node, "objects") + + for item in frame_item.get("results", []): + obj = ET.SubElement(objects_node, "object") + + cls_id = int(item.get("cls_class_id", -1)) + cls_name = str(item.get("cls_class_name", "") or "").strip() + xml_name = cls_name if cls_name else f"class_{cls_id}" + + ET.SubElement(obj, "name").text = xml_name + ET.SubElement(obj, "class_id").text = str(cls_id) + + box = item["bbox_xyxy"] + bndbox = ET.SubElement(obj, "bndbox") + ET.SubElement(bndbox, "xmin").text = str(int(round(box[0]))) + ET.SubElement(bndbox, "ymin").text = str(int(round(box[1]))) + ET.SubElement(bndbox, "xmax").text = str(int(round(box[2]))) + ET.SubElement(bndbox, "ymax").text = str(int(round(box[3]))) + + indent_xml(root) + tree = ET.ElementTree(root) + xml_path.parent.mkdir(parents=True, exist_ok=True) + tree.write(xml_path, encoding="utf-8", xml_declaration=True) diff --git a/app/video_pipeline.py b/app/video_pipeline.py new file mode 100644 index 0000000..aed56b8 --- /dev/null +++ b/app/video_pipeline.py @@ -0,0 +1,373 @@ +from __future__ import annotations + +import shutil +import time +from pathlib import Path +from typing import Dict, List, Optional + +import cv2 + +from app.config import settings +from app.pipeline import BirdCascadePipeline +from app.utils import ( + draw_results, + now_timestamp_minute, + safe_stem, + save_json, + write_video_xml, +) + + +def make_unique_result_dir(parent_dir: Path, timestamp: str, original_filename: str) -> tuple[Path, str, str]: + """ + 创建单次视频识别结果目录。 + + 目录名： + 时间戳_原文件名stem + + 例如： + runs/video/202605181431_DJI_20260423155914_0006_Z/ + + 返回： + result_dir, result_base_name, original_stem + """ + original_stem = safe_stem(original_filename) + base_name = f"{timestamp}_{original_stem}" + + result_dir = parent_dir / base_name + + if not result_dir.exists(): + result_dir.mkdir(parents=True, exist_ok=True) + return result_dir, base_name, original_stem + + idx = 1 + while True: + candidate_base = f"{base_name}_{idx:02d}" + candidate_dir = parent_dir / candidate_base + if not candidate_dir.exists(): + candidate_dir.mkdir(parents=True, exist_ok=True) + return candidate_dir, candidate_base, original_stem + idx += 1 + + +def iou_xyxy(a: List[float], b: List[float]) -> float: + ax1, ay1, ax2, ay2 = a + bx1, by1, bx2, by2 = b + + x1 = max(ax1, bx1) + y1 = max(ay1, by1) + x2 = min(ax2, bx2) + y2 = min(ay2, by2) + + inter = max(0.0, x2 - x1) * max(0.0, y2 - y1) + area_a = max(0.0, ax2 - ax1) * max(0.0, ay2 - ay1) + area_b = max(0.0, bx2 - bx1) * max(0.0, by2 - by1) + + return inter / (area_a + area_b - inter + 1e-7) + + +class VideoResultSmoother: + """ + 降低视频框闪烁： + 1. 非推理帧复用上一推理结果； + 2. 短暂漏检时保留上一结果若干帧； + 3. 对匹配目标的坐标做指数平滑； + 4. 对同一目标累计物种置信度，避免标签和颜色逐帧跳变。 + """ + + def __init__(self, hold_frames: int = 8, alpha: float = 0.65): + self.hold_frames = hold_frames + self.alpha = alpha + self.last_results: List[Dict] = [] + self.last_update_frame: int = -10**9 + + @staticmethod + def _public_results(results: List[Dict]) -> List[Dict]: + return [ + { + key: value + for key, value in item.items() + if not key.startswith("_smoother_") + } + for item in results + ] + + def update(self, current_results: List[Dict], frame_index: int) -> List[Dict]: + if len(current_results) == 0: + if frame_index - self.last_update_frame <= self.hold_frames: + return self._public_results(self.last_results) + self.last_results = [] + return [] + + if len(self.last_results) == 0: + self.last_results = [] + for current in current_results: + item = dict(current) + class_key = ( + int(item.get("cls_class_id", -1)), + str(item.get("cls_class_name", "") or ""), + ) + confidence = max(float(item.get("cls_confidence", 0.0)), 0.01) + item["_smoother_class_votes"] = {class_key: confidence} + self.last_results.append(item) + self.last_update_frame = frame_index + return self._public_results(self.last_results) + + smoothed_results = [] + used_last = set() + + for cur in current_results: + cur_box = cur["bbox_xyxy"] + + best_idx = -1 + best_iou = 0.0 + + for idx, prev in enumerate(self.last_results): + if idx in used_last: + continue + + score = iou_xyxy(cur_box, prev["bbox_xyxy"]) + if score > best_iou: + best_iou = score + best_idx = idx + + new_item = dict(cur) + + if best_idx >= 0 and best_iou >= 0.1: + prev = self.last_results[best_idx] + prev_box = prev["bbox_xyxy"] + + smooth_box = [] + for p, c in zip(prev_box, cur_box): + smooth_box.append(round(self.alpha * float(p) + (1 - self.alpha) * float(c), 2)) + + new_item["bbox_xyxy"] = smooth_box + + votes = { + key: float(value) * 0.90 + for key, value in prev.get( + "_smoother_class_votes", + {}, + ).items() + } + current_key = ( + int(cur.get("cls_class_id", -1)), + str(cur.get("cls_class_name", "") or ""), + ) + current_confidence = max( + float(cur.get("cls_confidence", 0.0)), + 0.01, + ) + votes[current_key] = votes.get(current_key, 0.0) + current_confidence + chosen_key = max(votes, key=votes.get) + new_item["_smoother_class_votes"] = votes + new_item["cls_class_id"] = chosen_key[0] + new_item["cls_class_name"] = chosen_key[1] + if chosen_key != current_key: + new_item["show_cls_name"] = bool( + prev.get("show_cls_name", False) + ) + used_last.add(best_idx) + else: + class_key = ( + int(cur.get("cls_class_id", -1)), + str(cur.get("cls_class_name", "") or ""), + ) + confidence = max(float(cur.get("cls_confidence", 0.0)), 0.01) + new_item["_smoother_class_votes"] = {class_key: confidence} + + smoothed_results.append(new_item) + + self.last_results = smoothed_results + self.last_update_frame = frame_index + + return self._public_results(smoothed_results) + + def get_for_skipped_frame(self, frame_index: int) -> List[Dict]: + if frame_index - self.last_update_frame <= self.hold_frames: + return self._public_results(self.last_results) + return [] + + +def process_video_file( + video_path: str, + pipeline: BirdCascadePipeline, + original_filename: str, + sample_interval: Optional[int] = None, + save_annotated_video: Optional[bool] = None, +) -> Dict: + """ + 视频结果自动保存到： + runs/video/时间戳_原视频名/ + + 目录内部生成： + 原视频名.mp4 + 时间戳_原视频名.json + 时间戳_原视频名.xml + 原视频名_source.原后缀 + frames/ + + 示例： + runs/video/202605181431_DJI_20260423155914_0006_Z/ + ├── DJI_20260423155914_0006_Z.mp4 + ├── 202605181431_DJI_20260423155914_0006_Z.json + ├── 202605181431_DJI_20260423155914_0006_Z.xml + ├── DJI_20260423155914_0006_Z_source.MP4 + └── frames/ + """ + sample_interval = settings.video_sample_interval if sample_interval is None else int(sample_interval) + save_annotated_video = settings.save_annotated_video if save_annotated_video is None else bool(save_annotated_video) + + video_root = Path(settings.output_dir) / settings.video_output_subdir + video_root.mkdir(parents=True, exist_ok=True) + + timestamp = now_timestamp_minute() + result_dir, result_base_name, original_stem = make_unique_result_dir( + parent_dir=video_root, + timestamp=timestamp, + original_filename=original_filename, + ) + + original_suffix = Path(original_filename).suffix or ".mp4" + + source_video_path = result_dir / f"{original_stem}_source{original_suffix}" + + # 输出视频固定为 mp4，文件名不带时间戳 + annotated_video_path = result_dir / f"{original_stem}.mp4" + + # JSON 和 XML 带时间戳 + json_path = result_dir / f"{result_base_name}.json" + xml_path = result_dir / f"{result_base_name}.xml" + + frames_dir = result_dir / "frames" + frames_dir.mkdir(parents=True, exist_ok=True) + + shutil.copy2(video_path, source_video_path) + + cap = cv2.VideoCapture(str(source_video_path)) + if not cap.isOpened(): + raise ValueError(f"视频打开失败: {source_video_path}") + + fps = cap.get(cv2.CAP_PROP_FPS) + if fps is None or fps <= 0: + fps = 25.0 + + width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH) or 0) + height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT) or 0) + total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT) or 0) + + writer = None + if save_annotated_video: + fourcc = cv2.VideoWriter_fourcc(*"mp4v") + writer = cv2.VideoWriter(str(annotated_video_path), fourcc, fps, (width, height)) + + smoother = VideoResultSmoother( + hold_frames=settings.smooth_hold_frames, + alpha=settings.box_smooth_alpha, + ) + + frame_results = [] + processed_frames = 0 + detected_frames = 0 + frame_index = 0 + + start_time = time.time() + + try: + while True: + ok, frame = cap.read() + if not ok: + break + + should_process = frame_index % sample_interval == 0 + + if should_process: + output = pipeline.predict(frame, return_image=False) + raw_results = output["results"] + processed_frames += 1 + + if settings.smooth_video_boxes: + results = smoother.update(raw_results, frame_index) + else: + results = raw_results + + if len(results) > 0: + detected_frames += 1 + + frame_results.append( + { + "frame_index": frame_index, + "time_sec": round(frame_index / fps, 3), + "num_objects": len(results), + "results": results, + } + ) + else: + if settings.smooth_video_boxes: + results = smoother.get_for_skipped_frame(frame_index) + else: + results = [] + + annotated_frame = draw_results(frame, results) + + if should_process and len(results) > 0: + frame_img_path = frames_dir / f"frame_{frame_index:08d}.jpg" + cv2.imwrite(str(frame_img_path), annotated_frame) + + if writer is not None: + writer.write(annotated_frame) + + frame_index += 1 + + finally: + cap.release() + if writer is not None: + writer.release() + + elapsed = round(time.time() - start_time, 3) + + summary = { + "success": True, + "type": "video", + "original_filename": original_filename, + "result_dir": str(result_dir), + "source_video": str(source_video_path), + "annotated_video": str(annotated_video_path) if annotated_video_path.exists() else None, + "json": str(json_path), + "xml": str(xml_path), + "frames_dir": str(frames_dir), + "video_info": { + "width": width, + "height": height, + "fps": fps, + "total_frames": total_frames, + }, + "sample_interval": sample_interval, + "smooth_video_boxes": settings.smooth_video_boxes, + "smooth_hold_frames": settings.smooth_hold_frames, + "box_smooth_alpha": settings.box_smooth_alpha, + "processed_frames": processed_frames, + "detected_frames": detected_frames, + "elapsed_sec": elapsed, + } + + json_data = { + "summary": summary, + "frames": frame_results, + } + + save_json(json_path, json_data) + + write_video_xml( + xml_path=xml_path, + video_filename=annotated_video_path.name, + video_path=str(annotated_video_path), + width=width, + height=height, + fps=fps, + total_frames=total_frames, + sample_interval=sample_interval, + frame_results=frame_results, + ) + + return summary diff --git a/assets/fonts/NotoSansCJK-Regular.ttc b/assets/fonts/NotoSansCJK-Regular.ttc new file mode 100644 index 0000000..38b534a Binary files /dev/null and b/assets/fonts/NotoSansCJK-Regular.ttc differ diff --git a/docker-compose.yml b/docker-compose.yml new file mode 100644 index 0000000..25745e4 --- /dev/null +++ b/docker-compose.yml @@ -0,0 +1,34 @@ +services: + bird-cascade-batch: + build: . + image: bird-cascade-batch:latest + volumes: + - ./models:/app/models:ro + - ./input:/app/input:ro + - ./runs:/app/runs + environment: + ENABLE_DETECTION: "true" + ENABLE_CLASSIFICATION: "true" + ENABLE_DENSITY: "true" + DET_MODEL_PATH: /app/models/YNDC_20251010_yolov5_optimize.onnx + CLS_MODEL_PATH: /app/models/HHK_20260525_mmcls_optimize.onnx + DET_CLASSES_PATH: /app/models/YNDC_20251010_yolov5_optimize.txt + CLS_CLASSES_PATH: /app/models/HHK_20260525_mmcls_optimize.txt + DENSITY_MODEL_PATH: /app/models/bird_density.onnx + DET_IMG_SIZE: "640" + CLS_IMG_SIZE: "224" + DENSITY_IMG_SIZE: "0" + OUTPUT_DIR: /app/runs + TARGET_DET_CLASS_IDS: "0" + VIDEO_SAMPLE_INTERVAL: "2" + DEVICE: cpu + UNIFORM_SPECIES_COLOR: "false" + SMOOTH_VIDEO_BOXES: "true" + SMOOTH_HOLD_FRAMES: "3" + BOX_SMOOTH_ALPHA: "0.25" + FONT_PATH: /app/assets/fonts/NotoSansCJK-Regular.ttc + command: + - --input + - /app/input + - --output + - /app/runs diff --git a/font_tmp/etc/fonts/conf.d/70-fonts-noto-cjk.conf b/font_tmp/etc/fonts/conf.d/70-fonts-noto-cjk.conf new file mode 120000 index 0000000..af9a22a --- /dev/null +++ b/font_tmp/etc/fonts/conf.d/70-fonts-noto-cjk.conf @@ -0,0 +1 @@ +/usr/share/fontconfig/conf.avail/70-fonts-noto-cjk.conf \ No newline at end of file diff --git a/font_tmp/usr/share/doc/fonts-noto-cjk/NEWS.Debian.gz b/font_tmp/usr/share/doc/fonts-noto-cjk/NEWS.Debian.gz new file mode 100644 index 0000000..e495ad7 Binary files /dev/null and b/font_tmp/usr/share/doc/fonts-noto-cjk/NEWS.Debian.gz differ diff --git a/font_tmp/usr/share/doc/fonts-noto-cjk/NEWS.gz b/font_tmp/usr/share/doc/fonts-noto-cjk/NEWS.gz new file mode 100644 index 0000000..c8e8c04 Binary files /dev/null and b/font_tmp/usr/share/doc/fonts-noto-cjk/NEWS.gz differ diff --git a/font_tmp/usr/share/doc/fonts-noto-cjk/changelog.Debian.gz b/font_tmp/usr/share/doc/fonts-noto-cjk/changelog.Debian.gz new file mode 100644 index 0000000..fd392c6 Binary files /dev/null and b/font_tmp/usr/share/doc/fonts-noto-cjk/changelog.Debian.gz differ diff --git a/font_tmp/usr/share/doc/fonts-noto-cjk/copyright b/font_tmp/usr/share/doc/fonts-noto-cjk/copyright new file mode 100644 index 0000000..a86cbaa --- /dev/null +++ b/font_tmp/usr/share/doc/fonts-noto-cjk/copyright @@ -0,0 +1,186 @@ +Format: https://www.debian.org/doc/packaging-manuals/copyright-format/1.0/ +Upstream-Name: noto-cjk +Upstream-Contact: https://github.com/googlei18n/noto-cjk/issues +Source: https://github.com/googlei18n/noto-cjk +Files-Excluded: + NotoSansCJKjp-Black.otf + NotoSansCJKjp-Bold.otf + NotoSansCJKjp-DemiLight.otf + NotoSansCJKjp-Light.otf + NotoSansCJKjp-Medium.otf + NotoSansCJKjp-Regular.otf + NotoSansCJKjp-Thin.otf + NotoSansCJKkr-Black.otf + NotoSansCJKkr-Bold.otf + NotoSansCJKkr-DemiLight.otf + NotoSansCJKkr-Light.otf + NotoSansCJKkr-Medium.otf + NotoSansCJKkr-Regular.otf + NotoSansCJKkr-Thin.otf + NotoSansCJKsc-Black.otf + NotoSansCJKsc-Bold.otf + NotoSansCJKsc-DemiLight.otf + NotoSansCJKsc-Light.otf + NotoSansCJKsc-Medium.otf + NotoSansCJKsc-Regular.otf + NotoSansCJKsc-Thin.otf + NotoSansCJKtc-Black.otf + NotoSansCJKtc-Bold.otf + NotoSansCJKtc-DemiLight.otf + NotoSansCJKtc-Light.otf + NotoSansCJKtc-Medium.otf + NotoSansCJKtc-Regular.otf + NotoSansCJKtc-Thin.otf + NotoSansCJK.ttc.zip + NotoSansJP-Black.otf + NotoSansJP-Bold.otf + NotoSansJP-DemiLight.otf + NotoSansJP-Light.otf + NotoSansJP-Medium.otf + NotoSansJP-Regular.otf + NotoSansJP-Thin.otf + NotoSansKR-Black.otf + NotoSansKR-Bold.otf + NotoSansKR-DemiLight.otf + NotoSansKR-Light.otf + NotoSansKR-Medium.otf + NotoSansKR-Regular.otf + NotoSansKR-Thin.otf + NotoSansMonoCJKjp-Bold.otf + NotoSansMonoCJKjp-Regular.otf + NotoSansMonoCJKkr-Bold.otf + NotoSansMonoCJKkr-Regular.otf + NotoSansMonoCJKsc-Bold.otf + NotoSansMonoCJKsc-Regular.otf + NotoSansMonoCJKtc-Bold.otf + NotoSansMonoCJKtc-Regular.otf + NotoSansSC-Black.otf + NotoSansSC-Bold.otf + NotoSansSC-DemiLight.otf + NotoSansSC-Light.otf + NotoSansSC-Medium.otf + NotoSansSC-Regular.otf + NotoSansSC-Thin.otf + NotoSansTC-Black.otf + NotoSansTC-Bold.otf + NotoSansTC-DemiLight.otf + NotoSansTC-Light.otf + NotoSansTC-Medium.otf + NotoSansTC-Regular.otf + NotoSansTC-Thin.otf + +Files: * +Copyright: 2010-2015, Google Corporation +License: SIL-1.1 + +Files: debian/* +Copyright: 2013, Vasudev Kamath + 2013-2015, Jonas Smedegaard + 2015, ChangZhuo Chen (陳昌倬) +License: GPL-3+ + +License: SIL-1.1 + ----------------------------------------------------------- + SIL OPEN FONT LICENSE Version 1.1 - 26 February 2007 + ----------------------------------------------------------- + . + PREAMBLE + . + The goals of the Open Font License (OFL) are to stimulate worldwide + development of collaborative font projects, to support the font + creation efforts of academic and linguistic communities, and to provide + a free and open framework in which fonts may be shared and improved in + partnership with others. + . + The OFL allows the licensed fonts to be used, studied, modified and + redistributed freely as long as they are not sold by themselves. The + fonts, including any derivative works, can be bundled, embedded, + redistributed and/or sold with any software provided that any reserved + names are not used by derivative works. The fonts and derivatives, + however, cannot be released under any other type of license. The + requirement for fonts to remain under this license does not apply to + any document created using the fonts or their derivatives. + . + DEFINITIONS + . + "Font Software" refers to the set of files released by the Copyright + Holder(s) under this license and clearly marked as such. This may + include source files, build scripts and documentation. + . + "Reserved Font Name" refers to any names specified as such after the + copyright statement(s). + . + "Original Version" refers to the collection of Font Software components + as distributed by the Copyright Holder(s). + . + "Modified Version" refers to any derivative made by adding to, + deleting, or substituting -- in part or in whole -- any of the + components of the Original Version, by changing formats or by porting + the Font Software to a new environment. + . + "Author" refers to any designer, engineer, programmer, technical writer + or other person who contributed to the Font Software. + . + PERMISSION & CONDITIONS + . + Permission is hereby granted, free of charge, to any person obtaining a + copy of the Font Software, to use, study, copy, merge, embed, modify, + redistribute, and sell modified and unmodified copies of the Font + Software, subject to the following conditions: + . + 1) Neither the Font Software nor any of its individual components, in + Original or Modified Versions, may be sold by itself. + . + 2) Original or Modified Versions of the Font Software may be bundled, + redistributed and/or sold with any software, provided that each copy + contains the above copyright notice and this license. These can be + included either as stand-alone text files, human-readable headers or + in the appropriate machine-readable metadata fields within text or + binary files as long as those fields can be easily viewed by the user. + . + 3) No Modified Version of the Font Software may use the Reserved Font + Name(s) unless explicit written permission is granted by the + corresponding Copyright Holder. This restriction only applies to the + primary font name as presented to the users. + . + 4) The name(s) of the Copyright Holder(s) or the Author(s) of the Font + Software shall not be used to promote, endorse or advertise any + Modified Version, except to acknowledge the contribution(s) of the + Copyright Holder(s) and the Author(s) or with their explicit written + permission. + . + 5) The Font Software, modified or unmodified, in part or in whole, must + be distributed entirely under this license, and must not be distributed + under any other license. The requirement for fonts to remain under this + license does not apply to any document created using the Font Software. + . + TERMINATION + . + This license becomes null and void if any of the above conditions are + not met. + . + DISCLAIMER + . + THE FONT SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, + EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO ANY WARRANTIES OF + MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT + OF COPYRIGHT, PATENT, TRADEMARK, OR OTHER RIGHT. IN NO EVENT SHALL THE + COPYRIGHT HOLDER BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, + INCLUDING ANY GENERAL, SPECIAL, INDIRECT, INCIDENTAL, OR CONSEQUENTIAL + DAMAGES, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING + FROM, OUT OF THE USE OR INABILITY TO USE THE FONT SOFTWARE OR FROM + OTHER DEALINGS IN THE FONT SOFTWARE. + +License: GPL-3+ + This program is free software: you can redistribute it and/or modify + it under the terms of the GNU General Public License as published by + the Free Software Foundation, either version 3 of the License, or + (at your option) any later version. + . + This program is distributed in the hope that it will be useful, + but WITHOUT ANY WARRANTY; without even the implied warranty of + MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + GNU General Public License for more details. + . + On Debian systems, the complete text of the GNU General + Public License version 3 can be found in "/usr/share/common-licenses/GPL-3". diff --git a/font_tmp/usr/share/fontconfig/conf.avail/70-fonts-noto-cjk.conf b/font_tmp/usr/share/fontconfig/conf.avail/70-fonts-noto-cjk.conf new file mode 100644 index 0000000..c72aa44 --- /dev/null +++ b/font_tmp/usr/share/fontconfig/conf.avail/70-fonts-noto-cjk.conf @@ -0,0 +1,183 @@ + + + + + + ja + + + serif + + + Noto Serif CJK JP + + + + + + ko + + + serif + + + Noto Serif CJK KR + + + + + + zh-cn + + + serif + + + Noto Serif CJK SC + + + + + + zh-tw + + + serif + + + Noto Serif CJK TC + + + + + + zh-hk + + + serif + + + Noto Serif CJK HK + + + + + + ja + + + sans-serif + + + Noto Sans CJK JP + + + + + + ko + + + sans-serif + + + Noto Sans CJK KR + + + + + + zh-cn + + + sans-serif + + + Noto Sans CJK SC + + + + + + zh-tw + + + sans-serif + + + Noto Sans CJK TC + + + + + + zh-hk + + + sans-serif + + + Noto Sans CJK HK + + + + + + ja + + + monospace + + + Noto Sans Mono CJK JP + + + + + + ko + + + monospace + + + Noto Sans Mono CJK KR + + + + + + zh-cn + + + monospace + + + Noto Sans Mono CJK SC + + + + + + zh-tw + + + monospace + + + Noto Sans Mono CJK TC + + + + + + zh-hk + + + monospace + + + Noto Sans Mono CJK HK + + + diff --git a/font_tmp/usr/share/fonts/opentype/noto/NotoSansCJK-Bold.ttc b/font_tmp/usr/share/fonts/opentype/noto/NotoSansCJK-Bold.ttc new file mode 100644 index 0000000..22ea5a3 Binary files /dev/null and b/font_tmp/usr/share/fonts/opentype/noto/NotoSansCJK-Bold.ttc differ diff --git a/font_tmp/usr/share/fonts/opentype/noto/NotoSansCJK-Regular.ttc b/font_tmp/usr/share/fonts/opentype/noto/NotoSansCJK-Regular.ttc new file mode 100644 index 0000000..38b534a Binary files /dev/null and b/font_tmp/usr/share/fonts/opentype/noto/NotoSansCJK-Regular.ttc differ diff --git a/font_tmp/usr/share/fonts/opentype/noto/NotoSerifCJK-Bold.ttc b/font_tmp/usr/share/fonts/opentype/noto/NotoSerifCJK-Bold.ttc new file mode 100644 index 0000000..0c33469 Binary files /dev/null and b/font_tmp/usr/share/fonts/opentype/noto/NotoSerifCJK-Bold.ttc differ diff --git a/font_tmp/usr/share/fonts/opentype/noto/NotoSerifCJK-Regular.ttc b/font_tmp/usr/share/fonts/opentype/noto/NotoSerifCJK-Regular.ttc new file mode 100644 index 0000000..9df7017 Binary files /dev/null and b/font_tmp/usr/share/fonts/opentype/noto/NotoSerifCJK-Regular.ttc differ diff --git a/fonts-noto-cjk_1%3a20190410+repack1-2_all.deb b/fonts-noto-cjk_1%3a20190410+repack1-2_all.deb new file mode 100644 index 0000000..fcf07ee Binary files /dev/null and b/fonts-noto-cjk_1%3a20190410+repack1-2_all.deb differ diff --git a/input/.gitkeep b/input/.gitkeep new file mode 100644 index 0000000..8b13789 --- /dev/null +++ b/input/.gitkeep @@ -0,0 +1 @@ + diff --git a/models/.gitkeep b/models/.gitkeep new file mode 100644 index 0000000..e69de29 diff --git a/models/HHK_20260525_mmcls_optimize.onnx b/models/HHK_20260525_mmcls_optimize.onnx new file mode 100644 index 0000000..c28cd75 --- /dev/null +++ b/models/HHK_20260525_mmcls_optimize.onnx @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d69178bd5f31db723a44f926cf20bd3c901e189a714d7baa1b741473776e2c5c +size 350591905 diff --git a/models/HHK_20260525_mmcls_optimize.txt b/models/HHK_20260525_mmcls_optimize.txt new file mode 100644 index 0000000..4c98559 --- /dev/null +++ b/models/HHK_20260525_mmcls_optimize.txt @@ -0,0 +1 @@ +普通翠鸟;针尾鸭;绿翅鸭;绿头鸭;斑嘴鸭;灰雁;鸿雁;豆雁;苍鹭;红头潜鸭;斑嘴鸭;斑背潜鸭;白眼潜鸭;;鹊鸭;环颈鸻;金眶鸻;棕头鸥;红嘴鸥;东方白鹳;小天鹅;大天鹅;疣鼻天鹅;白鹭;普通鸬鹚;扇尾沙锥;黑水鸡;灰鹤;丹顶鹤;蛎鹬;黑翅长脚鹬;黑尾鸥;黑鹳;罗纹鸭;赤颈鸭;赤膀鸭;普通秋沙鸭;普通秋沙鸭;夜鹭;卷羽鹈鹕;;普通鸬鹚;雉鸡;罗纹鸭;白琵鹭;凤头䴙䴘;反嘴鹬;琵嘴鸭;白眉鸭;斑嘴鸭;小䴙䴘;赤麻鸭;鹤鹬;林鹬;泽鹬;红脚鹬;灰头麦鸡;斑嘴鸭 \ No newline at end of file diff --git a/models/JXPYH_20260204_mmcls_optimize.onnx b/models/JXPYH_20260204_mmcls_optimize.onnx new file mode 100644 index 0000000..242cc6d --- /dev/null +++ b/models/JXPYH_20260204_mmcls_optimize.onnx @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:633d31c72ee08a698b3c3e2726488f9766b5f72f80025b6607720d144a35263b +size 350550454 diff --git a/models/JXPYH_20260204_mmcls_optimize.txt b/models/JXPYH_20260204_mmcls_optimize.txt new file mode 100644 index 0000000..5c2a848 --- /dev/null +++ b/models/JXPYH_20260204_mmcls_optimize.txt @@ -0,0 +1 @@ +八哥;绿头鸭;绿头鸭;斑嘴鸭;白额雁;灰雁;鸿雁;豆雁;白枕鹤;苍鹭;草鹭;池鹭;鸟;;环颈鸻;金眶鸻;须浮鸥;白翅浮鸥;东方白鹳;黑鹳;小天鹅;白鹭;白骨顶;扇尾沙锥;黑水鸡;普通燕鸻;黑领椋鸟;灰鹤;白头鹤;白鹤;罗纹鸭;赤颈鸭;夜鹭;普通鸬鹚;雉鸡;白琵鹭;凤头䴙䴘;反嘴鹬;珠颈斑鸠;丝光椋鸟;小䴙䴘;鹤鹬;青脚鹬;灰头麦鸡;凤头麦鸡 \ No newline at end of file diff --git a/models/YNDC_20251010_yolov5_optimize.onnx b/models/YNDC_20251010_yolov5_optimize.onnx new file mode 100644 index 0000000..c36c66d --- /dev/null +++ b/models/YNDC_20251010_yolov5_optimize.onnx @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f9e1a82171905caae8bad2ae960eeb9e57c108e98e2ffb2652208bf2f69ca6dd +size 49393300 diff --git a/models/YNDC_20251010_yolov5_optimize.txt b/models/YNDC_20251010_yolov5_optimize.txt new file mode 100644 index 0000000..fabfbd3 --- /dev/null +++ b/models/YNDC_20251010_yolov5_optimize.txt @@ -0,0 +1,5 @@ +鸟 +人 +船 +牛 +其他 \ No newline at end of file diff --git a/models/ZJLY_20260427_mmcls_optimize.onnx b/models/ZJLY_20260427_mmcls_optimize.onnx new file mode 100644 index 0000000..5e474cc --- /dev/null +++ b/models/ZJLY_20260427_mmcls_optimize.onnx @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:16ecbf1e9a47a264962ff54ab169e50fb7eb09e25572c87958d132ff6eda10e7 +size 350563205 diff --git a/models/ZJLY_20260427_mmcls_optimize.txt b/models/ZJLY_20260427_mmcls_optimize.txt new file mode 100644 index 0000000..af9fa48 --- /dev/null +++ b/models/ZJLY_20260427_mmcls_optimize.txt @@ -0,0 +1 @@ +八哥;普通翠鸟;绿翅鸭;绿头鸭;绿头鸭;斑嘴鸭;白额雁;灰雁;鸿雁;苍鹭;池鹭;红头潜鸭;凤头潜鸭;斑背潜鸭;;;斑鱼狗;须浮鸥;小天鹅;白鹭;白骨顶;黑水鸡;黑翅长脚鹬;水雉;棕背伯劳;罗纹鸭;赤颈鸭;白鹡鸰;夜鹭;;鹗;大山雀;;普通鸬鹚;喜鹊;白琵鹭;凤头䴙䴘;黑颈䴙䴘;反嘴鹬;琵嘴鸭;白眉鸭;珠颈斑鸠;丝光椋鸟;山斑鸠;火斑鸠;小䴙䴘;赤麻鸭;青脚鹬;斑鸫;灰头麦鸡; \ No newline at end of file diff --git a/models/bird_density.onnx b/models/bird_density.onnx new file mode 100644 index 0000000..64e2954 --- /dev/null +++ b/models/bird_density.onnx @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:1192c794352dcde88ada1e8285850d4e303c1959b3a67ab043a91793b19ad269 +size 86004342 diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..708998c --- /dev/null +++ b/requirements.txt @@ -0,0 +1,9 @@ +fastapi==0.115.6 +uvicorn[standard]==0.34.0 +python-multipart==0.0.20 +numpy==1.26.4 +opencv-python-headless==4.10.0.84 +onnxruntime==1.20.1 +Pillow==11.0.0 +requests==2.32.3 +pydantic==2.10.4 diff --git a/run_batch.sh b/run_batch.sh new file mode 100755 index 0000000..f3c3b85 --- /dev/null +++ b/run_batch.sh @@ -0,0 +1,9 @@ +#!/usr/bin/env bash +set -euo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +PYTHON_BIN="${PYTHON_BIN:-python}" + +cd "$SCRIPT_DIR" + +exec "$PYTHON_BIN" -m app.batch_runner "$@" diff --git a/tests/test_batch_runner_cli.py b/tests/test_batch_runner_cli.py new file mode 100644 index 0000000..c586509 --- /dev/null +++ b/tests/test_batch_runner_cli.py @@ -0,0 +1,233 @@ +import unittest +from unittest.mock import patch + +from app import batch_runner +from app.config import settings + + +class BatchRunnerCliTests(unittest.TestCase): + def test_readme_options_are_parsed(self): + args = batch_runner.parse_args( + [ + "--input", + "/tmp/input", + "--output", + "/tmp/output", + "--sample-interval", + "5", + "--no-recursive", + "--disable-detection", + "--disable-classification", + "--disable-density", + "--det-conf-thres", + "0.15", + "--det-iou-thres", + "0.40", + "--det-second-conf-thres", + "0.45", + "--min-box-area-ratio", + "0.0001", + "--max-box-area-ratio", + "0.10", + "--max-box-aspect-ratio", + "4.0", + "--detection-roi-y-min-ratio", + "0.08", + "--detection-roi-y-max-ratio", + "0.82", + "--cls-show-conf-thres", + "0.60", + "--large-box-area-ratio", + "0.003", + "--large-box-cls-show-conf-thres", + "0.30", + "--show-large-box-name-always", + "false", + "--dense-scene-count-thres", + "30", + "--hide-cls-name-in-dense-scene", + "false", + "--crop-padding-ratio", + "0.10", + "--max-detections", + "75", + "--uniform-species-color", + "true", + "--species-name-remap-min-confidence", + "0.80", + "--binary-dark-species-name", + "东方白鹳", + "--binary-light-species-name", + "白琵鹭", + "--binary-dark-pixel-threshold", + "85", + "--binary-dark-ratio-threshold", + "0.12", + "--dark-override-species-name", + "雁鸭类", + "--dark-override-species-class-id", + "60", + "--dark-override-pixel-threshold", + "75", + "--dark-override-ratio-threshold", + "0.30", + "--dark-override-min-box-area-ratio", + "0.00001", + "--dark-override-max-box-area-ratio", + "0.002", + "--dark-override-roi-y-min-ratio", + "0.60", + "--dark-override-roi-y-max-ratio", + "0.95", + "--region-override-species-name", + "白头鹤", + "--region-override-species-class-id", + "28", + "--region-override-source-species-names", + "白鹤,白头鹤,白枕鹤", + "--region-override-roi-y-min-ratio", + "0.24", + "--region-override-roi-y-max-ratio", + "0.45", + "--enable-crane-appearance-correction", + "--crane-white-ratio-threshold", + "0.33", + "--crane-dark-ratio-threshold", + "0.29", + "--crane-tail-white-ratio-threshold", + "0.14", + "--crane-uncertain-hide-label", + "false", + ] + ) + + self.assertEqual(args.input, "/tmp/input") + self.assertEqual(args.output, "/tmp/output") + self.assertEqual(args.sample_interval, 5) + self.assertFalse(args.recursive) + self.assertFalse(args.enable_detection) + self.assertFalse(args.enable_classification) + self.assertFalse(args.enable_density) + self.assertEqual(args.det_conf_thres, 0.15) + self.assertEqual(args.det_iou_thres, 0.40) + self.assertEqual(args.det_second_conf_thres, 0.45) + self.assertEqual(args.min_box_area_ratio, 0.0001) + self.assertEqual(args.max_box_area_ratio, 0.10) + self.assertEqual(args.max_box_aspect_ratio, 4.0) + self.assertEqual(args.detection_roi_y_min_ratio, 0.08) + self.assertEqual(args.detection_roi_y_max_ratio, 0.82) + self.assertEqual(args.cls_show_conf_thres, 0.60) + self.assertEqual(args.large_box_area_ratio, 0.003) + self.assertEqual(args.large_box_cls_show_conf_thres, 0.30) + self.assertFalse(args.show_large_box_name_always) + self.assertEqual(args.dense_scene_count_thres, 30) + self.assertFalse(args.hide_cls_name_in_dense_scene) + self.assertEqual(args.crop_padding_ratio, 0.10) + self.assertEqual(args.max_detections, 75) + self.assertTrue(args.uniform_species_color) + self.assertEqual(args.species_name_remap_min_confidence, 0.80) + self.assertEqual(args.binary_dark_species_name, "东方白鹳") + self.assertEqual(args.binary_light_species_name, "白琵鹭") + self.assertEqual(args.binary_dark_pixel_threshold, 85) + self.assertEqual(args.binary_dark_ratio_threshold, 0.12) + self.assertEqual(args.dark_override_species_name, "雁鸭类") + self.assertEqual(args.dark_override_species_class_id, 60) + self.assertEqual(args.dark_override_pixel_threshold, 75) + self.assertEqual(args.dark_override_ratio_threshold, 0.30) + self.assertEqual(args.dark_override_min_box_area_ratio, 0.00001) + self.assertEqual(args.dark_override_max_box_area_ratio, 0.002) + self.assertEqual(args.dark_override_roi_y_min_ratio, 0.60) + self.assertEqual(args.dark_override_roi_y_max_ratio, 0.95) + self.assertEqual(args.region_override_species_name, "白头鹤") + self.assertEqual(args.region_override_species_class_id, 28) + self.assertEqual( + args.region_override_source_species_names, + "白鹤,白头鹤,白枕鹤", + ) + self.assertEqual(args.region_override_roi_y_min_ratio, 0.24) + self.assertEqual(args.region_override_roi_y_max_ratio, 0.45) + self.assertTrue(args.enable_crane_appearance_correction) + self.assertEqual(args.crane_white_ratio_threshold, 0.33) + self.assertEqual(args.crane_dark_ratio_threshold, 0.29) + self.assertEqual(args.crane_tail_white_ratio_threshold, 0.14) + self.assertFalse(args.crane_uncertain_hide_label) + + def test_readme_defaults_enable_recursive_and_all_models(self): + args = batch_runner.parse_args(["--input", "/tmp/input"]) + + self.assertTrue(args.recursive) + self.assertEqual(args.enable_detection, settings.enable_detection) + self.assertEqual( + args.enable_classification, + settings.enable_classification, + ) + self.assertEqual(args.enable_density, settings.enable_density) + self.assertEqual(args.sample_interval, settings.video_sample_interval) + self.assertEqual(args.max_box_area_ratio, 1.0) + self.assertEqual( + args.uniform_species_color, + settings.uniform_species_color, + ) + + def test_main_applies_cli_values_before_running_batch(self): + setting_names = ( + "enable_detection", + "enable_classification", + "enable_density", + "det_conf_thres", + "hide_cls_name_in_dense_scene", + "max_detections", + "uniform_species_color", + ) + original_values = { + name: getattr(settings, name) + for name in setting_names + } + + try: + with patch.object( + batch_runner, + "run_batch", + return_value={"success": True}, + ) as run_batch_mock: + result = batch_runner.main( + [ + "--input", + "/tmp/input", + "--output", + "/tmp/output", + "--disable-detection", + "--disable-classification", + "--disable-density", + "--det-conf-thres", + "0.15", + "--hide-cls-name-in-dense-scene", + "false", + "--max-detections", + "75", + "--uniform-species-color", + "true", + ] + ) + + self.assertEqual(result, {"success": True}) + self.assertFalse(settings.enable_detection) + self.assertFalse(settings.enable_classification) + self.assertFalse(settings.enable_density) + self.assertEqual(settings.det_conf_thres, 0.15) + self.assertFalse(settings.hide_cls_name_in_dense_scene) + self.assertEqual(settings.max_detections, 75) + self.assertTrue(settings.uniform_species_color) + run_batch_mock.assert_called_once_with( + input_path="/tmp/input", + output_base="/tmp/output", + recursive=True, + sample_interval=2, + ) + finally: + for name, value in original_values.items(): + object.__setattr__(settings, name, value) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_classifier_probabilities.py b/tests/test_classifier_probabilities.py new file mode 100644 index 0000000..6df780d --- /dev/null +++ b/tests/test_classifier_probabilities.py @@ -0,0 +1,24 @@ +import unittest + +import numpy as np + +from app.inference.convnext_onnx import softmax, to_probabilities + + +class ClassifierProbabilityTests(unittest.TestCase): + def test_existing_probabilities_are_not_softmaxed_again(self): + output = np.array([0.01, 0.09, 0.90], dtype=np.float32) + probabilities = to_probabilities(output) + np.testing.assert_allclose(probabilities, output, atol=1e-6) + + def test_logits_are_converted_with_softmax(self): + logits = np.array([-1.0, 0.0, 3.0], dtype=np.float32) + np.testing.assert_allclose( + to_probabilities(logits), + softmax(logits), + atol=1e-6, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_pipeline_switches.py b/tests/test_pipeline_switches.py new file mode 100644 index 0000000..1deb5c1 --- /dev/null +++ b/tests/test_pipeline_switches.py @@ -0,0 +1,114 @@ +import unittest +from unittest.mock import patch + +import numpy as np + +from app.config import settings +from app.pipeline import BirdCascadePipeline + + +class PipelineSwitchTests(unittest.TestCase): + setting_names = ( + "enable_detection", + "enable_classification", + "enable_density", + "uniform_species_color", + ) + + def setUp(self): + self.original_values = { + name: getattr(settings, name) + for name in self.setting_names + } + + def tearDown(self): + for name, value in self.original_values.items(): + object.__setattr__(settings, name, value) + + def test_all_models_can_be_disabled_without_loading_onnx(self): + object.__setattr__(settings, "enable_detection", False) + object.__setattr__(settings, "enable_classification", False) + object.__setattr__(settings, "enable_density", False) + + with ( + patch("app.pipeline.YOLOv5ONNXDetector") as detector_mock, + patch("app.pipeline.ConvNeXtONNXClassifier") as classifier_mock, + patch("app.pipeline.DensityEstimatorONNX") as density_mock, + ): + pipeline = BirdCascadePipeline() + + self.assertIsNone(pipeline.detector) + self.assertIsNone(pipeline.classifier) + self.assertIsNone(pipeline.density_estimator) + detector_mock.assert_not_called() + classifier_mock.assert_not_called() + density_mock.assert_not_called() + + output = pipeline.predict( + np.zeros((32, 32, 3), dtype=np.uint8), + ) + self.assertEqual(output["results"], []) + self.assertIsNone(output["density"]) + + def test_detection_without_classification_keeps_detection_fields(self): + object.__setattr__(settings, "enable_detection", True) + object.__setattr__(settings, "enable_classification", False) + object.__setattr__(settings, "enable_density", False) + + fake_detector_result = [ + { + "bbox_xyxy": [4.0, 4.0, 20.0, 20.0], + "det_confidence": 0.9, + "det_class_id": 0, + "det_class_name": "bird", + } + ] + + with ( + patch( + "app.pipeline.read_classes", + return_value=["bird"], + ), + patch( + "app.pipeline.os.path.exists", + return_value=True, + ), + patch("app.pipeline.YOLOv5ONNXDetector") as detector_class, + patch("app.pipeline.ConvNeXtONNXClassifier") as classifier_class, + ): + detector_class.return_value.infer.return_value = fake_detector_result + pipeline = BirdCascadePipeline() + output = pipeline.predict( + np.zeros((32, 32, 3), dtype=np.uint8), + ) + + classifier_class.assert_not_called() + self.assertEqual(len(output["results"]), 1) + result = output["results"][0] + self.assertEqual(result["det_class_name"], "bird") + self.assertEqual(result["cls_class_id"], -1) + self.assertEqual(result["cls_class_name"], "") + self.assertEqual(result["cls_confidence"], 0.0) + self.assertFalse(result["show_cls_name"]) + + def test_uniform_species_color_uses_one_color_for_every_class(self): + pipeline = object.__new__(BirdCascadePipeline) + + object.__setattr__(settings, "uniform_species_color", False) + first_color = pipeline._get_color_by_class_id(1) + second_color = pipeline._get_color_by_class_id(2) + self.assertNotEqual(first_color, second_color) + + object.__setattr__(settings, "uniform_species_color", True) + self.assertEqual( + pipeline._get_color_by_class_id(1), + pipeline._get_color_by_class_id(2), + ) + self.assertEqual( + pipeline._get_color_by_class_id(999), + (0, 190, 0), + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_video_smoother.py b/tests/test_video_smoother.py new file mode 100644 index 0000000..f2a93f1 --- /dev/null +++ b/tests/test_video_smoother.py @@ -0,0 +1,38 @@ +import unittest + +from app.video_pipeline import VideoResultSmoother + + +def result(class_id, class_name, confidence, box): + return { + "bbox_xyxy": box, + "cls_class_id": class_id, + "cls_class_name": class_name, + "cls_confidence": confidence, + "show_cls_name": True, + } + + +class VideoResultSmootherTests(unittest.TestCase): + def test_one_low_confidence_label_does_not_flip_track(self): + smoother = VideoResultSmoother(hold_frames=4, alpha=0.5) + smoother.update([result(1, "白鹤", 0.95, [0, 0, 20, 20])], 0) + second = smoother.update( + [result(2, "白枕鹤", 0.20, [1, 1, 21, 21])], + 1, + ) + self.assertEqual(second[0]["cls_class_name"], "白鹤") + self.assertNotIn("_smoother_class_votes", second[0]) + + def test_matching_does_not_require_same_predicted_class(self): + smoother = VideoResultSmoother(hold_frames=4, alpha=0.5) + smoother.update([result(1, "白鹤", 0.9, [0, 0, 20, 20])], 0) + second = smoother.update( + [result(2, "白枕鹤", 0.9, [2, 2, 22, 22])], + 1, + ) + self.assertEqual(second[0]["bbox_xyxy"], [1.0, 1.0, 21.0, 21.0]) + + +if __name__ == "__main__": + unittest.main()