diff --git a/src/image_object_detection/camera_detector.py b/src/image_object_detection/camera_detector.py index cb0ad65ca4..466e923f15 100644 --- a/src/image_object_detection/camera_detector.py +++ b/src/image_object_detection/camera_detector.py @@ -1,110 +1,73 @@ # camera_detector.py import cv2 import time -import os import traceback class CameraDetector: - def __init__(self, detection_engine, camera_index=0): + def __init__(self, detection_engine, output_interval=1.0): self.engine = detection_engine - self.camera_index = camera_index - self.window_name = "YOLO 实时检测 - 摄像头" + self.output_interval = output_interval + self.last_output_time = 0 + self.frame_count = 0 + self.window_name = "YOLO_Live_Detection" # 👈 英文窗口名 - def detect_camera(self): - """对外统一接口：启动摄像头实时检测""" - self.start_detection() - - def start_detection(self): - """实际执行摄像头检测的逻辑""" - print("正在初始化摄像头...") - cap = self._open_camera() - if cap is None: + def start_detection(self, camera_index=0): + cap = cv2.VideoCapture(camera_index) + if not cap.isOpened(): + print(f"Error: Cannot open camera {camera_index}") return - cv2.namedWindow(self.window_name, cv2.WINDOW_NORMAL) - cv2.resizeWindow(self.window_name, 800, 600) - - print("摄像头已启动。按 'q' 或 ESC 退出，按 's' 保存当前帧。") - - prev_time = time.time() - saved_frame_index = 0 + print("Starting live detection. Press 'q' to quit, 's' to save frame.") + cv2.namedWindow(self.window_name, cv2.WINDOW_NORMAL) # 创建一次即可 try: while True: ret, frame = cap.read() if not ret: - print("警告: 无法读取摄像头帧（可能设备被拔出或占用）") + print("Warning: Failed to read frame from camera") break - annotated_frame, results = self._perform_detection(frame) - if annotated_frame is None: - annotated_frame = frame + current_time = time.time() + annotated_frame, results = self.engine.detect(frame) - curr_time = time.time() - fps = 1.0 / (curr_time - prev_time) if (curr_time - prev_time) > 0 else 0 - prev_time = curr_time - - num_objects = len(results[0].boxes) if results else 0 - info_text = f"FPS: {fps:.1f} | Objects: {num_objects}" - cv2.putText( - annotated_frame, - info_text, - (10, 30), - cv2.FONT_HERSHEY_SIMPLEX, - 0.7, - (0, 255, 0), - 2 - ) + if annotated_frame is None or annotated_frame.size == 0: + print("Warning: Invalid detection result") + continue cv2.imshow(self.window_name, annotated_frame) key = cv2.waitKey(1) & 0xFF - if key == ord('q') or key == 27: - print("用户请求退出。") + if key == ord('q'): break elif key == ord('s'): - saved_frame_index = self._save_frame(annotated_frame, saved_frame_index) - + self.save_frame(annotated_frame) + + self._print_fps_if_needed(current_time) + self.frame_count += 1 + + except KeyboardInterrupt: + print("\nDetection interrupted by user.") except Exception as e: - print(f"检测过程中发生未预期错误: {e}") + print(f"Unexpected error during detection: {e}") traceback.print_exc() finally: cap.release() cv2.destroyAllWindows() - print("摄像头已关闭，窗口已清理.") - - def _open_camera(self): - cap = cv2.VideoCapture(self.camera_index) - if not cap.isOpened(): - print(f"错误: 无法打开摄像头（索引 {self.camera_index}）") - print("请检查摄像头是否连接、是否被其他程序占用。") - return None - - cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1280) - cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 720) - return cap - - def _perform_detection(self, frame): - try: - # 抑制 YOLO 内部输出 - import sys, os - old_stdout = sys.stdout - sys.stdout = open(os.devnull, 'w') - annotated_frame, results = self.engine.detect(frame) - sys.stdout = old_stdout - return annotated_frame, results - except Exception as e: - print(f"单帧检测失败: {e}") - return None, [] - - def _save_frame(self, frame, index): - save_dir = "saved_frames" - os.makedirs(save_dir, exist_ok=True) - filename = os.path.join(save_dir, f"frame_{index:04d}.jpg") + print("Camera released and windows closed.") + + def _print_fps_if_needed(self, current_time): + if current_time - self.last_output_time >= self.output_interval: + fps = self.frame_count / (current_time - self.last_output_time) if self.last_output_time > 0 else 0 + print(f"FPS: {fps:.2f}") + self.last_output_time = current_time + self.frame_count = 0 + + def save_frame(self, frame): + import os + timestamp = int(time.time()) + filename = f"saved_frame_{timestamp}.jpg" success = cv2.imwrite(filename, frame) if success: - print(f"✅ 帧已保存: {filename}") - return index + 1 + print(f"Frame saved as {filename}") else: - print("❌ 保存帧失败！") - return index + print("Failed to save frame") diff --git a/src/image_object_detection/image_detector.py b/src/image_object_detection/image_detector.py index b161b45597..70f31256b2 100644 --- a/src/image_object_detection/image_detector.py +++ b/src/image_object_detection/image_detector.py @@ -2,111 +2,65 @@ import cv2 import numpy as np import traceback +import os class ImageDetector: def __init__(self, detection_engine): self.engine = detection_engine def detect_static_image(self, image_path): - """检测静态图像""" - print(f"正在加载图像: {image_path}") + print(f"Loading image: {image_path}") - if not self._check_image_exists(image_path): + if not os.path.exists(image_path): + print(f"Error: Image file not found - {image_path}") return - frame = self._load_image(image_path) + frame = cv2.imread(image_path) if frame is None or frame.size == 0: - print("错误: 无法读取图像文件或图像为空") + print("Error: Failed to load image") return - print("正在检测图像...") - results, annotated_frame = self._perform_detection(frame) + print("Running detection...") + annotated_frame, results = self.engine.detect(frame) - if annotated_frame is None: - print("错误: 检测未返回有效图像") + if annotated_frame is None or annotated_frame.size == 0: + print("Error: Invalid annotated frame") return - # 显示检测结果文本 self._display_results(results) - - # 显示图像窗口（带固定初始大小） - self._show_image(annotated_frame, 'YOLO 检测结果 - 静态图像') - - def _check_image_exists(self, image_path): - import os - if not os.path.exists(image_path): - print(f"错误: 图像文件不存在 - {image_path}") - return False - return True - - def _load_image(self, image_path): - try: - with open(image_path, "rb") as f: - bytes_data = bytearray(f.read()) - nparr = np.frombuffer(bytes_data, np.uint8) - frame = cv2.imdecode(nparr, cv2.IMREAD_COLOR) - return frame - except Exception as e: - print(f"加载图像时出错: {e}") - return None - - def _perform_detection(self, frame): - try: - annotated_frame, results = self.engine.detect(frame) - return results, annotated_frame - except Exception as e: - print(f"检测过程中发生错误: {e}") - traceback.print_exc() - return [], None + self._show_image(annotated_frame, "YOLO_Static_Detection") def _display_results(self, results): if not results: - print("未检测到任何对象（results 为空）") + print("No objects detected.") return result = results[0] boxes = result.boxes if len(boxes) == 0: - print("未检测到任何对象") + print("No objects detected.") return - print("检测完成，正在显示结果...") - print(f"检测到 {len(boxes)} 个对象:") - + print(f"Detected {len(boxes)} object(s):") names_list = self.engine.model.names for i, box in enumerate(boxes): try: cls_index = int(box.cls.item()) confidence = box.conf.item() - - if 0 <= cls_index < len(names_list): - cls_name = names_list[cls_index] - else: - cls_name = f"unknown_class_{cls_index}" - print(f" ⚠️ 警告：类别索引 {cls_index} 超出范围（共 {len(names_list)} 类）") - - print(f" {i+1}. {cls_name} (置信度: {confidence:.2f})") + cls_name = names_list[cls_index] if 0 <= cls_index < len(names_list) else f"unknown_{cls_index}" + print(f" {i+1}. {cls_name} (confidence: {confidence:.2f})") except Exception as e: - print(f" ⚠️ 解析第 {i+1} 个检测框时出错: {e}") + print(f" Warning: Failed to parse box {i+1}: {e}") - def _show_image(self, annotated_frame, window_name): - """显示图像，并设置窗口为可调整的小尺寸""" - if annotated_frame is None or annotated_frame.size == 0: - print("错误: 标注帧无效，无法显示") - return - - print("\n图像已显示。按任意键关闭窗口...") + def _show_image(self, annotated_frame, window_name="YOLO_Static_Detection"): try: - # 创建可调整大小的窗口 + cv2.destroyAllWindows() cv2.namedWindow(window_name, cv2.WINDOW_NORMAL) - # 设置初始窗口大小为 800x600（你可以按需修改） cv2.resizeWindow(window_name, 800, 600) - # 显示图像（OpenCV 会自动适应窗口） cv2.imshow(window_name, annotated_frame) cv2.waitKey(0) cv2.destroyAllWindows() - print("窗口已关闭。") except Exception as e: - print(f"显示图像时发生错误: {e}") + print(f"Failed to display image: {e}") traceback.print_exc() diff --git a/src/image_object_detection/main.py b/src/image_object_detection/main.py index 6c1bda8efd..dc6d559477 100644 --- a/src/image_object_detection/main.py +++ b/src/image_object_detection/main.py @@ -1,4 +1,7 @@ # main.py +import matplotlib +matplotlib.use('Agg') # 必须在最前，禁用 GUI 后端 + from ui_handler import UIHandler from config import Config @@ -8,4 +11,4 @@ def main(): ui_handler.run() if __name__ == "__main__": - main() \ No newline at end of file + main() diff --git a/src/image_object_detection/ui_handler.py b/src/image_object_detection/ui_handler.py index 53e4036e0d..1e3f1073c3 100644 --- a/src/image_object_detection/ui_handler.py +++ b/src/image_object_detection/ui_handler.py @@ -1,75 +1,49 @@ # ui_handler.py +import os import cv2 +import traceback from detection_engine import DetectionEngine from image_detector import ImageDetector from camera_detector import CameraDetector -import sys -import os class UIHandler: def __init__(self, config): self.config = config - self.detection_engine = DetectionEngine() - self.image_detector = ImageDetector(self.detection_engine) - self.camera_detector = CameraDetector(self.detection_engine) + self.engine = DetectionEngine( + model_path=config.model_path, + conf_threshold=config.confidence_threshold + ) + self.stop_flag = False - def print_debug(self, message): - """调试输出函数""" - print(f"[DEBUG] {message}") - sys.stdout.flush() # 强制刷新输出缓冲区 - - def display_menu(self): - """显示主菜单""" - print("=" * 60) - print(" 欢迎使用 YOLO 对象检测系统") - print("=" * 60) + def run(self): + print("=== YOLO Detection System ===") + print("1. Static Image Detection") + print("2. Live Camera Detection") + print("3. Exit") - print(f"\n测试图像路径: {self.config.test_image_path}") - print(f"图像文件存在: {os.path.exists(self.config.test_image_path)}") + choice = input("Please select an option (1-3): ").strip() - print("\n请选择检测模式:") - print("1. 检测静态图像") - print("2. 检测摄像头实时画面") - print("3. 退出程序") - - def handle_choice(self, choice): - """处理用户选择""" - if choice == '1': - print("\n您选择了: 检测静态图像") - self.image_detector.detect_static_image(self.config.test_image_path) - return True - elif choice == '2': - print("\n您选择了: 检测摄像头实时画面") - print("注意：要退出摄像头模式，请按 Ctrl+C ！") - self.camera_detector.detect_camera() - return True - elif choice == '3': - print("\n程序已退出。") - return False + if choice == "1": + self._run_static_detection() + elif choice == "2": + self._run_camera_detection() + elif choice == "3": + print("Exiting program.") else: - print("\n无效选择，请输入 1、2 或 3") - return True + print("Invalid option. Please enter 1, 2, or 3.") - def run(self): - """运行主程序""" - self.print_debug("程序开始执行") - self.display_menu() + def _run_static_detection(self): + image_path = self.config.test_image_path + if not os.path.exists(image_path): + print(f"Image not found: {image_path}") + return - running = True - while running: - try: - choice = input("\n请输入您的选择 (1/2/3): ").strip() - self.print_debug(f"用户输入: {choice}") - running = self.handle_choice(choice) - - except KeyboardInterrupt: - print("\n\n程序被用户中断。") - break - except EOFError: - print("\n输入结束，程序退出。") - break - except Exception as e: - print(f"\n发生错误: {e}") - break - - self.print_debug("程序结束") \ No newline at end of file + detector = ImageDetector(self.engine) + detector.detect_static_image(image_path) + + def _run_camera_detection(self): + detector = CameraDetector( + detection_engine=self.engine, + output_interval=self.config.output_interval + ) + detector.start_detection(camera_index=self.config.camera_index)