115 lines
3.9 KiB
Python
115 lines
3.9 KiB
Python
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",
|
|
} |