first
This commit is contained in:
orisys
@@ -0,0 +1,479 @@
|
||||
"""工件螺纹孔二阶段视频推理：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()
|
||||
Reference in New Issue
Block a user