309 lines
11 KiB
Python
309 lines
11 KiB
Python
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",
|
|
} |