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