105 lines
2.6 KiB
Python
105 lines
2.6 KiB
Python
# roi_filter.py
|
|
# =========================
|
|
# ROI 屏蔽 / 过滤模块
|
|
# =========================
|
|
|
|
import cv2
|
|
import yaml
|
|
import os
|
|
|
|
# =========================
|
|
# 全局变量（常量风格）
|
|
# =========================
|
|
CAMERA_ID = None
|
|
IGNORE_ROIS = []
|
|
|
|
|
|
# =========================
|
|
# 读取配置
|
|
# =========================
|
|
def load_ignore_rois(camera_id, config_path="config.yaml"):
|
|
"""
|
|
从配置文件中读取指定摄像头的 ROI 忽略区域
|
|
"""
|
|
if not os.path.exists(config_path):
|
|
print(f"[WARN] config file not found: {config_path}")
|
|
return []
|
|
|
|
with open(config_path, "r", encoding="utf-8") as f:
|
|
cfg = yaml.safe_load(f) or {}
|
|
|
|
return cfg.get("roi", {}).get("ignore_rois", {}).get(camera_id, [])
|
|
|
|
|
|
# =========================
|
|
# 初始化（由主程序调用一次）
|
|
# =========================
|
|
def init_ignore_rois(camera_id, config_path="config.yaml"):
|
|
"""
|
|
初始化 ROI 忽略区域（必须调用）
|
|
"""
|
|
global CAMERA_ID, IGNORE_ROIS
|
|
CAMERA_ID = camera_id
|
|
IGNORE_ROIS = load_ignore_rois(camera_id, config_path)
|
|
|
|
print(f"[ROI] camera_id={camera_id}, ignore_rois={IGNORE_ROIS}")
|
|
|
|
|
|
# =========================
|
|
# 判断是否在忽略区域
|
|
# =========================
|
|
def is_in_ignore_roi(bbox, mode="center"):
|
|
"""
|
|
判断检测框是否需要被屏蔽
|
|
|
|
:param bbox: (x1, y1, x2, y2)
|
|
:param mode:
|
|
- center : 中心点判断（推荐）
|
|
- iou : 有交集即忽略
|
|
:return: True = 忽略
|
|
"""
|
|
if not IGNORE_ROIS:
|
|
return False
|
|
|
|
x1, y1, x2, y2 = bbox
|
|
|
|
if mode == "center":
|
|
cx = (x1 + x2) // 2
|
|
cy = (y1 + y2) // 2
|
|
for rx1, ry1, rx2, ry2 in IGNORE_ROIS:
|
|
if rx1 <= cx <= rx2 and ry1 <= cy <= ry2:
|
|
return True
|
|
return False
|
|
|
|
if mode == "iou":
|
|
for rx1, ry1, rx2, ry2 in IGNORE_ROIS:
|
|
ix1 = max(x1, rx1)
|
|
iy1 = max(y1, ry1)
|
|
ix2 = min(x2, rx2)
|
|
iy2 = min(y2, ry2)
|
|
if ix1 < ix2 and iy1 < iy2:
|
|
return True
|
|
return False
|
|
|
|
return False
|
|
|
|
|
|
# =========================
|
|
# 调试：画出忽略区域
|
|
# =========================
|
|
def draw_ignore_rois(img):
|
|
"""
|
|
调试用：在画面中把 ROI 忽略区域画出来
|
|
"""
|
|
for x1, y1, x2, y2 in IGNORE_ROIS:
|
|
cv2.rectangle(img, (x1, y1), (x2, y2), (128, 128, 128), 2)
|
|
cv2.putText(
|
|
img,
|
|
"IGNORE ROI",
|
|
(x1 + 5, y1 + 20),
|
|
cv2.FONT_HERSHEY_SIMPLEX,
|
|
0.6,
|
|
(128, 128, 128),
|
|
2
|
|
)
|