测试文件提交

2025-11-04 11:39:07 +08:00
parent b10333a308
commit c5415156ed
1 changed files with 170 additions and 9 deletions
--- 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帧打印一次