orisys
589 lines
24 KiB
Python
589 lines
24 KiB
Python
"""实时摄像头推理：支持 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()
|