77 lines
2.9 KiB
Python
77 lines
2.9 KiB
Python
# roi_filter.py
|
|
import yaml
|
|
import cv2
|
|
import numpy as np
|
|
|
|
class ROIManager:
|
|
def __init__(self, config_path="../config.yaml"):
|
|
self.config_path = config_path
|
|
self.roi_config = self.load_config()
|
|
|
|
def load_config(self):
|
|
try:
|
|
with open(self.config_path, 'r', encoding='utf-8') as f:
|
|
config = yaml.safe_load(f)
|
|
roi_config = config.get('roi_filter', {})
|
|
return roi_config
|
|
except FileNotFoundError:
|
|
print(f"配置文件 {self.config_path} 不存在")
|
|
return {"enabled": False, "regions": []}
|
|
except yaml.YAMLError as e:
|
|
print(f"解析YAML配置文件出错: {e}")
|
|
return {"enabled": False, "regions": []}
|
|
|
|
def is_in_any_roi(self, center_x, center_y):
|
|
if not self.roi_config.get('enabled', False):
|
|
return False, None
|
|
|
|
regions = self.roi_config.get('regions', [])
|
|
for region in regions:
|
|
if (region.get('enabled', True) and
|
|
region['x_min'] <= center_x <= region['x_max'] and
|
|
region['y_min'] <= center_y <= region['y_max']):
|
|
return True, region.get('name', 'unknown')
|
|
return False, None
|
|
|
|
def filter_detections(self, detections, img_width, img_height):
|
|
if not self.roi_config.get('enabled', False):
|
|
return detections
|
|
|
|
filtered_detections = []
|
|
for det in detections:
|
|
center_x = ((det[0] + det[2]) / 2) / img_width
|
|
center_y = ((det[1] + det[3]) / 2) / img_height
|
|
|
|
in_roi, roi_name = self.is_in_any_roi(center_x, center_y)
|
|
if not in_roi:
|
|
filtered_detections.append(det)
|
|
else:
|
|
print(f"过滤检测框: 位置({center_x:.2f}, {center_y:.2f}) 在ROI区域 {roi_name}")
|
|
|
|
return filtered_detections
|
|
|
|
def draw_roi_regions(self, frame):
|
|
if not self.roi_config.get('enabled', False):
|
|
return frame
|
|
|
|
height, width = frame.shape[:2]
|
|
regions = self.roi_config.get('regions', [])
|
|
|
|
frame_with_roi = frame.copy()
|
|
|
|
for region in regions:
|
|
if region.get('enabled', True):
|
|
x_min = int(region['x_min'] * width)
|
|
x_max = int(region['x_max'] * width)
|
|
y_min = int(region['y_min'] * height)
|
|
y_max = int(region['y_max'] * height)
|
|
|
|
cv2.rectangle(frame_with_roi, (x_min, y_min), (x_max, y_max), (0, 0, 255), 2)
|
|
cv2.putText(frame_with_roi, region['name'], (x_min, y_min - 10),
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1)
|
|
|
|
return frame_with_roi
|
|
|
|
def draw_roi_on_frame(frame, config_path="../config.yaml"):
|
|
roi_manager = ROIManager(config_path)
|
|
return roi_manager.draw_roi_regions(frame) |