"""实时摄像头推理：支持 YOLO / Faster R-CNN 权重。""" from __future__ import annotations import threading import time import tkinter as tk from pathlib import Path from tkinter import filedialog, messagebox, ttk import cv2 import numpy as np from PIL import Image, ImageTk from detectors import guess_backend, load_detector ROOT = Path(__file__).resolve().parent COLOR_RESOLUTIONS = [(1280, 720), (848, 480), (640, 480), (1920, 1080)] def default_model_path(backend: str = "yolo") -> str: if backend == "faster_rcnn": candidates = [ ROOT / "runs" / "faster_rcnn" / "gongjian_frcnn" / "weights" / "best.pth", ROOT / "runs" / "faster_rcnn" / "gongjian_frcnn" / "weights" / "last.pth", ] for path in candidates: if path.exists(): return str(path) return str(ROOT / "runs" / "faster_rcnn" / "gongjian_frcnn" / "weights" / "best.pth") candidates = [ ROOT / "ultralytics" / "runs" / "detect" / "runs" / "yolov8_gongjian" / "weights" / "best.pt", ROOT / "ultralytics" / "runs" / "detect" / "runs" / "yolov8_exp" / "weights" / "best.pt", ROOT / "ultralytics" / "runs" / "detect" / "runs" / "yolov8_exp" / "weights" / "last.pt", ROOT / "runs" / "detect" / "yolov8_exp" / "weights" / "best.pt", ROOT / "yolov8n.pt", ] for path in candidates: if path.exists(): return str(path) return str(ROOT / "yolov8n.pt") def list_realsense() -> list[tuple[str, str]]: try: import pyrealsense2 as rs except Exception: return [] items = [] for dev in rs.context().query_devices(): name = dev.get_info(rs.camera_info.name) serial = dev.get_info(rs.camera_info.serial_number) items.append((f"{name} [{serial}]", serial)) return items class RealSenseRGB: """与 capture.py 一致：同时开左右红外 + RGB，界面只用 RGB。""" def __init__(self, width: int, height: int, fps: int = 30, serial: str | None = None) -> None: import pyrealsense2 as rs self.rs = rs self.pipeline = None self.size = (width, height) last_error: Exception | None = None attempts = [(width, height, fps)] + [(w, h, fps) for w, h in COLOR_RESOLUTIONS if (w, h) != (width, height)] # 推理只要 RGB：先试仅彩色（快、省带宽）；失败再试红外+RGB（与 capture 相同） modes = ("rgb_only", "rgb_ir") for mode in modes: for w, h, f in attempts: pipeline = rs.pipeline() config = rs.config() if serial: config.enable_device(serial) try: if mode == "rgb_ir": config.enable_stream(rs.stream.infrared, 1, w, h, rs.format.y8, f) config.enable_stream(rs.stream.infrared, 2, w, h, rs.format.y8, f) cw, ch = (w, h) if (w, h) in COLOR_RESOLUTIONS else (1280, 720) config.enable_stream(rs.stream.color, cw, ch, rs.format.bgr8, f) profile = pipeline.start(config) try: depth_sensor = profile.get_device().first_depth_sensor() if depth_sensor.supports(rs.option.emitter_enabled): depth_sensor.set_option(rs.option.emitter_enabled, 0.0) except Exception: pass for _ in range(5): pipeline.wait_for_frames(1500) self.pipeline = pipeline self.size = (cw, ch) return except Exception as exc: last_error = exc try: pipeline.stop() except Exception: pass raise RuntimeError( f"无法打开 RealSense RGB（请先关掉 capture.py / RealSense Viewer）: {last_error}" ) def read(self) -> tuple[bool, np.ndarray | None]: if self.pipeline is None: return False, None try: frames = self.pipeline.wait_for_frames(1000) except Exception: return False, None color = frames.get_color_frame() if not color: return False, None return True, np.asanyarray(color.get_data()) def release(self) -> None: if self.pipeline is None: return try: self.pipeline.stop() except Exception: pass self.pipeline = None class Webcam: def __init__(self, index: int, width: int, height: int, fps: int = 30) -> None: cap = None for backend in (cv2.CAP_DSHOW, cv2.CAP_MSMF, cv2.CAP_ANY): trial = cv2.VideoCapture(index, backend) if trial.isOpened(): cap = trial break trial.release() if cap is None or not cap.isOpened(): raise RuntimeError(f"无法打开摄像头 {index}。本机可试编号 1（0 常被 RealSense UVC 占住）") cap.set(cv2.CAP_PROP_FRAME_WIDTH, width) cap.set(cv2.CAP_PROP_FRAME_HEIGHT, height) cap.set(cv2.CAP_PROP_FPS, fps) ok, frame = False, None for _ in range(15): ok, frame = cap.read() if ok and frame is not None: break time.sleep(0.05) if not ok or frame is None: cap.release() raise RuntimeError(f"摄像头 {index} 已打开但读不到画面，请换编号或关掉占用程序") self.cap = cap def read(self) -> tuple[bool, np.ndarray | None]: ok, frame = self.cap.read() return ok, frame if ok else None def release(self) -> None: self.cap.release() def open_source(source: str, width: int, height: int, serial: str | None, cam_index: int): if source == "realsense": return RealSenseRGB(width, height, 30, serial) return Webcam(cam_index, width, height, 30) class InferApp(tk.Tk): def __init__(self) -> None: super().__init__() self.title("实时检测推理（YOLO / Faster R-CNN）") self.geometry("1180x820") self.minsize(980, 700) self.model = None self.cam = None self.running = False self.opening = False self.infer_enabled = False self.photo: ImageTk.PhotoImage | None = None self.latest: np.ndarray | None = None self.fps = 0.0 self.lock = threading.Lock() self.cam_lock = threading.Lock() self.fail_reads = 0 self.rs_devices = list_realsense() self.var_backend = tk.StringVar(value="yolo") self.var_model = tk.StringVar(value=default_model_path("yolo")) self.var_num_classes = tk.IntVar(value=2) self.var_source = tk.StringVar(value="realsense" if self.rs_devices else "webcam") self.var_rs = tk.StringVar(value=self.rs_devices[0][0] if self.rs_devices else "") # 本机普通摄像头能出图的多半是 1 self.var_camera = tk.IntVar(value=1) self.var_width = tk.IntVar(value=1280) self.var_height = tk.IntVar(value=720) self.var_conf = tk.DoubleVar(value=0.25) self.var_iou = tk.DoubleVar(value=0.7) self.var_imgsz = tk.IntVar(value=640) self.var_device = tk.StringVar(value="0") self.var_status = tk.StringVar(value="就绪：点「打开摄像头」先出画面，再点「开始推理」") self._build() self.protocol("WM_DELETE_WINDOW", self.on_close) self.after(33, self.refresh_view) def _build(self) -> None: self.columnconfigure(0, weight=1) self.rowconfigure(0, weight=1) view = ttk.Frame(self, padding=8) view.grid(row=0, column=0, sticky="nsew") view.columnconfigure(0, weight=1) view.rowconfigure(0, weight=1) self.canvas = tk.Label(view, bg="#111", fg="#eee", text="摄像头画面", width=80, height=20) self.canvas.grid(row=0, column=0, sticky="nsew") ttk.Label(view, textvariable=self.var_status).grid(row=1, column=0, sticky="w", pady=(8, 0)) panel = ttk.LabelFrame(self, text="推理参数", padding=10) panel.grid(row=1, column=0, sticky="ew", padx=8, pady=(0, 8)) for c in range(8): panel.columnconfigure(c, weight=1) ttk.Label(panel, text="算法").grid(row=0, column=0, sticky="w") ttk.Combobox( panel, textvariable=self.var_backend, values=["yolo", "faster_rcnn", "auto"], width=12, state="readonly", ).grid(row=0, column=1, sticky="w") ttk.Label(panel, text="RCNN类数(+背景)").grid(row=0, column=2, sticky="w") ttk.Entry(panel, textvariable=self.var_num_classes, width=6).grid(row=0, column=3, sticky="w") self.var_backend.trace_add("write", lambda *_: self.on_backend_change()) ttk.Label(panel, text="权重文件").grid(row=1, column=0, sticky="w", pady=(8, 0)) ttk.Entry(panel, textvariable=self.var_model).grid(row=1, column=1, columnspan=5, sticky="ew", padx=4, pady=(8, 0)) ttk.Button(panel, text="浏览…", command=self.browse_model).grid(row=1, column=6, sticky="ew", pady=(8, 0)) ttk.Button(panel, text="自动查找", command=self.autofill_model).grid(row=1, column=7, sticky="ew", padx=(4, 0), pady=(8, 0)) ttk.Label(panel, text="视频源").grid(row=2, column=0, sticky="w", pady=(8, 0)) src = ttk.Frame(panel) src.grid(row=2, column=1, columnspan=3, sticky="w", pady=(8, 0)) ttk.Radiobutton(src, text="RealSense RGB", variable=self.var_source, value="realsense").pack(side="left") ttk.Radiobutton(src, text="普通摄像头", variable=self.var_source, value="webcam").pack(side="left", padx=(10, 0)) ttk.Label(panel, text="RealSense").grid(row=3, column=0, sticky="w", pady=(8, 0)) self.cmb_rs = ttk.Combobox( panel, textvariable=self.var_rs, values=[x[0] for x in self.rs_devices], state="readonly", ) self.cmb_rs.grid(row=3, column=1, columnspan=3, sticky="ew", pady=(8, 0)) ttk.Button(panel, text="刷新设备", command=self.refresh_devices).grid(row=3, column=4, sticky="w", pady=(8, 0)) ttk.Label(panel, text="摄像头编号").grid(row=3, column=5, sticky="w", pady=(8, 0)) ttk.Spinbox(panel, from_=0, to=8, textvariable=self.var_camera, width=6).grid( row=3, column=6, sticky="w", pady=(8, 0) ) ttk.Label(panel, text="分辨率").grid(row=4, column=0, sticky="w", pady=(8, 0)) ttk.Entry(panel, textvariable=self.var_width, width=8).grid(row=4, column=1, sticky="w", pady=(8, 0)) ttk.Label(panel, text="x").grid(row=4, column=2, pady=(8, 0)) ttk.Entry(panel, textvariable=self.var_height, width=8).grid(row=4, column=3, sticky="w", pady=(8, 0)) ttk.Label(panel, text="imgsz").grid(row=4, column=4, sticky="w", pady=(8, 0)) ttk.Entry(panel, textvariable=self.var_imgsz, width=8).grid(row=4, column=5, sticky="w", pady=(8, 0)) ttk.Label(panel, text="设备").grid(row=4, column=6, sticky="w", pady=(8, 0)) ttk.Combobox(panel, textvariable=self.var_device, values=["0", "cpu"], width=8, state="readonly").grid( row=4, column=7, sticky="w", pady=(8, 0) ) ttk.Label(panel, text="置信度").grid(row=5, column=0, sticky="w", pady=(8, 0)) ttk.Scale(panel, from_=0.05, to=0.95, variable=self.var_conf, orient="horizontal").grid( row=5, column=1, columnspan=2, sticky="ew", pady=(8, 0) ) self.lbl_conf = ttk.Label(panel, text="0.25") self.lbl_conf.grid(row=5, column=3, sticky="w", pady=(8, 0)) self.var_conf.trace_add("write", lambda *_: self.lbl_conf.config(text=f"{float(self.var_conf.get()):.2f}")) ttk.Label(panel, text="IoU").grid(row=5, column=4, sticky="w", pady=(8, 0)) ttk.Entry(panel, textvariable=self.var_iou, width=8).grid(row=5, column=5, sticky="w", pady=(8, 0)) btns = ttk.Frame(panel) btns.grid(row=6, column=0, columnspan=8, sticky="ew", pady=(12, 0)) self.btn_open = ttk.Button(btns, text="打开摄像头", command=self.open_camera) self.btn_open.pack(side="left") self.btn_infer = ttk.Button(btns, text="开始推理", command=self.start_infer, state="disabled") self.btn_infer.pack(side="left", padx=8) self.btn_stop = ttk.Button(btns, text="停止", command=self.stop, state="disabled") self.btn_stop.pack(side="left") ttk.Button(btns, text="保存当前帧", command=self.save_frame).pack(side="left", padx=8) def refresh_devices(self) -> None: self.rs_devices = list_realsense() labels = [x[0] for x in self.rs_devices] self.cmb_rs["values"] = labels if labels: self.var_rs.set(labels[0]) self.var_status.set(f"找到 {len(labels)} 台 RealSense") else: self.var_rs.set("") self.var_status.set("未检测到 RealSense，请改用普通摄像头或检查驱动") def current_rs_serial(self) -> str | None: label = self.var_rs.get() for text, serial in self.rs_devices: if text == label: return serial return None def on_backend_change(self) -> None: backend = self.var_backend.get() if backend == "auto": return self.var_model.set(default_model_path(backend if backend != "auto" else "yolo")) def browse_model(self) -> None: initial = self.var_model.get() initial_dir = str(Path(initial).parent) if initial else str(ROOT) path = filedialog.askopenfilename( title="选择检测权重", initialdir=initial_dir, filetypes=[("PyTorch 权重", "*.pt *.pth"), ("所有文件", "*.*")], ) if path: self.var_model.set(path) if self.var_backend.get() == "auto": self.var_status.set(f"auto 将按文件识别为: {guess_backend(path)}") def autofill_model(self) -> None: backend = self.var_backend.get() if backend == "auto": backend = "yolo" path = default_model_path(backend) self.var_model.set(path) self.var_status.set(f"已填入: {path}") def release_cam(self) -> None: with self.cam_lock: if self.cam is not None: try: self.cam.release() except Exception: pass self.cam = None def open_camera(self) -> None: if self.running or self.opening: return self.opening = True self.btn_open.config(state="disabled") self.var_status.set("正在打开摄像头…（请稍候，不要重复点击）") self.update_idletasks() source = self.var_source.get() width = int(self.var_width.get()) height = int(self.var_height.get()) serial = self.current_rs_serial() cam_index = int(self.var_camera.get()) def worker() -> None: cam = None err = "" try: self.release_cam() if source == "realsense": devices = list_realsense() if not devices: raise RuntimeError("没有 RealSense 设备。可改选「普通摄像头」，编号试 1") use_serial = serial or devices[0][1] cam = open_source("realsense", width, height, use_serial, cam_index) else: cam = open_source("webcam", width, height, None, cam_index) ok, frame = cam.read() if not ok or frame is None: raise RuntimeError("摄像头打开后读不到画面") except Exception as exc: if cam is not None: try: cam.release() except Exception: pass cam = None err = str(exc) self.after(0, lambda: self._open_done(cam, err)) threading.Thread(target=worker, daemon=True).start() def _open_done(self, cam, err: str) -> None: self.opening = False if err or cam is None: self.btn_open.config(state="normal") self.btn_infer.config(state="disabled") self.btn_stop.config(state="disabled") self.var_status.set(err or "打开失败") messagebox.showerror("摄像头", err or "打开失败") return with self.cam_lock: self.cam = cam self.running = True self.infer_enabled = False self.fail_reads = 0 self.btn_open.config(state="disabled") self.btn_infer.config(state="normal") self.btn_stop.config(state="normal") w, h = getattr(cam, "size", (self.var_width.get(), self.var_height.get())) self.var_status.set(f"摄像头已打开 {w}x{h}（预览中）。确认画面后点「开始推理」") threading.Thread(target=self.loop_capture, daemon=True).start() def start_infer(self) -> None: if not self.running or self.cam is None: messagebox.showwarning("推理", "请先打开摄像头") return model_path = Path(self.var_model.get().strip()) if not model_path.is_absolute(): model_path = ROOT / model_path if not model_path.exists(): messagebox.showerror("权重", f"找不到权重文件:\n{model_path}") return self.btn_infer.config(state="disabled") self.var_status.set(f"正在加载模型: {model_path.name} …") # 取一帧热身，避免和采集线程同时 wait_for_frames warmup = None with self.lock: if self.latest is not None: warmup = self.latest.copy() def load() -> None: try: backend = self.var_backend.get() model = load_detector( backend=backend, weights=model_path, device=self.var_device.get(), num_classes=int(self.var_num_classes.get()), conf=float(self.var_conf.get()), ) if warmup is not None: model.plot( warmup, conf=float(self.var_conf.get()), iou=float(self.var_iou.get()), imgsz=int(self.var_imgsz.get()), ) except Exception as exc: self.after(0, lambda e=str(exc): self._infer_load_failed(e)) return self.model = model self.infer_enabled = True used = getattr(model, "backend", backend) self.after(0, lambda: self.var_status.set(f"推理中 backend={used} 模型={model_path.name}")) threading.Thread(target=load, daemon=True).start() def _infer_load_failed(self, msg: str) -> None: self.infer_enabled = False self.model = None self.btn_infer.config(state="normal") self.var_status.set(msg) messagebox.showerror("模型加载失败", msg) def stop(self) -> None: self.running = False self.infer_enabled = False self.var_status.set("正在停止…") def loop_capture(self) -> None: t0, n = time.perf_counter(), 0 while self.running: with self.cam_lock: cam = self.cam if cam is None: break ok, frame = cam.read() if not ok or frame is None: self.fail_reads += 1 if self.fail_reads >= 30: self.after(0, lambda: self.var_status.set("连续读帧失败，请停止后重开摄像头")) self.fail_reads = 0 time.sleep(0.02) continue self.fail_reads = 0 plotted = frame det = 0 if self.infer_enabled and self.model is not None: try: dets = self.model.detect( frame, conf=float(self.var_conf.get()), iou=float(self.var_iou.get()), imgsz=int(self.var_imgsz.get()), ) plotted = frame.copy() for d in dets: x1, y1, x2, y2 = d.xyxy cv2.rectangle(plotted, (x1, y1), (x2, y2), (0, 255, 0), 2) cv2.putText( plotted, f"{d.conf:.2f}", (x1, max(20, y1 - 6)), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2, ) det = len(dets) except Exception as exc: self.after(0, lambda e=str(exc): self.var_status.set(f"推理错误: {e}")) time.sleep(0.05) continue n += 1 if n >= 8: now = time.perf_counter() self.fps = n / max(1e-6, now - t0) t0, n = now, 0 show = plotted.copy() mode = "DET" if self.infer_enabled else "PREVIEW" cv2.putText( show, f"{mode} FPS {self.fps:.1f} det={det} conf>={float(self.var_conf.get()):.2f}", (16, 36), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0) if self.infer_enabled else (0, 200, 255), 2, ) with self.lock: self.latest = show self.release_cam() self.model = None self.infer_enabled = False self.after(0, self._stopped_ui) def _stopped_ui(self) -> None: self.btn_open.config(state="normal") self.btn_infer.config(state="disabled") self.btn_stop.config(state="disabled") self.var_status.set("已停止") def refresh_view(self) -> None: frame = None with self.lock: if self.latest is not None: frame = self.latest.copy() if frame is not None: try: rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) h, w = rgb.shape[:2] max_w = max(int(self.canvas.winfo_width()), 640) max_h = max(int(self.canvas.winfo_height()), 360) scale = min(max_w / w, max_h / h, 1.0) size = (max(1, int(w * scale)), max(1, int(h * scale))) img = Image.fromarray(rgb).resize(size, Image.Resampling.BILINEAR) self.photo = ImageTk.PhotoImage(img) self.canvas.config(image=self.photo, text="") except Exception as exc: self.var_status.set(f"画面刷新失败: {exc}") self.after(33, self.refresh_view) def save_frame(self) -> None: with self.lock: frame = None if self.latest is None else self.latest.copy() if frame is None: messagebox.showwarning("保存", "当前没有可保存的画面") return out_dir = ROOT / "runs" / "predict_realtime" out_dir.mkdir(parents=True, exist_ok=True) out = out_dir / f"shot_{int(time.time())}.jpg" cv2.imwrite(str(out), frame) self.var_status.set(f"已保存: {out}") def on_close(self) -> None: self.running = False self.infer_enabled = False self.release_cam() self.destroy() if __name__ == "__main__": InferApp().mainloop()