795 lines
28 KiB
Python
795 lines
28 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, 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",
|
||
}
|