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