495 lines
12 KiB
Python
495 lines
12 KiB
Python
import os
|
||
from dataclasses import dataclass
|
||
|
||
|
||
# =========================
|
||
# 项目根目录
|
||
# =========================
|
||
# 例如：
|
||
# 宿主机：
|
||
# /data2/zhangw_data/Bird_detect_cascade_service
|
||
#
|
||
# Docker：
|
||
# /Bird_detect_cascade_service
|
||
#
|
||
# 会自动适配，不需要写死绝对路径
|
||
# =========================
|
||
BASE_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
|
||
|
||
MODELS_DIR = os.path.join(BASE_DIR, "models")
|
||
RUNS_DIR = os.path.join(BASE_DIR, "runs")
|
||
|
||
|
||
def _get_str(name: str, default: str) -> str:
|
||
return os.getenv(name, default)
|
||
|
||
|
||
def _get_int(name: str, default: int) -> int:
|
||
value = os.getenv(name)
|
||
if value is None or value == "":
|
||
return default
|
||
return int(value)
|
||
|
||
|
||
def _get_float(name: str, default: float) -> float:
|
||
value = os.getenv(name)
|
||
if value is None or value == "":
|
||
return default
|
||
return float(value)
|
||
|
||
|
||
def _get_bool(name: str, default: bool) -> bool:
|
||
value = os.getenv(name)
|
||
if value is None or value == "":
|
||
return default
|
||
return value.lower() in ["1", "true", "yes", "y"]
|
||
|
||
|
||
def _get_int_list(name: str, default: str) -> list[int]:
|
||
value = os.getenv(name, default)
|
||
|
||
if value is None or value.strip() == "":
|
||
return []
|
||
|
||
return [
|
||
int(x.strip())
|
||
for x in value.split(",")
|
||
if x.strip() != ""
|
||
]
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class Settings:
|
||
|
||
# =========================
|
||
# 模型开关
|
||
# =========================
|
||
enable_detection: bool = _get_bool(
|
||
"ENABLE_DETECTION",
|
||
True,
|
||
)
|
||
|
||
enable_classification: bool = _get_bool(
|
||
"ENABLE_CLASSIFICATION",
|
||
True,
|
||
)
|
||
|
||
# =========================
|
||
# YOLOv5 模型
|
||
# =========================
|
||
det_model_path: str = _get_str(
|
||
"DET_MODEL_PATH",
|
||
os.path.join(
|
||
MODELS_DIR,
|
||
"YNDC_20251010_yolov5_optimize.onnx",
|
||
),
|
||
)
|
||
|
||
det_classes_path: str = _get_str(
|
||
"DET_CLASSES_PATH",
|
||
os.path.join(
|
||
MODELS_DIR,
|
||
"YNDC_20251010_yolov5_optimize.txt",
|
||
),
|
||
)
|
||
|
||
# =========================
|
||
# ConvNeXt 分类模型
|
||
# =========================
|
||
cls_model_path: str = _get_str(
|
||
"CLS_MODEL_PATH",
|
||
os.path.join(
|
||
MODELS_DIR,
|
||
"JXPYH_20260204_mmcls_optimize.onnx",
|
||
),
|
||
)
|
||
|
||
cls_classes_path: str = _get_str(
|
||
"CLS_CLASSES_PATH",
|
||
os.path.join(
|
||
MODELS_DIR,
|
||
"JXPYH_20260204_mmcls_optimize.txt",
|
||
),
|
||
)
|
||
|
||
# =========================
|
||
# 密度估计模型
|
||
# =========================
|
||
enable_density: bool = _get_bool(
|
||
"ENABLE_DENSITY",
|
||
True,
|
||
)
|
||
|
||
density_model_path: str = _get_str(
|
||
"DENSITY_MODEL_PATH",
|
||
os.path.join(
|
||
MODELS_DIR,
|
||
"bird_density.onnx",
|
||
),
|
||
)
|
||
|
||
# 如果模型输入固定，例如 400×400
|
||
# 可以设置 DENSITY_IMG_SIZE=400
|
||
# 默认 0 表示自动读取 ONNX 输入尺寸
|
||
density_img_size: int = _get_int(
|
||
"DENSITY_IMG_SIZE",
|
||
0,
|
||
)
|
||
|
||
# =========================
|
||
# 输入尺寸
|
||
# =========================
|
||
det_img_size: int = _get_int(
|
||
"DET_IMG_SIZE",
|
||
640,
|
||
)
|
||
|
||
cls_img_size: int = _get_int(
|
||
"CLS_IMG_SIZE",
|
||
224,
|
||
)
|
||
|
||
# =========================
|
||
# YOLO 阈值
|
||
# =========================
|
||
det_conf_thres: float = _get_float(
|
||
"DET_CONF_THRES",
|
||
0.25,
|
||
)
|
||
|
||
det_iou_thres: float = _get_float(
|
||
"DET_IOU_THRES",
|
||
0.45,
|
||
)
|
||
|
||
# 二次过滤 YOLO 检测框，避免把明显错误的框传给 ConvNeXt
|
||
det_second_conf_thres: float = _get_float(
|
||
"DET_SECOND_CONF_THRES",
|
||
0.35,
|
||
)
|
||
|
||
# 最小框面积占整图比例，太小的噪声框过滤
|
||
min_box_area_ratio: float = _get_float(
|
||
"MIN_BOX_AREA_RATIO",
|
||
0.00002,
|
||
)
|
||
|
||
# 最大框面积占整图比例，太大的误检框过滤
|
||
max_box_area_ratio: float = _get_float(
|
||
"MAX_BOX_AREA_RATIO",
|
||
1.0,
|
||
)
|
||
|
||
# 最大长宽比，过细过长的框过滤
|
||
max_box_aspect_ratio: float = _get_float(
|
||
"MAX_BOX_ASPECT_RATIO",
|
||
6.0,
|
||
)
|
||
|
||
# 可选检测区域：按检测框中心点的纵向位置过滤。
|
||
# 默认 0~1 表示整幅画面；可排除固定字幕、时间戳或岸边噪声区域。
|
||
detection_roi_y_min_ratio: float = _get_float(
|
||
"DETECTION_ROI_Y_MIN_RATIO",
|
||
0.0,
|
||
)
|
||
|
||
detection_roi_y_max_ratio: float = _get_float(
|
||
"DETECTION_ROI_Y_MAX_RATIO",
|
||
1.0,
|
||
)
|
||
|
||
# 分类置信度低于该值时，不显示物种名称
|
||
cls_show_conf_thres: float = _get_float(
|
||
"CLS_SHOW_CONF_THRES",
|
||
0.50,
|
||
)
|
||
|
||
# 大目标判断阈值：框面积占整图比例超过该值，认为是近景大目标
|
||
large_box_area_ratio: float = _get_float(
|
||
"LARGE_BOX_AREA_RATIO",
|
||
0.002,
|
||
)
|
||
|
||
# 大目标分类显示阈值，近景大目标即使分类置信度稍低，也显示物种名
|
||
large_box_cls_show_conf_thres: float = _get_float(
|
||
"LARGE_BOX_CLS_SHOW_CONF_THRES",
|
||
0.25,
|
||
)
|
||
|
||
# 数量达到该阈值时，认为是密集场景，不显示物种名
|
||
dense_scene_count_thres: int = _get_int(
|
||
"DENSE_SCENE_COUNT_THRES",
|
||
20,
|
||
)
|
||
|
||
# 大目标强制显示物种名，只要分类结果非空即可
|
||
show_large_box_name_always: bool = _get_bool(
|
||
"SHOW_LARGE_BOX_NAME_ALWAYS",
|
||
True,
|
||
)
|
||
|
||
# 是否在密集场景隐藏所有物种名
|
||
hide_cls_name_in_dense_scene: bool = _get_bool(
|
||
"HIDE_CLS_NAME_IN_DENSE_SCENE",
|
||
True,
|
||
)
|
||
|
||
# 所有物种是否统一使用同一种标签和检测框颜色
|
||
uniform_species_color: bool = _get_bool(
|
||
"UNIFORM_SPECIES_COLOR",
|
||
False,
|
||
)
|
||
|
||
# 可选：整段素材已人工确认只有一个物种时，覆盖细粒度模型标签。
|
||
# 留空表示使用模型输出。
|
||
forced_species_name: str = _get_str(
|
||
"FORCED_SPECIES_NAME",
|
||
"",
|
||
).strip()
|
||
|
||
forced_species_class_id: int = _get_int(
|
||
"FORCED_SPECIES_CLASS_ID",
|
||
-1,
|
||
)
|
||
|
||
# 只对达到该框面积占比的目标应用人工物种覆盖。
|
||
# 可用来统一主要大型鸟类，同时保留画面中的小型伴生鸟类。
|
||
forced_species_min_box_area_ratio: float = _get_float(
|
||
"FORCED_SPECIES_MIN_BOX_AREA_RATIO",
|
||
0.0,
|
||
)
|
||
|
||
# 视频级人工复核映射，格式示例：
|
||
# 东方白鹳=白鹤,白鹭=白鹤
|
||
species_name_remap: str = _get_str(
|
||
"SPECIES_NAME_REMAP",
|
||
"",
|
||
).strip()
|
||
|
||
# 只有原始细粒度分类置信度达到该阈值时才应用标签映射。
|
||
# 默认 0 保持原有行为；混合物种场景可提高该值以避免低置信度误映射。
|
||
species_name_remap_min_confidence: float = _get_float(
|
||
"SPECIES_NAME_REMAP_MIN_CONFIDENCE",
|
||
0.0,
|
||
)
|
||
|
||
# 留空表示允许全部；非空时，只有白名单中的物种显示名称和物种色。
|
||
allowed_species_names: str = _get_str(
|
||
"ALLOWED_SPECIES_NAMES",
|
||
"",
|
||
).strip()
|
||
|
||
# 可选双物种明暗规则。两个名称都非空时，不运行细粒度分类模型，
|
||
# 而是根据检测框内暗色像素比例区分两个外观差异明显的物种。
|
||
binary_dark_species_name: str = _get_str(
|
||
"BINARY_DARK_SPECIES_NAME",
|
||
"",
|
||
).strip()
|
||
|
||
binary_light_species_name: str = _get_str(
|
||
"BINARY_LIGHT_SPECIES_NAME",
|
||
"",
|
||
).strip()
|
||
|
||
binary_dark_pixel_threshold: int = _get_int(
|
||
"BINARY_DARK_PIXEL_THRESHOLD",
|
||
85,
|
||
)
|
||
|
||
binary_dark_ratio_threshold: float = _get_float(
|
||
"BINARY_DARK_RATIO_THRESHOLD",
|
||
0.12,
|
||
)
|
||
|
||
# 可选：对画面中的深色目标做人工标签覆盖。
|
||
# 适合模型把黑色雁鸭类误判成白骨顶等细粒度类别的场景。
|
||
dark_override_species_name: str = _get_str(
|
||
"DARK_OVERRIDE_SPECIES_NAME",
|
||
"",
|
||
).strip()
|
||
|
||
dark_override_species_class_id: int = _get_int(
|
||
"DARK_OVERRIDE_SPECIES_CLASS_ID",
|
||
-1,
|
||
)
|
||
|
||
dark_override_pixel_threshold: int = _get_int(
|
||
"DARK_OVERRIDE_PIXEL_THRESHOLD",
|
||
80,
|
||
)
|
||
|
||
dark_override_ratio_threshold: float = _get_float(
|
||
"DARK_OVERRIDE_RATIO_THRESHOLD",
|
||
0.35,
|
||
)
|
||
|
||
dark_override_min_box_area_ratio: float = _get_float(
|
||
"DARK_OVERRIDE_MIN_BOX_AREA_RATIO",
|
||
0.0,
|
||
)
|
||
|
||
dark_override_max_box_area_ratio: float = _get_float(
|
||
"DARK_OVERRIDE_MAX_BOX_AREA_RATIO",
|
||
1.0,
|
||
)
|
||
|
||
dark_override_roi_y_min_ratio: float = _get_float(
|
||
"DARK_OVERRIDE_ROI_Y_MIN_RATIO",
|
||
0.0,
|
||
)
|
||
|
||
dark_override_roi_y_max_ratio: float = _get_float(
|
||
"DARK_OVERRIDE_ROI_Y_MAX_RATIO",
|
||
1.0,
|
||
)
|
||
|
||
# 可选：按画面纵向区域做人工标签覆盖。
|
||
# 用于同一视频中上/中/下区域物种分布稳定的场景。
|
||
region_override_species_name: str = _get_str(
|
||
"REGION_OVERRIDE_SPECIES_NAME",
|
||
"",
|
||
).strip()
|
||
|
||
region_override_species_class_id: int = _get_int(
|
||
"REGION_OVERRIDE_SPECIES_CLASS_ID",
|
||
-1,
|
||
)
|
||
|
||
region_override_source_species_names: str = _get_str(
|
||
"REGION_OVERRIDE_SOURCE_SPECIES_NAMES",
|
||
"",
|
||
).strip()
|
||
|
||
region_override_roi_y_min_ratio: float = _get_float(
|
||
"REGION_OVERRIDE_ROI_Y_MIN_RATIO",
|
||
0.0,
|
||
)
|
||
|
||
region_override_roi_y_max_ratio: float = _get_float(
|
||
"REGION_OVERRIDE_ROI_Y_MAX_RATIO",
|
||
1.0,
|
||
)
|
||
|
||
# 可选：鹤类混淆纠正。针对白鹤/白头鹤/白枕鹤相互误判，
|
||
# 根据检测框内亮白像素和暗色像素比例做保守二次校验。
|
||
enable_crane_appearance_correction: bool = _get_bool(
|
||
"ENABLE_CRANE_APPEARANCE_CORRECTION",
|
||
False,
|
||
)
|
||
|
||
crane_white_ratio_threshold: float = _get_float(
|
||
"CRANE_WHITE_RATIO_THRESHOLD",
|
||
0.30,
|
||
)
|
||
|
||
crane_dark_ratio_threshold: float = _get_float(
|
||
"CRANE_DARK_RATIO_THRESHOLD",
|
||
0.28,
|
||
)
|
||
|
||
crane_tail_white_ratio_threshold: float = _get_float(
|
||
"CRANE_TAIL_WHITE_RATIO_THRESHOLD",
|
||
0.16,
|
||
)
|
||
|
||
crane_uncertain_hide_label: bool = _get_bool(
|
||
"CRANE_UNCERTAIN_HIDE_LABEL",
|
||
True,
|
||
)
|
||
|
||
# =========================
|
||
# 裁剪外扩比例
|
||
# =========================
|
||
crop_padding_ratio: float = _get_float(
|
||
"CROP_PADDING_RATIO",
|
||
0.08,
|
||
)
|
||
|
||
# =========================
|
||
# 最多检测目标数
|
||
# =========================
|
||
max_detections: int = _get_int(
|
||
"MAX_DETECTIONS",
|
||
50,
|
||
)
|
||
|
||
# =========================
|
||
# 推理设备
|
||
# cpu / cuda
|
||
# =========================
|
||
device: str = _get_str(
|
||
"DEVICE",
|
||
"cpu",
|
||
)
|
||
|
||
# =========================
|
||
# 输出目录
|
||
# =========================
|
||
output_dir: str = _get_str(
|
||
"OUTPUT_DIR",
|
||
RUNS_DIR,
|
||
)
|
||
|
||
image_output_subdir: str = _get_str(
|
||
"IMAGE_OUTPUT_SUBDIR",
|
||
"images",
|
||
)
|
||
|
||
video_output_subdir: str = _get_str(
|
||
"VIDEO_OUTPUT_SUBDIR",
|
||
"video",
|
||
)
|
||
|
||
# =========================
|
||
# 视频抽帧间隔
|
||
# =========================
|
||
video_sample_interval: int = _get_int(
|
||
"VIDEO_SAMPLE_INTERVAL",
|
||
2,
|
||
)
|
||
|
||
# =========================
|
||
# 是否保存画框视频
|
||
# =========================
|
||
save_annotated_video: bool = _get_bool(
|
||
"SAVE_ANNOTATED_VIDEO",
|
||
True,
|
||
)
|
||
|
||
# =========================
|
||
# 只保留 YOLO 类别
|
||
# 默认只保留 bird
|
||
# =========================
|
||
target_det_class_ids: list[int] = None
|
||
|
||
# =========================
|
||
# 视频框平滑
|
||
# =========================
|
||
smooth_video_boxes: bool = _get_bool(
|
||
"SMOOTH_VIDEO_BOXES",
|
||
True,
|
||
)
|
||
|
||
smooth_hold_frames: int = _get_int(
|
||
"SMOOTH_HOLD_FRAMES",
|
||
3,
|
||
)
|
||
|
||
box_smooth_alpha: float = _get_float(
|
||
"BOX_SMOOTH_ALPHA",
|
||
0.25,
|
||
)
|
||
|
||
def __post_init__(self):
|
||
object.__setattr__(
|
||
self,
|
||
"target_det_class_ids",
|
||
_get_int_list(
|
||
"TARGET_DET_CLASS_IDS",
|
||
"0",
|
||
),
|
||
)
|
||
|
||
|
||
settings = Settings()
|