测试文件提交
This commit is contained in:
44
yolo_test.py
44
yolo_test.py
@@ -125,16 +125,14 @@ def set_camera_properties(cap, brightness=None, contrast=None, saturation=None,
|
|||||||
logger.warning(f" ⚠️ 设置参数 {prop} 失败: {e}")
|
logger.warning(f" ⚠️ 设置参数 {prop} 失败: {e}")
|
||||||
|
|
||||||
|
|
||||||
def yolo_shibie(im_PIL, im_opencv_rgb, raw_frame_bgr, detections, model, show_original=True,
|
def yolo_shibie(im_PIL, detections, model, enhance_enabled=False, enhance_params=None):
|
||||||
enhance_enabled=False, enhance_params=None):
|
|
||||||
"""
|
"""
|
||||||
YOLO识别函数
|
YOLO识别函数
|
||||||
:param im_PIL: PIL图像对象
|
:param im_PIL: PIL图像对象
|
||||||
:param im_opencv_rgb: RGB格式的OpenCV图像(裁剪后)
|
|
||||||
:param raw_frame_bgr: 原始BGR格式的OpenCV图像(未裁剪,与raw_frame.jpg一致)
|
|
||||||
:param detections: 检测结果字典
|
:param detections: 检测结果字典
|
||||||
:param model: YOLO模型
|
:param model: YOLO模型
|
||||||
:param show_original: 是否同时显示原始帧
|
:param enhance_enabled: 是否启用图像增强
|
||||||
|
:param enhance_params: 图像增强参数
|
||||||
:return: 更新后的detections字典,如果用户退出则返回None
|
:return: 更新后的detections字典,如果用户退出则返回None
|
||||||
"""
|
"""
|
||||||
if im_PIL is None:
|
if im_PIL is None:
|
||||||
@@ -158,28 +156,7 @@ def yolo_shibie(im_PIL, im_opencv_rgb, raw_frame_bgr, detections, model, show_or
|
|||||||
except Exception as e:
|
except Exception as e:
|
||||||
print(f"⚠️ 图像增强失败: {e}")
|
print(f"⚠️ 图像增强失败: {e}")
|
||||||
|
|
||||||
# 显示画面
|
# 显示YOLO检测结果
|
||||||
if show_original and raw_frame_bgr is not None:
|
|
||||||
# 同时显示原始帧和检测结果(并排显示)
|
|
||||||
# 调整原始帧大小以匹配裁剪后的检测结果
|
|
||||||
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
|
|
||||||
crop_bottom = min(crop_top + h, raw_height)
|
|
||||||
crop_right = min(w, raw_width)
|
|
||||||
raw_cropped = raw_frame_bgr[crop_top:crop_bottom, 0:crop_right]
|
|
||||||
|
|
||||||
# 如果尺寸不匹配,调整原始帧大小
|
|
||||||
if raw_cropped.shape[:2] != (h, w):
|
|
||||||
raw_cropped = cv2.resize(raw_cropped, (w, h))
|
|
||||||
|
|
||||||
# 并排显示:原始帧(左) | 检测结果(右)
|
|
||||||
# 原始帧已经是BGR格式,检测结果也是BGR格式,可以直接拼接
|
|
||||||
combined = cv2.hconcat([raw_cropped, display_frame])
|
|
||||||
cv2.imshow("Original BGR (Left) | YOLO Detection (Right)", combined)
|
|
||||||
else:
|
|
||||||
# 只显示检测结果
|
|
||||||
cv2.imshow("YOLO Real-time Detection", display_frame)
|
cv2.imshow("YOLO Real-time Detection", display_frame)
|
||||||
|
|
||||||
# ✅ 提取检测信息
|
# ✅ 提取检测信息
|
||||||
@@ -271,7 +248,6 @@ def main():
|
|||||||
print("✅ 采集卡初始化成功")
|
print("✅ 采集卡初始化成功")
|
||||||
print("\n快捷键:")
|
print("\n快捷键:")
|
||||||
print(" 'q' 或 ESC - 退出")
|
print(" 'q' 或 ESC - 退出")
|
||||||
print(" 'o' - 切换原始帧对比模式")
|
|
||||||
print(" 'e' - 切换图像增强")
|
print(" 'e' - 切换图像增强")
|
||||||
print(" '1'/'2' - 调整锐化强度 (+/-0.1)")
|
print(" '1'/'2' - 调整锐化强度 (+/-0.1)")
|
||||||
print(" '3'/'4' - 调整对比度 (+/-0.1)")
|
print(" '3'/'4' - 调整对比度 (+/-0.1)")
|
||||||
@@ -279,7 +255,6 @@ def main():
|
|||||||
|
|
||||||
try:
|
try:
|
||||||
frame_count = 0
|
frame_count = 0
|
||||||
show_original = True # 默认同时显示原始帧和检测结果
|
|
||||||
enhance_enabled = False # 默认关闭图像增强
|
enhance_enabled = False # 默认关闭图像增强
|
||||||
|
|
||||||
# 图像增强参数
|
# 图像增强参数
|
||||||
@@ -309,11 +284,6 @@ def main():
|
|||||||
print("⚠️ PIL图像为空,跳过...")
|
print("⚠️ PIL图像为空,跳过...")
|
||||||
continue
|
continue
|
||||||
|
|
||||||
# 获取原始BGR帧(与test_capture_card.py保存的raw_frame.jpg一致)
|
|
||||||
raw_frame_bgr = None
|
|
||||||
if get_image.cap is not None and get_image.frame is not None:
|
|
||||||
raw_frame_bgr = get_image.frame.copy() # 原始BGR格式,未裁剪
|
|
||||||
|
|
||||||
# 初始化检测结果字典
|
# 初始化检测结果字典
|
||||||
detections = {
|
detections = {
|
||||||
'center': None, 'next': None,
|
'center': None, 'next': None,
|
||||||
@@ -323,17 +293,13 @@ def main():
|
|||||||
}
|
}
|
||||||
|
|
||||||
# 执行YOLO检测
|
# 执行YOLO检测
|
||||||
detections = yolo_shibie(im_PIL, im_opencv_rgb, raw_frame_bgr, detections, model,
|
detections = yolo_shibie(im_PIL, detections, model, enhance_enabled, enhance_params)
|
||||||
show_original, enhance_enabled, enhance_params)
|
|
||||||
|
|
||||||
# 检查按键
|
# 检查按键
|
||||||
key = cv2.waitKey(1) & 0xFF
|
key = cv2.waitKey(1) & 0xFF
|
||||||
if key in [27, ord('q'), ord('Q')]:
|
if key in [27, ord('q'), ord('Q')]:
|
||||||
print("\n用户退出")
|
print("\n用户退出")
|
||||||
break
|
break
|
||||||
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'):
|
elif key == ord('e') or key == ord('E'):
|
||||||
enhance_enabled = not enhance_enabled
|
enhance_enabled = not enhance_enabled
|
||||||
status = "开启" if enhance_enabled else "关闭"
|
status = "开启" if enhance_enabled else "关闭"
|
||||||
|
|||||||
Reference in New Issue
Block a user