- 添加完整的项目文档(README.md, design.md, CLAUDE.md) - 实现核心检测算法:ROI管理、峰值检测、帧间稳定 - 支持实时摄像头检测和视频文件处理 - 包含图像预处理:去雾、几何校正、图像增强 - 提供多种输出格式:JSON、CSV、矩阵、文本 - 实现双阈值检测算法适应雾天环境 - 添加ROI标定工具和配置文件管理 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <noreply@anthropic.com>
290 lines
9.9 KiB
Python
290 lines
9.9 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
ROI标定工具
|
||
为YantaiVisionX项目提供交互式ROI区域标定功能
|
||
支持3排×6列共18个LED灯的ROI区域标定
|
||
"""
|
||
|
||
import cv2
|
||
import numpy as np
|
||
import yaml
|
||
import sys
|
||
import os
|
||
from typing import List, Tuple, Optional, Dict, Any
|
||
|
||
# 添加项目根目录到路径
|
||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))
|
||
|
||
from src.roi_detection.roi_manager import ROIManager, ROIRegion
|
||
|
||
|
||
class ROICalibrationTool:
|
||
"""
|
||
交互式ROI标定工具
|
||
"""
|
||
|
||
def __init__(self):
|
||
self.image = None
|
||
self.display_image = None
|
||
self.roi_points = {} # roi_name -> center_point
|
||
self.current_roi_name = None
|
||
self.roi_size = (60, 60) # 默认ROI大小 (width, height)
|
||
self.core_ratio = 0.6 # 核心区域占ROI的比例
|
||
|
||
# 3排×6列的LED布局
|
||
self.rows = 3
|
||
self.cols = 6
|
||
self.roi_names = []
|
||
for r in range(1, self.rows + 1):
|
||
for c in range(1, self.cols + 1):
|
||
self.roi_names.append(f"R{r}C{c}")
|
||
|
||
self.current_roi_index = 0
|
||
self.window_name = "ROI标定工具"
|
||
|
||
def load_calibration_image(self, image_path: str) -> bool:
|
||
"""
|
||
加载标定图像
|
||
|
||
Args:
|
||
image_path: 图像文件路径
|
||
|
||
Returns:
|
||
bool: 加载是否成功
|
||
"""
|
||
try:
|
||
self.image = cv2.imread(image_path)
|
||
if self.image is None:
|
||
print(f"无法加载图像: {image_path}")
|
||
return False
|
||
|
||
self.display_image = self.image.copy()
|
||
print(f"图像加载成功: {image_path}")
|
||
print(f"图像尺寸: {self.image.shape[1]}x{self.image.shape[0]}")
|
||
return True
|
||
|
||
except Exception as e:
|
||
print(f"加载图像失败: {e}")
|
||
return False
|
||
|
||
def mouse_callback(self, event, x, y, flags, param):
|
||
"""
|
||
鼠标事件回调函数
|
||
"""
|
||
if event == cv2.EVENT_LBUTTONDOWN:
|
||
if self.current_roi_index < len(self.roi_names):
|
||
roi_name = self.roi_names[self.current_roi_index]
|
||
self.roi_points[roi_name] = (x, y)
|
||
print(f"标定 {roi_name}: 中心点({x}, {y})")
|
||
|
||
self.current_roi_index += 1
|
||
self.update_display()
|
||
|
||
if self.current_roi_index >= len(self.roi_names):
|
||
print("所有ROI区域标定完成!按'S'保存,按'R'重新开始")
|
||
|
||
def update_display(self):
|
||
"""
|
||
更新显示图像
|
||
"""
|
||
self.display_image = self.image.copy()
|
||
|
||
# 绘制已标定的ROI
|
||
for i, roi_name in enumerate(self.roi_names):
|
||
if roi_name in self.roi_points:
|
||
center_x, center_y = self.roi_points[roi_name]
|
||
|
||
# 计算ROI矩形
|
||
half_w = self.roi_size[0] // 2
|
||
half_h = self.roi_size[1] // 2
|
||
|
||
# 绘制ROI外框(蓝色)
|
||
cv2.rectangle(self.display_image,
|
||
(center_x - half_w, center_y - half_h),
|
||
(center_x + half_w, center_y + half_h),
|
||
(255, 0, 0), 2)
|
||
|
||
# 绘制核心区域(绿色)
|
||
core_half_w = int(half_w * self.core_ratio)
|
||
core_half_h = int(half_h * self.core_ratio)
|
||
cv2.rectangle(self.display_image,
|
||
(center_x - core_half_w, center_y - core_half_h),
|
||
(center_x + core_half_w, center_y + core_half_h),
|
||
(0, 255, 0), 1)
|
||
|
||
# 绘制中心点
|
||
cv2.circle(self.display_image, (center_x, center_y), 3, (0, 0, 255), -1)
|
||
|
||
# 添加ROI名称
|
||
cv2.putText(self.display_image, roi_name,
|
||
(center_x - 15, center_y - half_h - 5),
|
||
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 0), 1)
|
||
|
||
# 显示当前要标定的ROI提示
|
||
if self.current_roi_index < len(self.roi_names):
|
||
current_roi = self.roi_names[self.current_roi_index]
|
||
info_text = f"请点击标定 {current_roi} ({self.current_roi_index + 1}/{len(self.roi_names)})"
|
||
cv2.putText(self.display_image, info_text, (10, 30),
|
||
cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 255), 2)
|
||
|
||
# 显示操作提示
|
||
help_text = "操作: 鼠标左键-标定点 | R-重置 | S-保存 | Q-退出 | +/-调整ROI大小"
|
||
cv2.putText(self.display_image, help_text, (10, self.image.shape[0] - 10),
|
||
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
|
||
|
||
cv2.imshow(self.window_name, self.display_image)
|
||
|
||
def adjust_roi_size(self, delta: int):
|
||
"""
|
||
调整ROI大小
|
||
|
||
Args:
|
||
delta: 尺寸变化量
|
||
"""
|
||
new_w = max(20, self.roi_size[0] + delta)
|
||
new_h = max(20, self.roi_size[1] + delta)
|
||
self.roi_size = (new_w, new_h)
|
||
print(f"ROI大小调整为: {self.roi_size}")
|
||
self.update_display()
|
||
|
||
def reset_calibration(self):
|
||
"""
|
||
重置标定
|
||
"""
|
||
self.roi_points.clear()
|
||
self.current_roi_index = 0
|
||
print("标定已重置")
|
||
self.update_display()
|
||
|
||
def generate_roi_config(self) -> Dict[str, Any]:
|
||
"""
|
||
生成ROI配置
|
||
|
||
Returns:
|
||
Dict[str, Any]: ROI配置字典
|
||
"""
|
||
config = {
|
||
'led_matrix': {
|
||
'rows': self.rows,
|
||
'cols': self.cols,
|
||
'total_leds': self.rows * self.cols
|
||
},
|
||
'roi_regions': {}
|
||
}
|
||
|
||
for roi_name, center in self.roi_points.items():
|
||
center_x, center_y = center
|
||
half_w = self.roi_size[0] // 2
|
||
half_h = self.roi_size[1] // 2
|
||
|
||
# ROI边界框
|
||
roi_box = (center_x - half_w, center_y - half_h,
|
||
self.roi_size[0], self.roi_size[1])
|
||
|
||
# 核心区域
|
||
core_half_w = int(half_w * self.core_ratio)
|
||
core_half_h = int(half_h * self.core_ratio)
|
||
core_area = (center_x - core_half_w, center_y - core_half_h,
|
||
core_half_w * 2, core_half_h * 2)
|
||
|
||
config['roi_regions'][roi_name] = {
|
||
'center': [center_x, center_y],
|
||
'roi_box': list(roi_box),
|
||
'core_area': list(core_area)
|
||
}
|
||
|
||
return config
|
||
|
||
def save_roi_config(self, output_path: str = "config/roi_config.yaml") -> bool:
|
||
"""
|
||
保存ROI配置到文件
|
||
|
||
Args:
|
||
output_path: 输出文件路径
|
||
|
||
Returns:
|
||
bool: 保存是否成功
|
||
"""
|
||
if len(self.roi_points) != len(self.roi_names):
|
||
print(f"标定未完成,只标定了{len(self.roi_points)}/{len(self.roi_names)}个ROI")
|
||
return False
|
||
|
||
try:
|
||
config = self.generate_roi_config()
|
||
|
||
# 确保目录存在
|
||
os.makedirs(os.path.dirname(output_path), exist_ok=True)
|
||
|
||
with open(output_path, 'w', encoding='utf-8') as f:
|
||
yaml.dump(config, f, default_flow_style=False,
|
||
allow_unicode=True, indent=2)
|
||
|
||
print(f"ROI配置已保存到: {output_path}")
|
||
return True
|
||
|
||
except Exception as e:
|
||
print(f"保存ROI配置失败: {e}")
|
||
return False
|
||
|
||
def run_calibration(self, image_path: str):
|
||
"""
|
||
运行标定程序
|
||
|
||
Args:
|
||
image_path: 标定图像路径
|
||
"""
|
||
if not self.load_calibration_image(image_path):
|
||
return
|
||
|
||
cv2.namedWindow(self.window_name, cv2.WINDOW_NORMAL)
|
||
cv2.setMouseCallback(self.window_name, self.mouse_callback)
|
||
|
||
print("ROI标定工具启动")
|
||
print(f"需要标定{len(self.roi_names)}个ROI区域")
|
||
print("按顺序点击每个LED灯的中心位置")
|
||
print("ROI命名规则: R1C1, R1C2, ..., R3C6 (行列从1开始)")
|
||
|
||
self.update_display()
|
||
|
||
while True:
|
||
key = cv2.waitKey(1) & 0xFF
|
||
|
||
if key == ord('q') or key == 27: # Q键或ESC退出
|
||
break
|
||
elif key == ord('r'): # R键重置
|
||
self.reset_calibration()
|
||
elif key == ord('s'): # S键保存
|
||
if self.save_roi_config():
|
||
print("标定完成并保存成功!")
|
||
break
|
||
elif key == ord('+') or key == ord('='): # +键增大ROI
|
||
self.adjust_roi_size(5)
|
||
elif key == ord('-'): # -键减小ROI
|
||
self.adjust_roi_size(-5)
|
||
|
||
cv2.destroyAllWindows()
|
||
|
||
|
||
def main():
|
||
"""
|
||
主函数
|
||
"""
|
||
import argparse
|
||
|
||
parser = argparse.ArgumentParser(description="ROI标定工具")
|
||
parser.add_argument("image", help="标定图像路径")
|
||
parser.add_argument("-o", "--output", default="config/roi_config.yaml",
|
||
help="输出配置文件路径")
|
||
|
||
args = parser.parse_args()
|
||
|
||
if not os.path.exists(args.image):
|
||
print(f"图像文件不存在: {args.image}")
|
||
return
|
||
|
||
tool = ROICalibrationTool()
|
||
tool.run_calibration(args.image)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main() |