"""工件螺纹孔二阶段视频推理：YOLO / Faster R-CNN 定位 + HoughCircles。""" 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 from hole_detect import HoleStabilizer, detect_holes, load_config from infer_realtime import ( list_realsense, open_source, ) ROOT = Path(__file__).resolve().parent 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_gongjian" / "weights" / "last.pt", ROOT / "ultralytics" / "runs" / "detect" / "runs" / "yolov8_exp" / "weights" / "best.pt", ROOT / "yolov8n.pt", ] for path in candidates: if path.exists(): return str(path) return str(ROOT / "yolov8n.pt") def default_config_path() -> str: path = ROOT / "hole_detect_config.yaml" return str(path) class HoleInferApp(tk.Tk): def __init__(self) -> None: super().__init__() self.title("工件螺纹孔二阶段检测") self.geometry("1180x860") self.minsize(980, 720) self.model = None self.cfg: dict = {} 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.last_result: dict | None = None self.stabilizer = HoleStabilizer() 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_config = tk.StringVar(value=default_config_path()) 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 "") self.var_camera = tk.IntVar(value=1) self.var_width = tk.IntVar(value=1280) self.var_height = tk.IntVar(value=720) 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) self.reload_config(silent=True) 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)) ttk.Entry(panel, textvariable=self.var_config).grid(row=2, column=1, columnspan=5, sticky="ew", padx=4, pady=(8, 0)) ttk.Button(panel, text="浏览…", command=self.browse_config).grid(row=2, column=6, sticky="ew", pady=(8, 0)) ttk.Button(panel, text="重新加载", command=self.reload_config).grid(row=2, column=7, sticky="ew", padx=(4, 0), pady=(8, 0)) ttk.Label(panel, text="视频源").grid(row=3, column=0, sticky="w", pady=(8, 0)) src = ttk.Frame(panel) src.grid(row=3, 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="设备").grid(row=3, column=4, sticky="w", pady=(8, 0)) ttk.Combobox(panel, textvariable=self.var_device, values=["0", "cpu"], width=8, state="readonly").grid( row=3, column=5, sticky="w", pady=(8, 0) ) ttk.Label(panel, text="RealSense").grid(row=4, 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=4, column=1, columnspan=3, sticky="ew", pady=(8, 0)) ttk.Button(panel, text="刷新设备", command=self.refresh_devices).grid(row=4, column=4, sticky="w", pady=(8, 0)) ttk.Label(panel, text="摄像头编号").grid(row=4, column=5, sticky="w", pady=(8, 0)) ttk.Spinbox(panel, from_=0, to=8, textvariable=self.var_camera, width=6).grid( row=4, column=6, sticky="w", pady=(8, 0) ) ttk.Label(panel, text="分辨率").grid(row=5, column=0, sticky="w", pady=(8, 0)) ttk.Entry(panel, textvariable=self.var_width, width=8).grid(row=5, column=1, sticky="w", pady=(8, 0)) ttk.Label(panel, text="x").grid(row=5, column=2, pady=(8, 0)) ttk.Entry(panel, textvariable=self.var_height, width=8).grid(row=5, column=3, 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)) 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 browse_config(self) -> None: initial = self.var_config.get() initial_dir = str(Path(initial).parent) if initial else str(ROOT) path = filedialog.askopenfilename( title="选择检测配置", initialdir=initial_dir, filetypes=[("YAML", "*.yaml;*.yml"), ("所有文件", "*.*")], ) if path: self.var_config.set(path) self.reload_config() def reload_config(self, silent: bool = False) -> None: try: path = Path(self.var_config.get().strip()) if not path.is_absolute(): path = ROOT / path self.cfg = load_config(path) self.var_config.set(str(path)) self.stabilizer = HoleStabilizer(self.cfg) if not silent: self.var_status.set(f"已加载配置: {path.name}") except Exception as exc: self.cfg = {} self.stabilizer = HoleStabilizer() if not silent: messagebox.showerror("配置", str(exc)) self.var_status.set(f"配置加载失败: {exc}") 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") cam = open_source("realsense", width, height, serial or devices[0][1], 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") self.var_status.set("摄像头已打开。点「开始推理」加载权重并检测螺纹孔") 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.reload_config(silent=True) if not self.cfg: messagebox.showerror("配置", "请先加载有效的 hole_detect_config.yaml") return self.btn_infer.config(state="disabled") self.var_status.set(f"正在加载模型: {model_path.name} …") 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.cfg.get("conf", 0.25)), ) if warmup is not None: detect_holes(warmup, model, self.cfg, device=self.var_device.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 self.stabilizer.reset() 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 show = frame result = None if self.infer_enabled and self.model is not None: raw = detect_holes(frame, self.model, self.cfg, device=self.var_device.get()) result = self.stabilizer.update(frame, raw) show = result["draw_img"] self.last_result = result n += 1 if n >= 8: now = time.perf_counter() self.fps = n / max(1e-6, now - t0) t0, n = now, 0 overlay = show.copy() if result is None: mode = "PREVIEW" color = (0, 200, 255) info = "" elif result["status"]: mode = "OK" color = (0, 255, 0) pts = result["hole_points"] info = f" holes={pts}" else: mode = "FAIL" color = (0, 0, 255) info = f" {result['msg']}" cv2.putText( overlay, f"{mode} FPS {self.fps:.1f}{info}", (16, overlay.shape[0] - 20), cv2.FONT_HERSHEY_SIMPLEX, 0.7, color, 2, ) with self.lock: self.latest = overlay 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_holes" out_dir.mkdir(parents=True, exist_ok=True) out = out_dir / f"hole_{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__": HoleInferApp().mainloop()