From c5415156edecd2a4406841670548e06b0d9f24e5 Mon Sep 17 00:00:00 2001 From: ray <1416431931@qq.com> Date: Tue, 4 Nov 2025 11:39:07 +0800 Subject: [PATCH] =?UTF-8?q?=E6=B5=8B=E8=AF=95=E6=96=87=E4=BB=B6=E6=8F=90?= =?UTF-8?q?=E4=BA=A4?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- yolo_test.py | 179 ++++++++++++++++++++++++++++++++++++++++++++++++--- 1 file changed, 170 insertions(+), 9 deletions(-) diff --git a/yolo_test.py b/yolo_test.py index 36aea2e..b58c497 100644 --- a/yolo_test.py +++ b/yolo_test.py @@ -2,7 +2,9 @@ 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" @@ -18,7 +20,113 @@ except Exception as e: print(f"❌ 模型加载失败: {e}") exit(1) -def yolo_shibie(im_PIL, im_opencv_rgb, raw_frame_bgr, detections, model, show_original=True): +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, im_opencv_rgb, raw_frame_bgr, detections, model, show_original=True, + enhance_enabled=False, enhance_params=None): """ YOLO识别函数 :param im_PIL: PIL图像对象 @@ -42,11 +150,19 @@ def yolo_shibie(im_PIL, im_opencv_rgb, raw_frame_bgr, detections, model, show_or # ✅ 转换为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}") + # 显示画面 if show_original and raw_frame_bgr is not None: # 同时显示原始帧和检测结果（并排显示） # 调整原始帧大小以匹配裁剪后的检测结果 - h, w = frame_with_boxes_bgr.shape[:2] + h, w = display_frame.shape[:2] # 裁剪原始帧（与get_frame的处理一致：30:30+720, 0:1280） raw_height, raw_width = raw_frame_bgr.shape[:2] crop_top = 30 @@ -60,11 +176,11 @@ def yolo_shibie(im_PIL, im_opencv_rgb, raw_frame_bgr, detections, model, show_or # 并排显示：原始帧(左) | 检测结果(右) # 原始帧已经是BGR格式，检测结果也是BGR格式，可以直接拼接 - combined = cv2.hconcat([raw_cropped, frame_with_boxes_bgr]) - cv2.imshow("原始帧 (左, 与raw_frame.jpg一致) | YOLO检测结果 (右)", combined) + combined = cv2.hconcat([raw_cropped, display_frame]) + cv2.imshow("Original BGR (Left) | YOLO Detection (Right)", combined) else: # 只显示检测结果 - cv2.imshow("YOLO实时检测", frame_with_boxes_bgr) + cv2.imshow("YOLO Real-time Detection", display_frame) # ✅ 提取检测信息 if result.boxes is not None and len(result.boxes.xyxy) > 0: @@ -139,16 +255,43 @@ def main(): 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("按 'q' 或 ESC 键退出测试") + print("\n快捷键:") + print(" 'q' 或 ESC - 退出") + print(" 'o' - 切换原始帧对比模式") + print(" 'e' - 切换图像增强") + print(" '1'/'2' - 调整锐化强度 (+/-0.1)") + print(" '3'/'4' - 调整对比度 (+/-0.1)") print() try: frame_count = 0 show_original = True # 默认同时显示原始帧和检测结果 + enhance_enabled = False # 默认关闭图像增强 - print("提示: 按 'o' 键切换是否显示原始帧对比") - print() + # 图像增强参数 + enhance_params = { + 'sharpness': True, + 'contrast': True, + 'denoise': True, + 'sharp_strength': 1.5, + 'contrast_alpha': 1.2, + 'contrast_beta': 10, + 'denoise_method': 'bilateral' + } while True: # 获取帧 @@ -180,7 +323,8 @@ def main(): } # 执行YOLO检测 - detections = yolo_shibie(im_PIL, im_opencv_rgb, raw_frame_bgr, detections, model, show_original) + detections = yolo_shibie(im_PIL, im_opencv_rgb, raw_frame_bgr, detections, model, + show_original, enhance_enabled, enhance_params) # 检查按键 key = cv2.waitKey(1) & 0xFF @@ -190,6 +334,23 @@ def main(): elif key == ord('o') or key == ord('O'): show_original = not show_original print(f"切换显示模式: {'原始帧对比' if show_original else '仅检测结果'}") + 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帧打印一次