initial commit

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",
+        }