398 lines
15 KiB
Python
398 lines
15 KiB
Python
import cv2
|
||
from utils.get_image import GetImage
|
||
from ultralytics import YOLO
|
||
from config import config_manager
|
||
from utils.logger import logger
|
||
import os
|
||
import numpy as np
|
||
|
||
# 检查模型文件是否存在
|
||
model_path = r"best0.pt"
|
||
if not os.path.exists(model_path):
|
||
print(f"❌ 模型文件不存在: {model_path}")
|
||
exit(1)
|
||
|
||
# 加载YOLO模型
|
||
try:
|
||
model = YOLO(model_path).to('cuda')
|
||
print(f"✅ 模型加载成功: {model_path}")
|
||
except Exception as e:
|
||
print(f"❌ 模型加载失败: {e}")
|
||
exit(1)
|
||
|
||
def enhance_sharpness(image, strength=1.5):
|
||
"""
|
||
增强图像锐度
|
||
:param image: 输入图像（BGR格式）
|
||
:param strength: 锐化强度（1.0-3.0，默认1.5）
|
||
:return: 锐化后的图像
|
||
"""
|
||
# 创建锐化核
|
||
kernel = np.array([[-1, -1, -1],
|
||
[-1, 9*strength, -1],
|
||
[-1, -1, -1]]) / (9*strength - 8)
|
||
sharpened = cv2.filter2D(image, -1, kernel)
|
||
return sharpened
|
||
|
||
|
||
def enhance_contrast(image, alpha=1.2, beta=10):
|
||
"""
|
||
增强对比度和亮度
|
||
:param image: 输入图像
|
||
:param alpha: 对比度控制（1.0-3.0，默认1.2）
|
||
:param beta: 亮度控制（-100到100，默认10）
|
||
:return: 增强后的图像
|
||
"""
|
||
return cv2.convertScaleAbs(image, alpha=alpha, beta=beta)
|
||
|
||
|
||
def denoise_image(image, method='bilateral'):
|
||
"""
|
||
去噪处理
|
||
:param image: 输入图像
|
||
:param method: 去噪方法 ('bilateral', 'gaussian', 'fastNlMeans')
|
||
:return: 去噪后的图像
|
||
"""
|
||
if method == 'bilateral':
|
||
# 双边滤波，保留边缘的同时去噪
|
||
return cv2.bilateralFilter(image, 9, 75, 75)
|
||
elif method == 'gaussian':
|
||
# 高斯模糊去噪
|
||
return cv2.GaussianBlur(image, (5, 5), 0)
|
||
elif method == 'fastNlMeans':
|
||
# 非局部均值去噪（效果最好但较慢）
|
||
return cv2.fastNlMeansDenoisingColored(image, None, 10, 10, 7, 21)
|
||
return image
|
||
|
||
|
||
def apply_enhancements(image, sharpness=True, contrast=True, denoise=True,
|
||
sharp_strength=1.5, contrast_alpha=1.2, contrast_beta=10,
|
||
denoise_method='bilateral'):
|
||
"""
|
||
应用所有图像增强
|
||
:param image: 输入图像（BGR格式）
|
||
:param sharpness: 是否锐化
|
||
:param contrast: 是否增强对比度
|
||
:param denoise: 是否去噪
|
||
:param sharp_strength: 锐化强度
|
||
:param contrast_alpha: 对比度系数
|
||
:param contrast_beta: 亮度调整
|
||
:param denoise_method: 去噪方法
|
||
:return: 增强后的图像
|
||
"""
|
||
enhanced = image.copy()
|
||
|
||
if denoise:
|
||
enhanced = denoise_image(enhanced, denoise_method)
|
||
|
||
if contrast:
|
||
enhanced = enhance_contrast(enhanced, contrast_alpha, contrast_beta)
|
||
|
||
if sharpness:
|
||
enhanced = enhance_sharpness(enhanced, sharp_strength)
|
||
|
||
return enhanced
|
||
|
||
|
||
def set_camera_properties(cap, brightness=None, contrast=None, saturation=None,
|
||
sharpness=None, gain=None, exposure=None):
|
||
"""
|
||
设置采集卡硬件参数
|
||
:param cap: VideoCapture对象
|
||
:param brightness: 亮度 (0-100)
|
||
:param contrast: 对比度 (0-100)
|
||
:param saturation: 饱和度 (0-100)
|
||
:param sharpness: 锐度 (0-100)
|
||
:param gain: 增益 (0-100)
|
||
:param exposure: 曝光 (通常为负值，如-6)
|
||
"""
|
||
props = {
|
||
cv2.CAP_PROP_BRIGHTNESS: brightness,
|
||
cv2.CAP_PROP_CONTRAST: contrast,
|
||
cv2.CAP_PROP_SATURATION: saturation,
|
||
cv2.CAP_PROP_SHARPNESS: sharpness,
|
||
cv2.CAP_PROP_GAIN: gain,
|
||
cv2.CAP_PROP_EXPOSURE: exposure,
|
||
}
|
||
|
||
for prop, value in props.items():
|
||
if value is not None:
|
||
try:
|
||
cap.set(prop, value)
|
||
actual = cap.get(prop)
|
||
logger.info(f" 设置 {prop.name if hasattr(prop, 'name') else prop}: {value} -> 实际: {actual:.2f}")
|
||
except Exception as e:
|
||
logger.warning(f" ⚠️ 设置参数 {prop} 失败: {e}")
|
||
|
||
|
||
def yolo_shibie(im_PIL, detections, model, enhance_enabled=False, enhance_params=None):
|
||
"""
|
||
YOLO识别函数
|
||
:param im_PIL: PIL图像对象
|
||
:param detections: 检测结果字典
|
||
:param model: YOLO模型
|
||
:param enhance_enabled: 是否启用图像增强
|
||
:param enhance_params: 图像增强参数
|
||
:return: 更新后的detections字典，如果用户退出则返回None
|
||
"""
|
||
if im_PIL is None:
|
||
return detections
|
||
|
||
try:
|
||
results = model(im_PIL)
|
||
result = results[0]
|
||
|
||
# ✅ 获取绘制好框的图像（RGB格式）
|
||
frame_with_boxes_rgb = result.plot()
|
||
|
||
# ✅ 转换为BGR格式用于OpenCV显示
|
||
frame_with_boxes_bgr = cv2.cvtColor(frame_with_boxes_rgb, cv2.COLOR_RGB2BGR)
|
||
|
||
# 应用图像增强（如果启用）
|
||
display_frame = frame_with_boxes_bgr.copy()
|
||
if enhance_enabled and enhance_params:
|
||
try:
|
||
display_frame = apply_enhancements(display_frame, **enhance_params)
|
||
except Exception as e:
|
||
print(f"⚠️ 图像增强失败: {e}")
|
||
|
||
# 显示YOLO检测结果
|
||
cv2.imshow("YOLO Real-time Detection", display_frame)
|
||
|
||
# ✅ 提取检测信息
|
||
if result.boxes is not None and len(result.boxes.xyxy) > 0:
|
||
# 用于存储多个候选npc4（如果检测到多个）
|
||
npc4_candidates = []
|
||
|
||
for i in range(len(result.boxes.xyxy)):
|
||
try:
|
||
left = float(result.boxes.xyxy[i][0])
|
||
top = float(result.boxes.xyxy[i][1])
|
||
right = float(result.boxes.xyxy[i][2])
|
||
bottom = float(result.boxes.xyxy[i][3])
|
||
cls_id = int(result.boxes.cls[i])
|
||
label = result.names[cls_id]
|
||
|
||
# 获取置信度（如果可用）
|
||
confidence = float(result.boxes.conf[i]) if hasattr(result.boxes, 'conf') and len(result.boxes.conf) > i else 1.0
|
||
|
||
# npc1-npc4 使用底部位置（与main.py保持一致）
|
||
if label in ['npc1', 'npc2', 'npc3', 'npc4']:
|
||
player_x = int(left + (right - left) / 2)
|
||
player_y = int(bottom) + 30 # 使用底部位置，与main.py保持一致
|
||
position = [player_x, player_y]
|
||
|
||
# 特殊处理npc4：如果检测到多个，收集所有候选
|
||
if label == 'npc4':
|
||
npc4_candidates.append({
|
||
'position': position,
|
||
'confidence': confidence,
|
||
'box': [left, top, right, bottom],
|
||
'area': (right - left) * (bottom - top) # 检测框面积
|
||
})
|
||
else:
|
||
# npc1-npc3直接赋值（如果已经有值，保留置信度更高的）
|
||
if detections[label] is None or (hasattr(result.boxes, 'conf') and
|
||
confidence > 0.5):
|
||
detections[label] = position
|
||
|
||
# 其他目标使用中心点
|
||
elif label in ['center', 'next', 'boss', 'zhaozi']:
|
||
player_x = int(left + (right - left) / 2) + 3
|
||
player_y = int(top + (bottom - top) / 2) + 40
|
||
detections[label] = [player_x, player_y]
|
||
|
||
# 道具和怪物可以多个
|
||
elif label in ['daojv', 'gw']:
|
||
player_x = int(left + (right - left) / 2) + 3
|
||
player_y = int(top + (bottom - top) / 2) + 40
|
||
# 确保列表存在
|
||
if label not in detections:
|
||
detections[label] = []
|
||
detections[label].append([player_x, player_y])
|
||
|
||
except Exception as e:
|
||
print(f"⚠️ 处理检测框时出错: {e}")
|
||
continue
|
||
|
||
# 处理npc4：如果检测到多个，选择最合适的
|
||
if npc4_candidates:
|
||
# 按置信度排序，选择置信度最高的
|
||
npc4_candidates.sort(key=lambda x: x['confidence'], reverse=True)
|
||
|
||
# 选择最佳候选（置信度最高且面积合理）
|
||
best_npc4 = None
|
||
for candidate in npc4_candidates:
|
||
# 置信度阈值：至少0.3（可根据实际情况调整）
|
||
if candidate['confidence'] >= 0.3:
|
||
# 检查检测框面积是否合理（避免过小的误检）
|
||
area = candidate['area']
|
||
if area > 100: # 最小面积阈值
|
||
best_npc4 = candidate
|
||
break
|
||
|
||
if best_npc4:
|
||
detections['npc4'] = best_npc4['position']
|
||
# 可选：输出调试信息
|
||
# print(f"✅ 检测到npc4: 位置={best_npc4['position']}, 置信度={best_npc4['confidence']:.2f}")
|
||
elif len(npc4_candidates) == 1:
|
||
# 如果只有一个候选，即使置信度较低也使用
|
||
detections['npc4'] = npc4_candidates[0]['position']
|
||
|
||
except Exception as e:
|
||
print(f"⚠️ YOLO检测出错: {e}")
|
||
|
||
return detections
|
||
|
||
|
||
def main():
|
||
"""主函数"""
|
||
print("="*60)
|
||
print("YOLO实时检测测试")
|
||
print("="*60)
|
||
|
||
# 从配置加载采集卡设置
|
||
active_group = config_manager.get_active_group()
|
||
|
||
if active_group is None:
|
||
print("⚠️ 没有活动的配置组，使用默认设置")
|
||
print("提示: 可以运行 python gui_config.py 设置配置")
|
||
cam_index = 0
|
||
width = 1920
|
||
height = 1080
|
||
else:
|
||
print(f"📋 使用配置组: {active_group['name']}")
|
||
cam_index = active_group['camera_index']
|
||
width = active_group['camera_width']
|
||
height = active_group['camera_height']
|
||
|
||
print(f" 采集卡索引: {cam_index}")
|
||
print(f" 分辨率: {width}x{height}")
|
||
print()
|
||
|
||
# 初始化采集卡
|
||
print("🔧 正在初始化采集卡...")
|
||
get_image = GetImage(
|
||
cam_index=cam_index,
|
||
width=width,
|
||
height=height
|
||
)
|
||
|
||
if get_image.cap is None:
|
||
print("❌ 采集卡初始化失败")
|
||
print("请检查：")
|
||
print("1. 采集卡是否正确连接")
|
||
print("2. 采集卡索引是否正确")
|
||
print("3. 采集卡驱动是否安装")
|
||
return
|
||
|
||
# 设置采集卡硬件参数以提高清晰度（可选）
|
||
print("\n🔧 设置采集卡参数以提高清晰度...")
|
||
print("提示: 可以根据实际情况调整这些参数")
|
||
set_camera_properties(
|
||
get_image.cap,
|
||
brightness=50, # 亮度 (0-100)
|
||
contrast=50, # 对比度 (0-100)
|
||
saturation=55, # 饱和度 (0-100)
|
||
sharpness=60, # 锐度 (0-100，提高清晰度)
|
||
gain=None, # 增益 (根据实际情况调整)
|
||
exposure=None # 曝光 (根据实际情况调整，通常为负值)
|
||
)
|
||
|
||
print("✅ 采集卡初始化成功")
|
||
print("\n快捷键:")
|
||
print(" 'q' 或 ESC - 退出")
|
||
print(" 'e' - 切换图像增强")
|
||
print(" '1'/'2' - 调整锐化强度 (+/-0.1)")
|
||
print(" '3'/'4' - 调整对比度 (+/-0.1)")
|
||
print()
|
||
|
||
try:
|
||
frame_count = 0
|
||
enhance_enabled = False # 默认关闭图像增强
|
||
|
||
# 图像增强参数
|
||
enhance_params = {
|
||
'sharpness': True,
|
||
'contrast': True,
|
||
'denoise': True,
|
||
'sharp_strength': 1.5,
|
||
'contrast_alpha': 1.2,
|
||
'contrast_beta': 10,
|
||
'denoise_method': 'bilateral'
|
||
}
|
||
|
||
while True:
|
||
# 获取帧
|
||
frame_data = get_image.get_frame()
|
||
|
||
if frame_data is None:
|
||
print("⚠️ 无法获取帧，跳过...")
|
||
continue
|
||
|
||
# frame_data 是 [im_opencv_rgb, im_PIL] 格式
|
||
# im_opencv_rgb 已经是RGB格式（经过BGR2RGB转换）
|
||
im_opencv_rgb, im_PIL = frame_data
|
||
|
||
if im_PIL is None:
|
||
print("⚠️ PIL图像为空，跳过...")
|
||
continue
|
||
|
||
# 初始化检测结果字典
|
||
detections = {
|
||
'center': None, 'next': None,
|
||
'npc1': None, 'npc2': None, 'npc3': None, 'npc4': None,
|
||
'boss': None, 'zhaozi': None,
|
||
'daojv': [], 'gw': []
|
||
}
|
||
|
||
# 执行YOLO检测
|
||
detections = yolo_shibie(im_PIL, detections, model, enhance_enabled, enhance_params)
|
||
|
||
# 检查按键
|
||
key = cv2.waitKey(1) & 0xFF
|
||
if key in [27, ord('q'), ord('Q')]:
|
||
print("\n用户退出")
|
||
break
|
||
elif key == ord('e') or key == ord('E'):
|
||
enhance_enabled = not enhance_enabled
|
||
status = "开启" if enhance_enabled else "关闭"
|
||
print(f"图像增强: {status} (锐化={enhance_params['sharp_strength']:.1f}, "
|
||
f"对比度={enhance_params['contrast_alpha']:.1f})")
|
||
elif key == ord('1'):
|
||
enhance_params['sharp_strength'] = min(3.0, enhance_params['sharp_strength'] + 0.1)
|
||
print(f"锐化强度: {enhance_params['sharp_strength']:.1f}")
|
||
elif key == ord('2'):
|
||
enhance_params['sharp_strength'] = max(0.5, enhance_params['sharp_strength'] - 0.1)
|
||
print(f"锐化强度: {enhance_params['sharp_strength']:.1f}")
|
||
elif key == ord('3'):
|
||
enhance_params['contrast_alpha'] = min(3.0, enhance_params['contrast_alpha'] + 0.1)
|
||
print(f"对比度: {enhance_params['contrast_alpha']:.1f}")
|
||
elif key == ord('4'):
|
||
enhance_params['contrast_alpha'] = max(0.5, enhance_params['contrast_alpha'] - 0.1)
|
||
print(f"对比度: {enhance_params['contrast_alpha']:.1f}")
|
||
|
||
frame_count += 1
|
||
if frame_count % 30 == 0: # 每30帧打印一次
|
||
print(f"📊 已处理 {frame_count} 帧")
|
||
# 打印有检测到的目标
|
||
detected_items = {k: v for k, v in detections.items() if v is not None and v != []}
|
||
if detected_items:
|
||
print(f" 检测到: {detected_items}")
|
||
|
||
except KeyboardInterrupt:
|
||
print("\n\n用户中断测试")
|
||
except Exception as e:
|
||
print(f"\n❌ 测试过程中发生错误: {e}")
|
||
import traceback
|
||
traceback.print_exc()
|
||
finally:
|
||
# 清理资源
|
||
get_image.release()
|
||
cv2.destroyAllWindows()
|
||
print("🔚 测试结束")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|