105 lines
3.0 KiB
Python
105 lines
3.0 KiB
Python
import os
|
|
from dataclasses import dataclass
|
|
|
|
|
|
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:
|
|
# Docker 内默认模型路径
|
|
det_model_path: str = _get_str(
|
|
"DET_MODEL_PATH",
|
|
"/app/models/YNDC_20251010_yolov5_optimize.onnx",
|
|
)
|
|
cls_model_path: str = _get_str(
|
|
"CLS_MODEL_PATH",
|
|
"/app/models/HHK_20260525_mmcls_optimize.onnx",
|
|
)
|
|
det_classes_path: str = _get_str(
|
|
"DET_CLASSES_PATH",
|
|
"/app/models/YNDC_20251010_yolov5_optimize.txt",
|
|
)
|
|
cls_classes_path: str = _get_str(
|
|
"CLS_CLASSES_PATH",
|
|
"/app/models/HHK_20260525_mmcls_optimize.txt",
|
|
)
|
|
|
|
# 输入尺寸
|
|
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)
|
|
|
|
# 目标裁剪外扩比例
|
|
crop_padding_ratio: float = _get_float("CROP_PADDING_RATIO", 0.08)
|
|
|
|
# 最多保留目标数量
|
|
max_detections: int = _get_int("MAX_DETECTIONS", 50)
|
|
|
|
# 当前环境没有 onnxruntime-gpu 时，建议用 cpu
|
|
device: str = _get_str("DEVICE", "cpu")
|
|
|
|
# 输出目录
|
|
output_dir: str = _get_str("OUTPUT_DIR", "/app/runs")
|
|
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 类别 0，也就是“鸟”
|
|
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() |