orisys
1058 lines
48 KiB
Python
1058 lines
48 KiB
Python
"""RealSense D435/D405 双目采集：左右红外、RGB、可选深度，抽清晰帧后删除原视频。"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import shutil
|
||
import threading
|
||
import time
|
||
import tkinter as tk
|
||
from datetime import datetime
|
||
from pathlib import Path
|
||
from tkinter import filedialog, messagebox, ttk
|
||
|
||
import cv2
|
||
import numpy as np
|
||
import pyrealsense2 as rs
|
||
from PIL import Image, ImageTk
|
||
|
||
ROOT = Path(__file__).resolve().parent
|
||
DEFAULT_OUT = ROOT / "datasets" / "my_dataset" / "images" / "train"
|
||
RESOLUTIONS = [(1280, 720), (848, 480), (640, 480)]
|
||
COLOR_RESOLUTIONS = [(1280, 720), (1920, 1080), (848, 480), (640, 480)]
|
||
|
||
|
||
def sharpness(frame: np.ndarray) -> float:
|
||
gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) if frame.ndim == 3 else frame
|
||
return float(cv2.Laplacian(gray, cv2.CV_64F).var())
|
||
|
||
|
||
def pair_score(left: np.ndarray, right: np.ndarray, color: np.ndarray | None = None) -> float:
|
||
"""左右红外平均清晰度；有 RGB 时加权，避免弱纹理红外分数过低抽不到帧。"""
|
||
s_ir = 0.5 * (sharpness(left) + sharpness(right))
|
||
if color is None:
|
||
return s_ir
|
||
s_rgb = sharpness(color)
|
||
return 0.35 * s_ir + 0.65 * s_rgb
|
||
|
||
|
||
def to_bgr(frame: np.ndarray) -> np.ndarray:
|
||
if frame.ndim == 2:
|
||
return cv2.cvtColor(frame, cv2.COLOR_GRAY2BGR)
|
||
return frame
|
||
|
||
|
||
def list_realsense_devices() -> list[tuple[str, str]]:
|
||
devices = []
|
||
for dev in rs.context().query_devices():
|
||
name = dev.get_info(rs.camera_info.name)
|
||
serial = dev.get_info(rs.camera_info.serial_number)
|
||
devices.append((f"{name} [{serial}]", serial))
|
||
return devices
|
||
|
||
|
||
def extract_clear_pairs(
|
||
left_path: Path,
|
||
right_path: Path,
|
||
out_dir: Path,
|
||
prefix: str,
|
||
min_sharpness: float,
|
||
interval: float,
|
||
max_frames: int,
|
||
rgb_path: Path | None = None,
|
||
depth_vis_path: Path | None = None,
|
||
depth_raw_dir: Path | None = None,
|
||
) -> tuple[int, str]:
|
||
"""抽取清晰成对帧。返回 (保存数量, 说明)。阈值过高时自动按相对清晰度回退。"""
|
||
cap_l = cv2.VideoCapture(str(left_path))
|
||
cap_r = cv2.VideoCapture(str(right_path))
|
||
cap_c = cv2.VideoCapture(str(rgb_path)) if rgb_path is not None else None
|
||
cap_d = cv2.VideoCapture(str(depth_vis_path)) if depth_vis_path is not None else None
|
||
if not cap_l.isOpened() or not cap_r.isOpened():
|
||
cap_l.release()
|
||
cap_r.release()
|
||
if cap_c is not None:
|
||
cap_c.release()
|
||
if cap_d is not None:
|
||
cap_d.release()
|
||
raise RuntimeError("无法读取左右目视频")
|
||
if rgb_path is not None and (cap_c is None or not cap_c.isOpened()):
|
||
raise RuntimeError("无法读取 RGB 视频")
|
||
if depth_vis_path is not None and (cap_d is None or not cap_d.isOpened()):
|
||
raise RuntimeError("无法读取深度视频")
|
||
|
||
fps = cap_l.get(cv2.CAP_PROP_FPS) or 30
|
||
min_gap = max(1, int(round(interval * fps)))
|
||
left_dir = out_dir / "left"
|
||
right_dir = out_dir / "right"
|
||
rgb_dir = out_dir / "rgb"
|
||
depth_dir = out_dir / "depth"
|
||
depth_vis_dir = out_dir / "depth_vis"
|
||
left_dir.mkdir(parents=True, exist_ok=True)
|
||
right_dir.mkdir(parents=True, exist_ok=True)
|
||
if rgb_path is not None:
|
||
rgb_dir.mkdir(parents=True, exist_ok=True)
|
||
if depth_vis_path is not None:
|
||
depth_dir.mkdir(parents=True, exist_ok=True)
|
||
depth_vis_dir.mkdir(parents=True, exist_ok=True)
|
||
|
||
# 每隔 interval 取最清晰一帧；阈值稍后过滤，过高则自动回退
|
||
windows: list[tuple[float, np.ndarray, np.ndarray, np.ndarray | None, np.ndarray | None, int]] = []
|
||
best: tuple[float, np.ndarray, np.ndarray, np.ndarray | None, np.ndarray | None, int] | None = None
|
||
last_saved_idx = -min_gap
|
||
idx = 0
|
||
all_scores: list[float] = []
|
||
|
||
while True:
|
||
ok_l, frame_l = cap_l.read()
|
||
ok_r, frame_r = cap_r.read()
|
||
frame_c = None
|
||
frame_d = None
|
||
if cap_c is not None:
|
||
ok_c, frame_c = cap_c.read()
|
||
if not ok_c:
|
||
break
|
||
if cap_d is not None:
|
||
ok_d, frame_d = cap_d.read()
|
||
if not ok_d:
|
||
break
|
||
if not ok_l or not ok_r:
|
||
break
|
||
score = pair_score(frame_l, frame_r, frame_c)
|
||
all_scores.append(score)
|
||
if best is None or score > best[0]:
|
||
best = (
|
||
score,
|
||
frame_l.copy(),
|
||
frame_r.copy(),
|
||
None if frame_c is None else frame_c.copy(),
|
||
None if frame_d is None else frame_d.copy(),
|
||
idx,
|
||
)
|
||
if best is not None and idx - last_saved_idx >= min_gap:
|
||
windows.append(best)
|
||
last_saved_idx = idx
|
||
best = None
|
||
idx += 1
|
||
|
||
if best is not None:
|
||
windows.append(best)
|
||
|
||
cap_l.release()
|
||
cap_r.release()
|
||
if cap_c is not None:
|
||
cap_c.release()
|
||
if cap_d is not None:
|
||
cap_d.release()
|
||
|
||
if not windows:
|
||
return 0, "视频里没有可读帧"
|
||
|
||
scores = np.asarray(all_scores, dtype=np.float64)
|
||
p50 = float(np.median(scores))
|
||
p80 = float(np.percentile(scores, 80))
|
||
thr = float(min_sharpness)
|
||
selected = [w for w in windows if w[0] >= thr]
|
||
if not selected:
|
||
adaptive = max(0.5, p50 * 0.5)
|
||
selected = [w for w in windows if w[0] >= adaptive]
|
||
if not selected:
|
||
selected = sorted(windows, key=lambda x: x[0], reverse=True)
|
||
keep = len(selected) if max_frames <= 0 else min(len(selected), max(1, max_frames))
|
||
selected = selected[:keep]
|
||
note = (
|
||
f"阈值 {thr:.0f} 过高（本段综合分 median={p50:.0f} p80={p80:.0f}），"
|
||
f"已按相对清晰度自动保存"
|
||
)
|
||
else:
|
||
note = f"按阈值 {thr:.0f} 筛选（本段 median={p50:.0f} p80={p80:.0f}）"
|
||
|
||
selected.sort(key=lambda x: x[5])
|
||
if max_frames > 0 and len(selected) > max_frames:
|
||
selected = sorted(selected, key=lambda x: x[0], reverse=True)[:max_frames]
|
||
selected.sort(key=lambda x: x[5])
|
||
|
||
saved = 0
|
||
for item in selected:
|
||
saved += 1
|
||
name = f"{prefix}_{saved:04d}_s{item[0]:.0f}"
|
||
cv2.imwrite(str(left_dir / f"{name}.jpg"), item[1])
|
||
cv2.imwrite(str(right_dir / f"{name}.jpg"), item[2])
|
||
if item[3] is not None:
|
||
cv2.imwrite(str(rgb_dir / f"{name}.jpg"), item[3])
|
||
if item[4] is not None:
|
||
cv2.imwrite(str(depth_vis_dir / f"{name}.jpg"), item[4])
|
||
raw = None if depth_raw_dir is None else depth_raw_dir / f"{item[5]:06d}.png"
|
||
if raw is not None and raw.exists():
|
||
shutil.copy2(raw, depth_dir / f"{name}.png")
|
||
return saved, note
|
||
|
||
|
||
class CaptureApp(tk.Tk):
|
||
def __init__(self) -> None:
|
||
super().__init__()
|
||
self.title("RealSense D435 / D405 双目采集")
|
||
self.geometry("1480x900")
|
||
self.minsize(1200, 760)
|
||
|
||
self.pipeline: rs.pipeline | None = None
|
||
self.align: rs.align | None = None
|
||
self.colorizer = rs.colorizer()
|
||
self.has_rgb = False
|
||
self.has_depth = False
|
||
self.writer_left: cv2.VideoWriter | None = None
|
||
self.writer_right: cv2.VideoWriter | None = None
|
||
self.writer_rgb: cv2.VideoWriter | None = None
|
||
self.writer_depth: cv2.VideoWriter | None = None
|
||
self.preview_job: str | None = None
|
||
self.photo_left: ImageTk.PhotoImage | None = None
|
||
self.photo_right: ImageTk.PhotoImage | None = None
|
||
self.photo_rgb: ImageTk.PhotoImage | None = None
|
||
self.photo_depth: ImageTk.PhotoImage | None = None
|
||
self.busy = False
|
||
self.recording = False
|
||
self.stop_flag = False
|
||
self.rec_start = 0.0
|
||
self.rec_frames = 0
|
||
self.video_left: Path | None = None
|
||
self.video_right: Path | None = None
|
||
self.video_rgb: Path | None = None
|
||
self.video_depth: Path | None = None
|
||
self.depth_raw_dir: Path | None = None
|
||
self.devices: list[tuple[str, str]] = []
|
||
self.depth_sensor = None
|
||
self.color_sensor = None
|
||
self.shared_exposure = False
|
||
|
||
self._build()
|
||
self.protocol("WM_DELETE_WINDOW", self.on_close)
|
||
self.after(200, self.refresh_devices)
|
||
|
||
def _build(self) -> None:
|
||
self.columnconfigure(0, weight=1)
|
||
self.rowconfigure(0, weight=1)
|
||
|
||
self.preview = ttk.Frame(self, padding=8)
|
||
self.preview.grid(row=0, column=0, sticky="nsew")
|
||
|
||
self.lbl_left = ttk.Label(self.preview, text="左目 Infrared 1", font=("Segoe UI", 11, "bold"))
|
||
self.lbl_right = ttk.Label(self.preview, text="右目 Infrared 2", font=("Segoe UI", 11, "bold"))
|
||
self.lbl_rgb = ttk.Label(self.preview, text="RGB Color", font=("Segoe UI", 11, "bold"))
|
||
self.lbl_depth = ttk.Label(self.preview, text="Depth 深度图", font=("Segoe UI", 11, "bold"))
|
||
self.view_left = tk.Label(self.preview, bg="#111", fg="#eee", text="左目")
|
||
self.view_right = tk.Label(self.preview, bg="#111", fg="#eee", text="右目")
|
||
self.view_rgb = tk.Label(self.preview, bg="#111", fg="#eee", text="RGB")
|
||
self.view_depth = tk.Label(self.preview, bg="#111", fg="#eee", text="深度图")
|
||
self.status = ttk.Label(self.preview, text="就绪")
|
||
|
||
panel = ttk.LabelFrame(self, text="RealSense 采集参数", padding=10)
|
||
panel.grid(row=1, column=0, sticky="ew", padx=8, pady=(0, 8))
|
||
for col in range(8):
|
||
panel.columnconfigure(col, weight=1)
|
||
|
||
self.var_device = tk.StringVar()
|
||
self.var_width = tk.IntVar(value=1280)
|
||
self.var_height = tk.IntVar(value=720)
|
||
self.var_fps = tk.IntVar(value=30)
|
||
self.var_duration = tk.DoubleVar(value=10)
|
||
self.var_sharp = tk.DoubleVar(value=5)
|
||
self.var_interval = tk.DoubleVar(value=0.2)
|
||
self.var_max = tk.IntVar(value=0)
|
||
self.var_swap = tk.BooleanVar(value=False)
|
||
self.var_emitter_off = tk.BooleanVar(value=True)
|
||
self.var_rgb = tk.BooleanVar(value=True)
|
||
self.var_align_rgb = tk.BooleanVar(value=True)
|
||
self.var_depth = tk.BooleanVar(value=False)
|
||
self.var_out = tk.StringVar(value=str(DEFAULT_OUT))
|
||
self.var_name = tk.StringVar(value="")
|
||
self.var_overwrite = tk.StringVar(value="rename")
|
||
self.var_preview_rec = tk.BooleanVar(value=False)
|
||
self.var_prev_left = tk.BooleanVar(value=True)
|
||
self.var_prev_right = tk.BooleanVar(value=False)
|
||
self.var_prev_rgb = tk.BooleanVar(value=True)
|
||
self.var_prev_depth = tk.BooleanVar(value=False)
|
||
self.var_ir_auto = tk.BooleanVar(value=True)
|
||
self.var_rgb_auto = tk.BooleanVar(value=True)
|
||
self.var_ir_exposure = tk.DoubleVar(value=8500) # µs，D435 红外常见默认附近
|
||
self.var_rgb_exposure = tk.DoubleVar(value=156) # µs，RGB 常见默认附近
|
||
self.save_prefix = ""
|
||
|
||
ttk.Label(panel, text="设备").grid(row=0, column=0, sticky="w")
|
||
self.cmb_device = ttk.Combobox(panel, textvariable=self.var_device, state="readonly")
|
||
self.cmb_device.grid(row=0, column=1, columnspan=2, sticky="ew")
|
||
ttk.Button(panel, text="刷新设备", command=self.refresh_devices).grid(row=0, column=3, sticky="ew", padx=(8, 0))
|
||
ttk.Checkbutton(panel, text="同时采集 RGB", variable=self.var_rgb, command=self.on_stream_toggle).grid(
|
||
row=0, column=4, sticky="w"
|
||
)
|
||
ttk.Checkbutton(panel, text="同时采集深度图", variable=self.var_depth, command=self.on_depth_toggle).grid(
|
||
row=0, column=5, sticky="w"
|
||
)
|
||
ttk.Checkbutton(panel, text="RGB对齐到左目", variable=self.var_align_rgb, command=self.on_stream_toggle).grid(
|
||
row=0, column=6, sticky="w"
|
||
)
|
||
ttk.Checkbutton(panel, text="交换左右", variable=self.var_swap).grid(row=0, column=7, sticky="w")
|
||
|
||
ttk.Label(panel, text="分辨率").grid(row=1, column=0, sticky="w", pady=(8, 0))
|
||
ttk.Combobox(panel, textvariable=self.var_width, values=[1280, 848, 640], width=8, state="readonly").grid(
|
||
row=1, column=1, sticky="w", pady=(8, 0)
|
||
)
|
||
ttk.Label(panel, text="x").grid(row=1, column=2, pady=(8, 0))
|
||
ttk.Combobox(panel, textvariable=self.var_height, values=[720, 480], width=8, state="readonly").grid(
|
||
row=1, column=3, sticky="w", pady=(8, 0)
|
||
)
|
||
ttk.Label(panel, text="FPS").grid(row=1, column=4, sticky="w", pady=(8, 0))
|
||
ttk.Combobox(panel, textvariable=self.var_fps, values=[30, 15, 6], width=8, state="readonly").grid(
|
||
row=1, column=5, sticky="w", pady=(8, 0)
|
||
)
|
||
self.chk_emitter = ttk.Checkbutton(panel, text="关闭红外点阵(D435)", variable=self.var_emitter_off)
|
||
self.chk_emitter.grid(row=1, column=6, columnspan=2, sticky="w", pady=(8, 0))
|
||
|
||
exp = ttk.LabelFrame(panel, text="曝光（单位 µs，取消自动后可手调）", padding=6)
|
||
exp.grid(row=2, column=0, columnspan=8, sticky="ew", pady=(10, 0))
|
||
self.chk_ir_auto = ttk.Checkbutton(
|
||
exp, text="红外自动曝光", variable=self.var_ir_auto, command=self.on_exposure_mode_change
|
||
)
|
||
self.chk_ir_auto.pack(side="left")
|
||
ttk.Label(exp, text="红外曝光").pack(side="left", padx=(10, 2))
|
||
self.spn_ir_exp = ttk.Spinbox(
|
||
exp,
|
||
from_=1,
|
||
to=165000,
|
||
increment=100,
|
||
textvariable=self.var_ir_exposure,
|
||
width=8,
|
||
command=self.apply_exposure,
|
||
)
|
||
self.spn_ir_exp.pack(side="left")
|
||
self.chk_rgb_auto = ttk.Checkbutton(
|
||
exp, text="RGB自动曝光", variable=self.var_rgb_auto, command=self.on_exposure_mode_change
|
||
)
|
||
self.chk_rgb_auto.pack(side="left", padx=(16, 0))
|
||
ttk.Label(exp, text="RGB曝光").pack(side="left", padx=(10, 2))
|
||
self.spn_rgb_exp = ttk.Spinbox(
|
||
exp,
|
||
from_=1,
|
||
to=10000,
|
||
increment=10,
|
||
textvariable=self.var_rgb_exposure,
|
||
width=8,
|
||
command=self.apply_exposure,
|
||
)
|
||
self.spn_rgb_exp.pack(side="left")
|
||
ttk.Button(exp, text="应用曝光", command=self.apply_exposure).pack(side="left", padx=(12, 0))
|
||
self.lbl_exp_hint = ttk.Label(exp, text="打开相机后改手动值再点应用即可生效")
|
||
self.lbl_exp_hint.pack(side="left", padx=(10, 0))
|
||
self.spn_ir_exp.bind("<Return>", lambda _e: self.apply_exposure())
|
||
self.spn_rgb_exp.bind("<Return>", lambda _e: self.apply_exposure())
|
||
self.on_exposure_mode_change(apply=False)
|
||
|
||
ttk.Label(panel, text="录制秒数").grid(row=3, column=0, sticky="w", pady=(8, 0))
|
||
ttk.Entry(panel, textvariable=self.var_duration, width=8).grid(row=3, column=1, sticky="w", pady=(8, 0))
|
||
ttk.Label(panel, text="清晰度阈值").grid(row=3, column=2, sticky="w", pady=(8, 0))
|
||
sharp_box = ttk.Frame(panel)
|
||
sharp_box.grid(row=3, column=3, sticky="w", pady=(8, 0))
|
||
ttk.Entry(sharp_box, textvariable=self.var_sharp, width=6).pack(side="left")
|
||
ttk.Button(sharp_box, text="宽松", width=4, command=lambda: self.var_sharp.set(2)).pack(side="left", padx=(4, 0))
|
||
ttk.Button(sharp_box, text="适中", width=4, command=lambda: self.var_sharp.set(5)).pack(side="left")
|
||
ttk.Button(sharp_box, text="严格", width=4, command=lambda: self.var_sharp.set(20)).pack(side="left")
|
||
ttk.Label(panel, text="抽帧间隔(秒)").grid(row=3, column=4, sticky="w", pady=(8, 0))
|
||
ttk.Entry(panel, textvariable=self.var_interval, width=8).grid(row=3, column=5, sticky="w", pady=(8, 0))
|
||
ttk.Label(panel, text="最多保存").grid(row=3, column=6, sticky="w", pady=(8, 0))
|
||
ttk.Entry(panel, textvariable=self.var_max, width=8).grid(row=3, column=7, 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_out).grid(row=4, column=1, columnspan=6, sticky="ew", pady=(8, 0))
|
||
ttk.Button(panel, text="浏览", command=self.browse_out).grid(row=4, column=7, sticky="e", pady=(8, 0))
|
||
|
||
ttk.Label(panel, text="命名前缀").grid(row=5, column=0, sticky="w", pady=(8, 0))
|
||
ttk.Entry(panel, textvariable=self.var_name).grid(row=5, column=1, columnspan=2, sticky="ew", pady=(8, 0))
|
||
self.lbl_name_hint = ttk.Label(panel, text="空则用时间戳。示例: 批次A_20260818_173000_0001_s112.jpg")
|
||
self.lbl_name_hint.grid(row=5, column=3, columnspan=5, sticky="w", pady=(8, 0))
|
||
|
||
ttk.Label(panel, text="同名文件").grid(row=6, column=0, sticky="w", pady=(8, 0))
|
||
ow = ttk.Frame(panel)
|
||
ow.grid(row=6, column=1, columnspan=7, sticky="w", pady=(8, 0))
|
||
ttk.Radiobutton(ow, text="自动改名(不覆盖)", variable=self.var_overwrite, value="rename").pack(side="left")
|
||
ttk.Radiobutton(ow, text="覆盖已有同前缀文件", variable=self.var_overwrite, value="overwrite").pack(side="left", padx=(12, 0))
|
||
ttk.Radiobutton(ow, text="每次询问", variable=self.var_overwrite, value="ask").pack(side="left", padx=(12, 0))
|
||
|
||
prev = ttk.LabelFrame(panel, text="录制时 UI 预览（关掉可防止卡死）", padding=6)
|
||
prev.grid(row=7, column=0, columnspan=8, sticky="ew", pady=(10, 0))
|
||
ttk.Checkbutton(prev, text="录制时开启实时渲染", variable=self.var_preview_rec).pack(side="left")
|
||
ttk.Checkbutton(prev, text="左红外", variable=self.var_prev_left).pack(side="left", padx=(12, 0))
|
||
ttk.Checkbutton(prev, text="右红外", variable=self.var_prev_right).pack(side="left", padx=(8, 0))
|
||
ttk.Checkbutton(prev, text="RGB", variable=self.var_prev_rgb).pack(side="left", padx=(8, 0))
|
||
ttk.Checkbutton(prev, text="深度", variable=self.var_prev_depth).pack(side="left", padx=(8, 0))
|
||
ttk.Label(prev, text="未勾选的窗口录制时不刷新画面").pack(side="left", padx=(12, 0))
|
||
|
||
prog = ttk.Frame(panel)
|
||
prog.grid(row=8, column=0, columnspan=8, sticky="ew", pady=(10, 0))
|
||
prog.columnconfigure(1, weight=1)
|
||
ttk.Label(prog, text="进度").grid(row=0, column=0, sticky="w")
|
||
self.progress = ttk.Progressbar(prog, mode="determinate", maximum=100)
|
||
self.progress.grid(row=0, column=1, sticky="ew", padx=8)
|
||
self.lbl_progress = ttk.Label(prog, text="未录制", width=22)
|
||
self.lbl_progress.grid(row=0, column=2, sticky="e")
|
||
|
||
btns = ttk.Frame(panel)
|
||
btns.grid(row=9, column=0, columnspan=8, sticky="ew", pady=(12, 0))
|
||
self.btn_open = ttk.Button(btns, text="打开相机", command=self.open_camera_ui)
|
||
self.btn_open.pack(side="left")
|
||
self.btn_start = ttk.Button(btns, text="开始采集", command=self.start_capture)
|
||
self.btn_start.pack(side="left", padx=8)
|
||
self.btn_stop = ttk.Button(btns, text="停止", command=self.stop_capture, state="disabled")
|
||
self.btn_stop.pack(side="left")
|
||
self.var_name.trace_add("write", lambda *_: self.refresh_name_hint())
|
||
self.refresh_name_hint()
|
||
|
||
self.apply_preview_layout()
|
||
|
||
def apply_preview_layout(self) -> None:
|
||
for child in self.preview.winfo_children():
|
||
child.grid_forget()
|
||
show_depth = self.var_depth.get()
|
||
if show_depth:
|
||
for col in range(2):
|
||
self.preview.columnconfigure(col, weight=1)
|
||
self.preview.columnconfigure(2, weight=0)
|
||
self.preview.rowconfigure(1, weight=1)
|
||
self.preview.rowconfigure(3, weight=1)
|
||
self.lbl_left.grid(row=0, column=0, sticky="w")
|
||
self.lbl_right.grid(row=0, column=1, sticky="w")
|
||
self.view_left.grid(row=1, column=0, sticky="nsew", padx=(0, 4), pady=(0, 4))
|
||
self.view_right.grid(row=1, column=1, sticky="nsew", padx=(4, 0), pady=(0, 4))
|
||
self.lbl_rgb.grid(row=2, column=0, sticky="w")
|
||
self.lbl_depth.grid(row=2, column=1, sticky="w")
|
||
self.view_rgb.grid(row=3, column=0, sticky="nsew", padx=(0, 4), pady=(4, 0))
|
||
self.view_depth.grid(row=3, column=1, sticky="nsew", padx=(4, 0), pady=(4, 0))
|
||
self.status.grid(row=4, column=0, columnspan=2, sticky="w", pady=(8, 0))
|
||
else:
|
||
for col in range(3):
|
||
self.preview.columnconfigure(col, weight=1)
|
||
self.preview.rowconfigure(1, weight=1)
|
||
self.preview.rowconfigure(3, weight=0)
|
||
self.lbl_left.grid(row=0, column=0, sticky="w")
|
||
self.lbl_right.grid(row=0, column=1, sticky="w")
|
||
self.lbl_rgb.grid(row=0, column=2, sticky="w")
|
||
self.view_left.grid(row=1, column=0, sticky="nsew", padx=(0, 4))
|
||
self.view_right.grid(row=1, column=1, sticky="nsew", padx=4)
|
||
self.view_rgb.grid(row=1, column=2, sticky="nsew", padx=(4, 0))
|
||
self.status.grid(row=2, column=0, columnspan=3, sticky="w", pady=(8, 0))
|
||
|
||
def on_depth_toggle(self) -> None:
|
||
if self.var_depth.get():
|
||
self.var_emitter_off.set(False)
|
||
self.chk_emitter.state(["disabled"])
|
||
else:
|
||
self.chk_emitter.state(["!disabled"])
|
||
self.apply_preview_layout()
|
||
self.on_stream_toggle()
|
||
|
||
def on_stream_toggle(self) -> None:
|
||
if not self.recording and not self.busy and self.devices:
|
||
self.open_camera_ui()
|
||
|
||
def on_exposure_mode_change(self, apply: bool = True) -> None:
|
||
ir_auto = self.var_ir_auto.get()
|
||
rgb_auto = self.var_rgb_auto.get()
|
||
self.spn_ir_exp.state(["disabled"] if ir_auto else ["!disabled"])
|
||
# D405 等机型 RGB/红外共用曝光时，RGB 控件跟随红外
|
||
if self.shared_exposure:
|
||
self.chk_rgb_auto.state(["disabled"])
|
||
self.spn_rgb_exp.state(["disabled"])
|
||
self.var_rgb_auto.set(ir_auto)
|
||
if not ir_auto:
|
||
self.var_rgb_exposure.set(self.var_ir_exposure.get())
|
||
self.lbl_exp_hint.config(text="当前设备 RGB/红外共用曝光，请调左侧红外项")
|
||
else:
|
||
self.chk_rgb_auto.state(["!disabled"])
|
||
self.spn_rgb_exp.state(["disabled"] if rgb_auto else ["!disabled"])
|
||
self.lbl_exp_hint.config(text="打开相机后改手动值再点应用即可生效")
|
||
if apply:
|
||
self.apply_exposure()
|
||
|
||
@staticmethod
|
||
def _set_sensor_option(sensor, option, value: float) -> bool:
|
||
if sensor is None or not sensor.supports(option):
|
||
return False
|
||
rng = sensor.get_option_range(option)
|
||
clamped = max(rng.min, min(rng.max, float(value)))
|
||
if rng.step > 0:
|
||
clamped = rng.min + round((clamped - rng.min) / rng.step) * rng.step
|
||
clamped = max(rng.min, min(rng.max, clamped))
|
||
sensor.set_option(option, clamped)
|
||
return True
|
||
|
||
def _bind_exposure_sensors(self, device, with_rgb: bool) -> None:
|
||
self.depth_sensor = device.first_depth_sensor()
|
||
self.color_sensor = None
|
||
try:
|
||
self.color_sensor = device.first_color_sensor()
|
||
except Exception:
|
||
self.color_sensor = None
|
||
if self.color_sensor is None and with_rgb:
|
||
for sensor in device.query_sensors():
|
||
name = ""
|
||
if sensor.supports(rs.camera_info.name):
|
||
name = sensor.get_info(rs.camera_info.name).lower()
|
||
if "rgb" in name or "color" in name:
|
||
self.color_sensor = sensor
|
||
break
|
||
# D405：彩色与红外同属 Stereo Module，无独立 RGB sensor
|
||
self.shared_exposure = with_rgb and (
|
||
self.color_sensor is None or self.color_sensor is self.depth_sensor
|
||
)
|
||
if self.shared_exposure:
|
||
self.color_sensor = None
|
||
self.on_exposure_mode_change(apply=False)
|
||
|
||
def apply_exposure(self) -> None:
|
||
if self.pipeline is None:
|
||
self.set_status("相机未打开：曝光设置会在下次打开时生效")
|
||
return
|
||
notes: list[str] = []
|
||
try:
|
||
if self.depth_sensor is not None:
|
||
if self.var_ir_auto.get():
|
||
ok = self._set_sensor_option(self.depth_sensor, rs.option.enable_auto_exposure, 1.0)
|
||
notes.append("红外AE开" if ok else "红外AE不支持")
|
||
if self.shared_exposure and ok:
|
||
notes.append("RGB共用AE")
|
||
else:
|
||
self._set_sensor_option(self.depth_sensor, rs.option.enable_auto_exposure, 0.0)
|
||
ok = self._set_sensor_option(
|
||
self.depth_sensor, rs.option.exposure, float(self.var_ir_exposure.get())
|
||
)
|
||
if ok:
|
||
cur = self.depth_sensor.get_option(rs.option.exposure)
|
||
self.var_ir_exposure.set(cur)
|
||
if self.shared_exposure:
|
||
self.var_rgb_exposure.set(cur)
|
||
notes.append(f"共用曝光 {cur:.0f}µs")
|
||
else:
|
||
notes.append(f"红外曝光 {cur:.0f}µs")
|
||
else:
|
||
notes.append("红外手动曝光失败")
|
||
if self.color_sensor is not None and not self.shared_exposure:
|
||
if self.var_rgb_auto.get():
|
||
ok = self._set_sensor_option(self.color_sensor, rs.option.enable_auto_exposure, 1.0)
|
||
notes.append("RGB AE开" if ok else "RGB AE不支持")
|
||
else:
|
||
self._set_sensor_option(self.color_sensor, rs.option.enable_auto_exposure, 0.0)
|
||
ok = self._set_sensor_option(
|
||
self.color_sensor, rs.option.exposure, float(self.var_rgb_exposure.get())
|
||
)
|
||
if ok:
|
||
cur = self.color_sensor.get_option(rs.option.exposure)
|
||
self.var_rgb_exposure.set(cur)
|
||
notes.append(f"RGB曝光 {cur:.0f}µs")
|
||
else:
|
||
notes.append("RGB手动曝光失败")
|
||
except Exception as exc:
|
||
self.set_status(f"曝光设置失败: {exc}")
|
||
messagebox.showerror("曝光", str(exc))
|
||
return
|
||
if notes:
|
||
self.set_status("曝光已应用: " + " / ".join(notes))
|
||
|
||
def set_status(self, text: str) -> None:
|
||
self.status.config(text=text)
|
||
|
||
def refresh_name_hint(self) -> None:
|
||
custom = self.var_name.get().strip() or "时间戳"
|
||
stamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||
prefix = f"{custom}_{stamp}" if self.var_name.get().strip() else stamp
|
||
self.lbl_name_hint.config(text=f"左右/RGB/深度同名。示例: {prefix}_0001_s112.jpg")
|
||
|
||
def make_prefix(self) -> str:
|
||
stamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||
custom = self.var_name.get().strip()
|
||
return f"{custom}_{stamp}" if custom else stamp
|
||
|
||
def existing_prefix_files(self, out_dir: Path, prefix: str) -> list[Path]:
|
||
hits: list[Path] = []
|
||
for folder in ("left", "right", "rgb", "depth", "depth_vis"):
|
||
d = out_dir / folder
|
||
if not d.exists():
|
||
continue
|
||
hits.extend(sorted(d.glob(f"{prefix}_*")))
|
||
return hits
|
||
|
||
def unique_prefix(self, out_dir: Path, prefix: str) -> str:
|
||
candidate = prefix
|
||
n = 2
|
||
while self.existing_prefix_files(out_dir, candidate):
|
||
candidate = f"{prefix}_{n}"
|
||
n += 1
|
||
return candidate
|
||
|
||
def resolve_save_prefix(self, out_dir: Path, prefix: str) -> str | None:
|
||
existing = self.existing_prefix_files(out_dir, prefix)
|
||
if not existing:
|
||
return prefix
|
||
mode = self.var_overwrite.get()
|
||
msg = (
|
||
f"目录中已有 {len(existing)} 个同前缀文件：\n{prefix}_*.jpg/png\n\n"
|
||
"覆盖会删除这些旧文件。左右/RGB/深度使用同一文件名。"
|
||
)
|
||
if mode == "rename":
|
||
new_prefix = self.unique_prefix(out_dir, prefix)
|
||
self.set_status(f"已有同名，自动改为 {new_prefix}")
|
||
return new_prefix
|
||
if mode == "overwrite":
|
||
if not messagebox.askyesno("覆盖确认", msg + "\n\n确定覆盖？"):
|
||
return None
|
||
for path in existing:
|
||
path.unlink(missing_ok=True)
|
||
return prefix
|
||
choice = messagebox.askyesnocancel(
|
||
"同名文件",
|
||
msg + "\n\n是 = 覆盖\n否 = 自动改名\n取消 = 不开始录制",
|
||
)
|
||
if choice is None:
|
||
return None
|
||
if choice:
|
||
for path in existing:
|
||
path.unlink(missing_ok=True)
|
||
return prefix
|
||
return self.unique_prefix(out_dir, prefix)
|
||
|
||
def set_progress(self, ratio: float, text: str) -> None:
|
||
self.progress["value"] = max(0.0, min(100.0, ratio * 100.0))
|
||
self.lbl_progress.config(text=text)
|
||
|
||
def browse_out(self) -> None:
|
||
path = filedialog.askdirectory(initialdir=self.var_out.get() or str(ROOT))
|
||
if path:
|
||
self.var_out.set(path)
|
||
|
||
def current_serial(self) -> str:
|
||
label = self.var_device.get()
|
||
for text, serial in self.devices:
|
||
if text == label:
|
||
return serial
|
||
raise RuntimeError("请先选择 RealSense 设备")
|
||
|
||
def refresh_devices(self) -> None:
|
||
if self.recording or self.busy:
|
||
return
|
||
self.devices = list_realsense_devices()
|
||
labels = [item[0] for item in self.devices]
|
||
self.cmb_device["values"] = labels
|
||
if labels:
|
||
self.var_device.set(labels[0])
|
||
self.set_status(f"找到 {len(labels)} 台 RealSense，可点「打开相机」")
|
||
self.after(100, self.open_camera_ui)
|
||
else:
|
||
self.var_device.set("")
|
||
self.set_status("未检测到 D435/D405。请插上相机，并确认已安装 Intel RealSense 驱动")
|
||
|
||
def release_writers(self) -> None:
|
||
for attr in ("writer_left", "writer_right", "writer_rgb", "writer_depth"):
|
||
writer = getattr(self, attr)
|
||
if writer is not None:
|
||
writer.release()
|
||
setattr(self, attr, None)
|
||
|
||
def stop_pipeline(self) -> None:
|
||
if self.preview_job is not None:
|
||
self.after_cancel(self.preview_job)
|
||
self.preview_job = None
|
||
self.release_writers()
|
||
self.align = None
|
||
self.has_rgb = False
|
||
self.has_depth = False
|
||
self.depth_sensor = None
|
||
self.color_sensor = None
|
||
self.shared_exposure = False
|
||
if self.pipeline is not None:
|
||
try:
|
||
self.pipeline.stop()
|
||
except Exception:
|
||
pass
|
||
self.pipeline = None
|
||
|
||
def _try_start(self, serial: str, w: int, h: int, f: int, with_rgb: bool, with_depth: bool) -> rs.pipeline:
|
||
pipeline = rs.pipeline()
|
||
config = rs.config()
|
||
config.enable_device(serial)
|
||
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)
|
||
if with_depth:
|
||
config.enable_stream(rs.stream.depth, w, h, rs.format.z16, f)
|
||
if with_rgb:
|
||
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)
|
||
self.pipeline = pipeline
|
||
self.has_rgb = with_rgb
|
||
self.has_depth = with_depth
|
||
need_align = (with_rgb and self.var_align_rgb.get()) or with_depth
|
||
self.align = rs.align(rs.stream.infrared) if need_align else None
|
||
self.var_width.set(w)
|
||
self.var_height.set(h)
|
||
device = profile.get_device()
|
||
self._bind_exposure_sensors(device, with_rgb)
|
||
if self.depth_sensor is not None and self.depth_sensor.supports(rs.option.emitter_enabled):
|
||
# 采深度时强制打开点阵；否则按勾选关闭
|
||
emitter_on = 1.0 if with_depth else (0.0 if self.var_emitter_off.get() else 1.0)
|
||
self.depth_sensor.set_option(rs.option.emitter_enabled, emitter_on)
|
||
self.apply_exposure()
|
||
return pipeline
|
||
|
||
def start_pipeline(self) -> None:
|
||
serial = self.current_serial()
|
||
width, height, fps = int(self.var_width.get()), int(self.var_height.get()), int(self.var_fps.get())
|
||
attempts = [(width, height, fps)] + [(w, h, fps) for w, h in RESOLUTIONS if (w, h) != (width, height)]
|
||
want_rgb = self.var_rgb.get()
|
||
want_depth = self.var_depth.get()
|
||
last_error = None
|
||
rgb_options = [True, False] if want_rgb else [False]
|
||
depth_options = [True, False] if want_depth else [False]
|
||
for with_depth in depth_options:
|
||
for with_rgb in rgb_options:
|
||
for w, h, f in attempts:
|
||
try:
|
||
self._try_start(serial, w, h, f, with_rgb, with_depth)
|
||
if want_depth and not with_depth:
|
||
self.set_status("深度流无法打开，已回退")
|
||
if want_rgb and not with_rgb:
|
||
self.set_status("RGB 流无法打开，已回退")
|
||
return
|
||
except Exception as exc:
|
||
last_error = exc
|
||
try:
|
||
if self.pipeline is not None:
|
||
self.pipeline.stop()
|
||
except Exception:
|
||
pass
|
||
self.pipeline = None
|
||
raise RuntimeError(f"无法打开 RealSense: {last_error}")
|
||
|
||
def open_camera_ui(self) -> None:
|
||
if self.recording or self.busy:
|
||
return
|
||
if not self.devices:
|
||
self.refresh_devices()
|
||
if not self.devices:
|
||
messagebox.showerror("RealSense", "没有检测到 D435/D405")
|
||
return
|
||
self.stop_pipeline()
|
||
try:
|
||
self.start_pipeline()
|
||
except Exception as exc:
|
||
self.view_left.config(image="", text=str(exc))
|
||
self.view_right.config(image="", text="")
|
||
self.view_rgb.config(image="", text="")
|
||
self.view_depth.config(image="", text="")
|
||
self.set_status(str(exc))
|
||
messagebox.showerror("RealSense", str(exc))
|
||
return
|
||
parts = ["左右红外"]
|
||
if self.has_rgb:
|
||
parts.append("RGB" + ("对齐左目" if self.var_align_rgb.get() else ""))
|
||
if self.has_depth:
|
||
parts.append("深度(点阵已开)")
|
||
self.set_status(f"{self.var_width.get()}x{self.var_height.get()} @{self.var_fps.get()} " + " + ".join(parts))
|
||
if not self.has_rgb:
|
||
self.view_rgb.config(image="", text="未开启 RGB")
|
||
if not self.has_depth:
|
||
self.view_depth.config(image="", text="未开启深度")
|
||
self.loop_preview()
|
||
|
||
def read_frames(self) -> tuple[np.ndarray, np.ndarray, np.ndarray | None, np.ndarray | None, np.ndarray | None] | None:
|
||
if self.pipeline is None:
|
||
return None
|
||
try:
|
||
frames = self.pipeline.wait_for_frames(1000)
|
||
except Exception:
|
||
return None
|
||
if self.align is not None:
|
||
frames = self.align.process(frames)
|
||
ir_l = frames.get_infrared_frame(1)
|
||
ir_r = frames.get_infrared_frame(2)
|
||
if not ir_l or not ir_r:
|
||
return None
|
||
left = to_bgr(np.asanyarray(ir_l.get_data()))
|
||
right = to_bgr(np.asanyarray(ir_r.get_data()))
|
||
if self.var_swap.get():
|
||
left, right = right, left
|
||
color = None
|
||
depth_raw = None
|
||
depth_vis = None
|
||
if self.has_rgb:
|
||
color_frame = frames.get_color_frame()
|
||
if color_frame:
|
||
color = np.asanyarray(color_frame.get_data())
|
||
if self.has_depth:
|
||
depth_frame = frames.get_depth_frame()
|
||
if depth_frame:
|
||
depth_raw = np.asanyarray(depth_frame.get_data())
|
||
depth_vis = np.asanyarray(self.colorizer.colorize(depth_frame).get_data())
|
||
return left, right, color, depth_raw, depth_vis
|
||
|
||
def loop_preview(self) -> None:
|
||
pack = self.read_frames()
|
||
if pack is not None:
|
||
left, right, color, depth_raw, depth_vis = pack
|
||
if self.recording:
|
||
duration = max(0.001, float(self.var_duration.get()))
|
||
elapsed = time.perf_counter() - self.rec_start
|
||
remain = max(0.0, duration - elapsed)
|
||
self.set_progress(elapsed / duration, f"录制 {elapsed:.1f}/{duration:.1f}s")
|
||
if self.writer_left is not None and self.writer_right is not None:
|
||
self.writer_left.write(left)
|
||
self.writer_right.write(right)
|
||
if self.writer_rgb is not None and color is not None:
|
||
self.writer_rgb.write(color)
|
||
if self.writer_depth is not None and depth_vis is not None:
|
||
self.writer_depth.write(depth_vis)
|
||
if self.depth_raw_dir is not None and depth_raw is not None:
|
||
cv2.imwrite(str(self.depth_raw_dir / f"{self.rec_frames:06d}.png"), depth_raw)
|
||
self.rec_frames += 1
|
||
if self.var_preview_rec.get():
|
||
if self.var_prev_left.get():
|
||
self.show_frame(self.view_left, left, "left")
|
||
if self.var_prev_right.get():
|
||
self.show_frame(self.view_right, right, "right")
|
||
if self.var_prev_rgb.get() and color is not None:
|
||
self.show_frame(self.view_rgb, color, "rgb")
|
||
if self.var_prev_depth.get() and depth_vis is not None:
|
||
self.show_frame(self.view_depth, depth_vis, "depth")
|
||
if remain <= 0 or self.stop_flag:
|
||
self.after(10, self.finish_capture)
|
||
return
|
||
else:
|
||
score_l, score_r = sharpness(left), sharpness(right)
|
||
score = pair_score(left, right, color)
|
||
thr = self.var_sharp.get()
|
||
show_l, show_r = left.copy(), right.copy()
|
||
ok = score >= thr
|
||
cv2.putText(
|
||
show_l,
|
||
f"L {score_l:.0f} score={score:.0f}/{thr:.0f}",
|
||
(18, 40),
|
||
cv2.FONT_HERSHEY_SIMPLEX,
|
||
0.7,
|
||
(0, 220, 0) if ok else (0, 180, 255),
|
||
2,
|
||
)
|
||
cv2.putText(
|
||
show_r,
|
||
f"R {score_r:.0f} score={score:.0f}/{thr:.0f}",
|
||
(18, 40),
|
||
cv2.FONT_HERSHEY_SIMPLEX,
|
||
0.7,
|
||
(0, 220, 0) if ok else (0, 180, 255),
|
||
2,
|
||
)
|
||
self.show_frame(self.view_left, show_l, "left")
|
||
self.show_frame(self.view_right, show_r, "right")
|
||
if color is not None:
|
||
show_c = color.copy()
|
||
s_rgb = sharpness(color)
|
||
cv2.putText(
|
||
show_c,
|
||
f"RGB {s_rgb:.0f} score={score:.0f}/{thr:.0f}",
|
||
(18, 40),
|
||
cv2.FONT_HERSHEY_SIMPLEX,
|
||
0.7,
|
||
(0, 220, 0) if ok else (0, 180, 255),
|
||
2,
|
||
)
|
||
self.show_frame(self.view_rgb, show_c, "rgb")
|
||
if depth_vis is not None:
|
||
show_d = depth_vis.copy()
|
||
cv2.putText(show_d, "DEPTH", (18, 40), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 220, 0), 2)
|
||
self.show_frame(self.view_depth, show_d, "depth")
|
||
delay = 30 if self.recording and not self.var_preview_rec.get() else 10
|
||
self.preview_job = self.after(delay, self.loop_preview)
|
||
|
||
def show_frame(self, widget: tk.Label, frame: np.ndarray, side: str) -> None:
|
||
rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
||
h, w = rgb.shape[:2]
|
||
max_w = max(widget.winfo_width(), 240)
|
||
max_h = max(widget.winfo_height(), 180)
|
||
scale = min(max_w / w, max_h / h)
|
||
if scale <= 0:
|
||
return
|
||
size = (max(1, int(w * scale)), max(1, int(h * scale)))
|
||
img = Image.fromarray(rgb).resize(size, Image.Resampling.BILINEAR)
|
||
photo = ImageTk.PhotoImage(img)
|
||
if side == "left":
|
||
self.photo_left = photo
|
||
elif side == "right":
|
||
self.photo_right = photo
|
||
elif side == "rgb":
|
||
self.photo_rgb = photo
|
||
else:
|
||
self.photo_depth = photo
|
||
widget.config(image=photo, text="")
|
||
|
||
def start_capture(self) -> None:
|
||
if self.pipeline is None:
|
||
self.open_camera_ui()
|
||
if self.pipeline is None:
|
||
return
|
||
if self.recording or self.busy:
|
||
return
|
||
duration = float(self.var_duration.get())
|
||
if duration <= 0:
|
||
messagebox.showwarning("参数", "录制秒数必须大于 0")
|
||
return
|
||
pack = self.read_frames()
|
||
if pack is None:
|
||
messagebox.showerror("采集", "读不到画面")
|
||
return
|
||
left, right, color, _, depth_vis = pack
|
||
out_dir = Path(self.var_out.get())
|
||
prefix = self.resolve_save_prefix(out_dir, self.make_prefix())
|
||
if prefix is None:
|
||
self.set_status("已取消录制")
|
||
return
|
||
self.save_prefix = prefix
|
||
fps = max(1, int(self.var_fps.get()))
|
||
tmp_dir = ROOT / "datasets" / "_tmp"
|
||
tmp_dir.mkdir(parents=True, exist_ok=True)
|
||
self.video_left = tmp_dir / f"{prefix}_left.avi"
|
||
self.video_right = tmp_dir / f"{prefix}_right.avi"
|
||
self.video_rgb = tmp_dir / f"{prefix}_rgb.avi" if self.has_rgb else None
|
||
self.video_depth = tmp_dir / f"{prefix}_depth.avi" if self.has_depth else None
|
||
self.depth_raw_dir = tmp_dir / f"{prefix}_depth_raw" if self.has_depth else None
|
||
if self.depth_raw_dir is not None:
|
||
self.depth_raw_dir.mkdir(parents=True, exist_ok=True)
|
||
h_l, w_l = left.shape[:2]
|
||
h_r, w_r = right.shape[:2]
|
||
self.writer_left = cv2.VideoWriter(str(self.video_left), cv2.VideoWriter_fourcc(*"XVID"), fps, (w_l, h_l))
|
||
self.writer_right = cv2.VideoWriter(str(self.video_right), cv2.VideoWriter_fourcc(*"XVID"), fps, (w_r, h_r))
|
||
if self.has_rgb and color is not None and self.video_rgb is not None:
|
||
h_c, w_c = color.shape[:2]
|
||
self.writer_rgb = cv2.VideoWriter(str(self.video_rgb), cv2.VideoWriter_fourcc(*"XVID"), fps, (w_c, h_c))
|
||
if self.has_depth and depth_vis is not None and self.video_depth is not None:
|
||
h_d, w_d = depth_vis.shape[:2]
|
||
self.writer_depth = cv2.VideoWriter(str(self.video_depth), cv2.VideoWriter_fourcc(*"XVID"), fps, (w_d, h_d))
|
||
if not self.writer_left.isOpened() or not self.writer_right.isOpened():
|
||
self.release_writers()
|
||
messagebox.showerror("采集", "无法写入视频")
|
||
return
|
||
|
||
self.recording = True
|
||
self.stop_flag = False
|
||
self.rec_start = time.perf_counter()
|
||
self.rec_frames = 0
|
||
self.btn_start.config(state="disabled")
|
||
self.btn_open.config(state="disabled")
|
||
self.btn_stop.config(state="normal")
|
||
extra = ""
|
||
if self.has_rgb:
|
||
extra += " + RGB"
|
||
if self.has_depth:
|
||
extra += " + 深度"
|
||
preview_hint = "预览已关" if not self.var_preview_rec.get() else "预览已开"
|
||
self.set_progress(0, f"录制 0.0/{duration:.1f}s")
|
||
self.set_status(f"正在同步录制左右红外{extra} {duration:.1f} 秒… 命名 {prefix} {preview_hint}")
|
||
|
||
def stop_capture(self) -> None:
|
||
if self.recording:
|
||
self.stop_flag = True
|
||
|
||
def finish_capture(self) -> None:
|
||
if not self.recording:
|
||
return
|
||
self.recording = False
|
||
self.btn_stop.config(state="disabled")
|
||
self.release_writers()
|
||
|
||
left_path, right_path = self.video_left, self.video_right
|
||
rgb_path, depth_path, depth_raw_dir = self.video_rgb, self.video_depth, self.depth_raw_dir
|
||
if left_path is None or right_path is None or not left_path.exists() or not right_path.exists():
|
||
self.btn_start.config(state="normal")
|
||
self.btn_open.config(state="normal")
|
||
self.set_status("没有录到左右目视频")
|
||
return
|
||
if self.has_rgb and (rgb_path is None or not rgb_path.exists()):
|
||
rgb_path = None
|
||
if self.has_depth and (depth_path is None or not depth_path.exists()):
|
||
depth_path = None
|
||
|
||
self.busy = True
|
||
self.set_progress(1.0, "正在抽帧…")
|
||
self.set_status("正在抽取清晰帧…")
|
||
prefix = self.save_prefix or left_path.stem.replace("_left", "")
|
||
out_dir = Path(self.var_out.get())
|
||
min_sharp = float(self.var_sharp.get())
|
||
interval = float(self.var_interval.get())
|
||
max_frames = int(self.var_max.get())
|
||
has_rgb = self.has_rgb
|
||
has_depth = self.has_depth
|
||
|
||
def work() -> None:
|
||
error = ""
|
||
saved = 0
|
||
note = ""
|
||
try:
|
||
saved, note = extract_clear_pairs(
|
||
left_path,
|
||
right_path,
|
||
out_dir,
|
||
prefix,
|
||
min_sharp,
|
||
interval,
|
||
max_frames,
|
||
rgb_path,
|
||
depth_path,
|
||
depth_raw_dir,
|
||
)
|
||
except Exception as exc:
|
||
error = str(exc)
|
||
finally:
|
||
left_path.unlink(missing_ok=True)
|
||
right_path.unlink(missing_ok=True)
|
||
if rgb_path is not None:
|
||
rgb_path.unlink(missing_ok=True)
|
||
if depth_path is not None:
|
||
depth_path.unlink(missing_ok=True)
|
||
if depth_raw_dir is not None and depth_raw_dir.exists():
|
||
shutil.rmtree(depth_raw_dir, ignore_errors=True)
|
||
|
||
def done() -> None:
|
||
self.busy = False
|
||
self.btn_start.config(state="normal")
|
||
self.btn_open.config(state="normal")
|
||
if error:
|
||
self.set_progress(0, "抽帧失败")
|
||
self.set_status(error)
|
||
messagebox.showerror("抽帧失败", error)
|
||
return
|
||
folders = f"{out_dir}/left 、 {out_dir}/right"
|
||
if has_rgb:
|
||
folders += f" 、 {out_dir}/rgb"
|
||
if has_depth:
|
||
folders += f" 、 {out_dir}/depth （16位）与 {out_dir}/depth_vis"
|
||
if saved == 0:
|
||
msg = "仍未抽到帧，请检查相机画面或缩短抽帧间隔后再试。原视频已删除"
|
||
else:
|
||
msg = f"已保存 {saved} 组，前缀 {prefix}，到 {folders}，原视频已删除"
|
||
if note:
|
||
msg += f"\n\n{note}"
|
||
self.set_progress(1.0, f"完成 {saved} 组")
|
||
self.set_status(msg.replace("\n", " "))
|
||
messagebox.showinfo("采集完成", msg + f"\n\n文件名格式: {prefix}_0001_s清晰度.jpg")
|
||
|
||
self.after(0, done)
|
||
|
||
threading.Thread(target=work, daemon=True).start()
|
||
|
||
def on_close(self) -> None:
|
||
self.recording = False
|
||
self.stop_pipeline()
|
||
for path in (self.video_left, self.video_right, self.video_rgb, self.video_depth):
|
||
if path is not None:
|
||
path.unlink(missing_ok=True)
|
||
if self.depth_raw_dir is not None and self.depth_raw_dir.exists():
|
||
shutil.rmtree(self.depth_raw_dir, ignore_errors=True)
|
||
self.destroy()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
CaptureApp().mainloop()
|