initial commit

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