Orbbec_D2C/main.py

#!/usr/bin/env python3
"""
Orbbec 数据采集 GUI
- 启动时选择彩色 / 深度分辨率
- 数据采集：实时预览 + 一键拍照（RGB / IR / 深度 / D2C深度 / 点云）
- Orbbec Viewer 深度图 D2C：批量软件 D2C 转换
"""

import os
import sys
import json
import time
import threading
import tempfile
import shutil
import contextlib
import io
import numpy as np
import cv2
from pathlib import Path
from datetime import datetime

# 确保 utils.py 可被 import
sys.path.insert(0, str(Path(__file__).parent))

from PyQt5.QtWidgets import (
    QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout,
    QLabel, QPushButton, QLineEdit, QCheckBox, QTabWidget,
    QComboBox, QDialog, QDialogButtonBox, QGroupBox, QFileDialog,
    QTextEdit, QFrame, QMessageBox, QSizePolicy,
)
from PyQt5.QtCore import Qt, QThread, pyqtSignal, QTimer
from PyQt5.QtGui import QImage, QPixmap, QFont

from d2c_batch import batch_convert

DATA_DIR = Path("data")
PARAMS_FILE = "camera_params.json"
MIN_DEPTH_MM = 200
MAX_DEPTH_MM = 5000


# ─────────────────────────────────────────────────────────────────
#  Helpers
# ─────────────────────────────────────────────────────────────────

def ndarray_to_pixmap(img_bgr: np.ndarray, max_w: int, max_h: int) -> QPixmap:
    """BGR numpy array → scaled QPixmap."""
    if img_bgr is None or img_bgr.size == 0:
        return QPixmap()
    h, w = img_bgr.shape[:2]
    scale = min(max_w / max(w, 1), max_h / max(h, 1), 1.0)
    if scale < 1.0:
        img_bgr = cv2.resize(img_bgr, (int(w * scale), int(h * scale)),
                             interpolation=cv2.INTER_AREA)
    rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
    rh, rw = rgb.shape[:2]
    qimg = QImage(rgb.data.tobytes(), rw, rh, rw * 3, QImage.Format_RGB888)
    return QPixmap.fromImage(qimg)


def imwrite_utf8(path: str, img: np.ndarray) -> bool:
    """cv2.imwrite 的 Unicode 路径安全版本（Windows 中文路径兼容）。"""
    ext = Path(path).suffix.lower() or ".png"
    ok, buf = cv2.imencode(ext, img)
    if ok:
        Path(path).write_bytes(buf.tobytes())
    return ok


def save_ply_utf8(target_path: str, pc_frame, save_func) -> None:
    """将点云先写到 ASCII 临时路径，再移动到 Unicode 目标路径。"""
    tmpdir = tempfile.mkdtemp()
    try:
        tmp_ply = os.path.join(tmpdir, "pc.ply")
        save_func(tmp_ply, pc_frame)
        shutil.move(tmp_ply, target_path)
    finally:
        shutil.rmtree(tmpdir, ignore_errors=True)


def depth_to_vis(depth_uint16: np.ndarray, scale: float) -> np.ndarray:
    """uint16 深度图 → JET 伪彩 BGR 图（无效像素为黑色）。"""
    depth_mm = depth_uint16.astype(np.float32) * scale
    norm = np.clip((depth_mm - MIN_DEPTH_MM) / (MAX_DEPTH_MM - MIN_DEPTH_MM), 0.0, 1.0)
    gray = (norm * 255).astype(np.uint8)
    vis = cv2.applyColorMap(gray, cv2.COLORMAP_JET)
    vis[depth_uint16 == 0] = 0
    return vis


def process_ir_frame(ir_frame) -> np.ndarray | None:
    """IR frame → 8-bit BGR 图像。"""
    if ir_frame is None:
        return None
    try:
        from pyorbbecsdk import OBFormat
        vf = ir_frame.as_video_frame()
        data = np.asanyarray(vf.get_data())
        w, h, fmt = vf.get_width(), vf.get_height(), vf.get_format()
        if fmt == OBFormat.Y8:
            arr = np.resize(data, (h, w)).astype(np.uint8)
        elif fmt == OBFormat.MJPG:
            arr = cv2.imdecode(data, cv2.IMREAD_GRAYSCALE)
            if arr is None:
                return None
        else:
            arr = np.frombuffer(data, dtype=np.uint16).reshape(h, w)
            arr = cv2.normalize(arr, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8)
        return cv2.cvtColor(arr, cv2.COLOR_GRAY2BGR)
    except Exception:
        return None


# ─────────────────────────────────────────────────────────────────
#  Profile Selection Dialog
# ─────────────────────────────────────────────────────────────────

class ProfileDialog(QDialog):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.setWindowTitle("选择分辨率")
        self.setMinimumWidth(540)
        self.pipeline = None
        self.color_profiles: list = []
        self.depth_profiles: list = []
        self.selected_color = None
        self.selected_depth = None
        self._build_ui()
        self._load_profiles()

    def _build_ui(self):
        layout = QVBoxLayout(self)
        layout.setSpacing(14)

        grp_c = QGroupBox("彩色 (RGB) 分辨率")
        lc = QVBoxLayout(grp_c)
        self.color_combo = QComboBox()
        lc.addWidget(self.color_combo)
        layout.addWidget(grp_c)

        grp_d = QGroupBox("深度分辨率")
        ld = QVBoxLayout(grp_d)
        self.depth_combo = QComboBox()
        ld.addWidget(self.depth_combo)
        layout.addWidget(grp_d)

        self.status_lbl = QLabel("正在连接相机，请稍候…")
        self.status_lbl.setAlignment(Qt.AlignCenter)
        layout.addWidget(self.status_lbl)

        btns = QDialogButtonBox(QDialogButtonBox.Ok | QDialogButtonBox.Cancel)
        btns.accepted.connect(self._on_ok)
        btns.rejected.connect(self.reject)
        layout.addWidget(btns)

    def _load_profiles(self):
        try:
            from pyorbbecsdk import Pipeline, OBSensorType
            self.pipeline = Pipeline()

            color_list = self.pipeline.get_stream_profile_list(OBSensorType.COLOR_SENSOR)
            for i in range(color_list.get_count()):
                p = color_list.get_stream_profile_by_index(i)
                vp = p.as_video_stream_profile()
                if vp is None:
                    continue
                fmt = str(vp.get_format()).split(".")[-1]
                self.color_combo.addItem(
                    f"{vp.get_width()}×{vp.get_height()} @ {vp.get_fps()} fps  [{fmt}]")
                self.color_profiles.append(vp)

            depth_list = self.pipeline.get_stream_profile_list(OBSensorType.DEPTH_SENSOR)
            for i in range(depth_list.get_count()):
                p = depth_list.get_stream_profile_by_index(i)
                vp = p.as_video_stream_profile()
                if vp is None:
                    continue
                fmt = str(vp.get_format()).split(".")[-1]
                self.depth_combo.addItem(
                    f"{vp.get_width()}×{vp.get_height()} @ {vp.get_fps()} fps  [{fmt}]")
                self.depth_profiles.append(vp)

            self.status_lbl.setText(
                f"已找到  {len(self.color_profiles)} 个彩色配置，"
                f"{len(self.depth_profiles)} 个深度配置")
        except Exception as exc:
            self.status_lbl.setText(f"相机连接失败：{exc}")

    def _on_ok(self):
        if not self.color_profiles or not self.depth_profiles:
            QMessageBox.warning(self, "错误", "没有可用的相机配置，请检查设备连接。")
            return
        self.selected_color = self.color_profiles[self.color_combo.currentIndex()]
        self.selected_depth = self.depth_profiles[self.depth_combo.currentIndex()]
        self.accept()


# ─────────────────────────────────────────────────────────────────
#  Camera Worker Thread
# ─────────────────────────────────────────────────────────────────

_DISCONNECT_TIMEOUTS = 6   # 6 × 500 ms = 3 s 无帧视为断开

class CameraWorker(QThread):
    preview_ready  = pyqtSignal(object, object)   # color_bgr, depth_vis
    capture_done   = pyqtSignal(str)              # save_dir
    params_ready   = pyqtSignal(dict)
    log_msg        = pyqtSignal(str)
    disconnected   = pyqtSignal()

    def __init__(self, pipeline, color_profile, depth_profile):
        super().__init__()
        self._pipeline          = pipeline
        self._color_profile     = color_profile
        self._depth_profile     = depth_profile
        self._lock              = threading.Lock()
        self._running           = True
        self._d2c               = False
        self._do_capture        = False
        self._capture_dir       = ""
        self._camera_params     = None
        self._has_dual_ir       = False
        self._last_emit_ts      = 0.0
        self._serial_number     = "unknown"
        self._consec_timeouts   = 0
        self._disconnect_emitted = False

    # ── Public controls ──

    def set_d2c(self, enabled: bool):
        with self._lock:
            self._d2c = enabled

    def trigger_capture(self, save_dir: str):
        with self._lock:
            self._capture_dir = save_dir
            self._do_capture  = True

    def stop(self):
        self._running = False
        self.wait(6000)

    # ── Thread body ──

    def run(self):
        try:
            from pyorbbecsdk import (
                Config, AlignFilter, PointCloudFilter,
                OBStreamType, OBSensorType, OBFrameType, OBFormat,
                save_point_cloud_to_ply,
            )
            from utils import frame_to_bgr_image
        except ImportError as exc:
            self.log_msg.emit(f"导入 pyorbbecsdk 失败：{exc}")
            return

        config = Config()
        config.enable_stream(self._color_profile)
        config.enable_stream(self._depth_profile)

        # 获取设备 SN + 探测 IR 传感器类型
        try:
            device = self._pipeline.get_device()
            try:
                self._serial_number = device.get_device_info().get_serial_number()
            except Exception:
                pass
            sensor_list = device.get_sensor_list()
            for i in range(len(sensor_list)):
                st = sensor_list[i].get_type()
                if st in (OBSensorType.LEFT_IR_SENSOR, OBSensorType.RIGHT_IR_SENSOR):
                    self._has_dual_ir = True
            if self._has_dual_ir:
                config.enable_video_stream(OBSensorType.LEFT_IR_SENSOR)
                config.enable_video_stream(OBSensorType.RIGHT_IR_SENSOR)
            else:
                config.enable_video_stream(OBSensorType.IR_SENSOR)
        except Exception as exc:
            self.log_msg.emit(f"IR 传感器初始化警告：{exc}")

        try:
            self._pipeline.start(config)
        except Exception as exc:
            self.log_msg.emit(f"Pipeline 启动失败：{exc}")
            return

        align_filter       = AlignFilter(align_to_stream=OBStreamType.COLOR_STREAM)
        point_cloud_filter = PointCloudFilter()

        while self._running:
            try:
                frames = self._pipeline.wait_for_frames(500)
                if frames is None:
                    self._consec_timeouts += 1
                    if (self._consec_timeouts >= _DISCONNECT_TIMEOUTS
                            and not self._disconnect_emitted):
                        self._disconnect_emitted = True
                        self._running = False
                        self.disconnected.emit()
                    continue
                self._consec_timeouts = 0

                with self._lock:
                    d2c       = self._d2c
                    do_cap    = self._do_capture
                    cap_dir   = self._capture_dir
                    if do_cap:
                        self._do_capture = False

                # 首帧提取相机参数（含 SN）
                if self._camera_params is None:
                    cf = frames.get_color_frame()
                    df = frames.get_depth_frame()
                    if cf and df:
                        try:
                            params = self._extract_params(cf, df)
                            params['device'] = {'serial_number': self._serial_number}
                            self._camera_params = params
                            self.params_ready.emit(self._camera_params)
                        except Exception:
                            pass

                # 执行拍照
                if do_cap and cap_dir:
                    try:
                        self._run_capture(
                            frames, cap_dir,
                            align_filter, point_cloud_filter,
                            frame_to_bgr_image, save_point_cloud_to_ply,
                            OBFrameType, OBFormat,
                        )
                    except Exception as exc:
                        self.log_msg.emit(f"拍照错误：{exc}")
                    continue

                # ── 预览 ──
                view_frames = frames
                if d2c:
                    af = align_filter.process(frames)
                    if af:
                        view_frames = af.as_frame_set()

                cf = view_frames.get_color_frame()
                df = view_frames.get_depth_frame()
                if not cf or not df:
                    continue

                color_bgr = frame_to_bgr_image(cf)
                if color_bgr is None:
                    continue

                depth_data = np.frombuffer(df.get_data(), dtype=np.uint16).reshape(
                    df.get_height(), df.get_width())
                depth_vis = depth_to_vis(depth_data, df.get_depth_scale())

                # 限制发射帧率 ~30 fps
                now = time.monotonic()
                if now - self._last_emit_ts >= 0.033:
                    self._last_emit_ts = now
                    self.preview_ready.emit(color_bgr.copy(), depth_vis.copy())

            except Exception as exc:
                if self._running:
                    self.log_msg.emit(f"相机错误：{exc}")
                    # 将异常也计入超时计数，快速检测断线
                    self._consec_timeouts += _DISCONNECT_TIMEOUTS
                    if not self._disconnect_emitted:
                        self._disconnect_emitted = True
                        self._running = False
                        self.disconnected.emit()
                    time.sleep(0.1)

        try:
            self._pipeline.stop()
        except Exception:
            pass

    # ── Capture helper ──

    def _run_capture(self, frames, save_dir,
                     align_filter, point_cloud_filter,
                     frame_to_bgr_image, save_point_cloud_to_ply,
                     OBFrameType, OBFormat):
        Path(save_dir).mkdir(parents=True, exist_ok=True)

        cf = frames.get_color_frame()
        df = frames.get_depth_frame()

        base = Path(save_dir)

        # ① 彩色图
        if cf:
            bgr = frame_to_bgr_image(cf)
            if bgr is not None:
                imwrite_utf8(str(base / "color.png"), bgr)

        # ② 红外图
        if self._has_dual_ir:
            lir = frames.get_frame(OBFrameType.LEFT_IR_FRAME)
            rir = frames.get_frame(OBFrameType.RIGHT_IR_FRAME)
            img = process_ir_frame(lir)
            if img is not None:
                imwrite_utf8(str(base / "ir_left.png"), img)
            img = process_ir_frame(rir)
            if img is not None:
                imwrite_utf8(str(base / "ir_right.png"), img)
        else:
            ir = frames.get_frame(OBFrameType.IR_FRAME)
            img = process_ir_frame(ir)
            if img is not None:
                imwrite_utf8(str(base / "ir.png"), img)

        # ③ 原始深度图
        if df:
            raw = np.frombuffer(df.get_data(), dtype=np.uint16).reshape(
                df.get_height(), df.get_width())
            imwrite_utf8(str(base / "depth.png"), raw)

        # ④ D2C 深度图 + 伪彩 + 点云
        aligned = align_filter.process(frames)
        if aligned:
            aset = aligned.as_frame_set()
            adf  = aset.get_depth_frame()
            if adf:
                ad = np.frombuffer(adf.get_data(), dtype=np.uint16).reshape(
                    adf.get_height(), adf.get_width())
                imwrite_utf8(str(base / "depth_d2c.png"), ad)
                vis = depth_to_vis(ad, adf.get_depth_scale())
                imwrite_utf8(str(base / "depth_d2c_vis.png"), vis)

            try:
                point_cloud_filter.set_create_point_format(OBFormat.POINT)
                pc = point_cloud_filter.process(aligned)
                if pc:
                    save_ply_utf8(str(base / "point_cloud.ply"), pc,
                                  save_point_cloud_to_ply)
            except Exception as exc:
                self.log_msg.emit(f"点云警告：{exc}")

            try:
                point_cloud_filter.set_create_point_format(OBFormat.RGB_POINT)
                pc = point_cloud_filter.process(aligned)
                if pc:
                    save_ply_utf8(str(base / "point_cloud_color.ply"), pc,
                                  save_point_cloud_to_ply)
            except Exception as exc:
                self.log_msg.emit(f"彩色点云警告：{exc}")

        self.capture_done.emit(save_dir)

    # ── Static helpers ──

    @staticmethod
    def _extract_params(color_frame, depth_frame) -> dict:
        cvp = color_frame.get_stream_profile().as_video_stream_profile()
        dvp = depth_frame.get_stream_profile().as_video_stream_profile()

        def intr(i):
            return {"fx": float(i.fx), "fy": float(i.fy),
                    "cx": float(i.cx), "cy": float(i.cy),
                    "width": int(i.width), "height": int(i.height)}

        def dist(d):
            return {"k1": float(d.k1), "k2": float(d.k2),
                    "k3": float(d.k3), "k4": float(d.k4),
                    "k5": float(d.k5), "k6": float(d.k6),
                    "p1": float(d.p1), "p2": float(d.p2)}

        ext = dvp.get_extrinsic_to(cvp)
        return {
            "color": {
                "width":  color_frame.get_width(),
                "height": color_frame.get_height(),
                "fps":    cvp.get_fps(),
                "format": str(cvp.get_format()).split(".")[-1],
                "intrinsic":  intr(cvp.get_intrinsic()),
                "distortion": dist(cvp.get_distortion()),
            },
            "depth": {
                "width":  depth_frame.get_width(),
                "height": depth_frame.get_height(),
                "fps":    dvp.get_fps(),
                "format": str(dvp.get_format()).split(".")[-1],
                "scale":  float(depth_frame.get_depth_scale()),
                "intrinsic":  intr(dvp.get_intrinsic()),
                "distortion": dist(dvp.get_distortion()),
            },
            "extrinsic": {
                "rot":       ext.rot.tolist(),
                "transform": ext.transform.tolist(),
            },
        }


# ─────────────────────────────────────────────────────────────────
#  Batch D2C Worker Thread
# ─────────────────────────────────────────────────────────────────

class BatchD2CWorker(QThread):
    log  = pyqtSignal(str)
    done = pyqtSignal()

    def __init__(self, params: dict, depth_dir: str, splat: bool = True):
        super().__init__()
        self._params    = params
        self._depth_dir = depth_dir
        self._splat     = splat

    def run(self):
        class _LineCapture(io.StringIO):
            def __init__(self, sig):
                super().__init__()
                self._sig = sig
            def write(self, s):
                for line in s.splitlines():
                    if line.strip():
                        self._sig.emit(line)

        cap = _LineCapture(self.log)
        with contextlib.redirect_stdout(cap):
            try:
                batch_convert(self._params, self._depth_dir, splat=self._splat)
            except Exception as exc:
                self.log.emit(f"Error: {exc}")
        self.done.emit()


# ─────────────────────────────────────────────────────────────────
#  Main Window
# ─────────────────────────────────────────────────────────────────

class MainWindow(QMainWindow):
    back_to_profiles = pyqtSignal()

    def __init__(self, pipeline, color_profile, depth_profile):
        super().__init__()
        self.setWindowTitle("Orbbec 数据采集 & D2C")
        self.resize(1300, 820)

        self._session_dir: Path | None = None
        self._batch_worker: BatchD2CWorker | None = None
        self._params_filename = "camera_params.json"   # 由 SN+分辨率确定后更新
        self._restart_mode    = False                  # 正在重启，不触发 app.quit

        # 保存选中的流参数，供断线重连时匹配
        self._saved_color_info = {
            "width":  color_profile.get_width(),
            "height": color_profile.get_height(),
            "fps":    color_profile.get_fps(),
        }
        self._saved_depth_info = {
            "width":  depth_profile.get_width(),
            "height": depth_profile.get_height(),
            "fps":    depth_profile.get_fps(),
        }

        self._worker = CameraWorker(pipeline, color_profile, depth_profile)
        self._worker.preview_ready.connect(self._update_preview)
        self._worker.capture_done.connect(self._on_capture_done)
        self._worker.params_ready.connect(self._on_params_ready)
        self._worker.log_msg.connect(self._append_log)
        self._worker.disconnected.connect(self._on_camera_disconnected)

        self._build_ui()
        self._worker.start()

    # ── UI construction ──

    def _build_ui(self):
        central = QWidget()
        self.setCentralWidget(central)
        root = QVBoxLayout(central)
        root.setContentsMargins(0, 0, 0, 0)
        root.setSpacing(0)

        tabs = QTabWidget()
        tabs.addTab(self._build_capture_tab(), "  数据采集  ")
        tabs.addTab(self._build_d2c_tab(),     "  Orbbec Viewer 深度图 D2C  ")
        root.addWidget(tabs)

    def _build_capture_tab(self) -> QWidget:
        w = QWidget()
        layout = QVBoxLayout(w)
        layout.setContentsMargins(12, 10, 12, 10)
        layout.setSpacing(8)

        # ── 顶部控制行 ──
        ctrl = QHBoxLayout()

        ctrl.addWidget(QLabel("保存目录名称："))
        today = datetime.now()
        default_name = f"{today.year}年{today.month}月{today.day}日"
        self.session_edit = QLineEdit(default_name)
        self.session_edit.setMinimumWidth(180)
        ctrl.addWidget(self.session_edit)

        create_btn = QPushButton("创建 / 载入目录")
        create_btn.setFixedHeight(30)
        create_btn.clicked.connect(self._create_session)
        ctrl.addWidget(create_btn)

        self.session_status = QLabel("（未创建）")
        self.session_status.setStyleSheet("color: gray; font-size: 12px;")
        ctrl.addWidget(self.session_status)

        ctrl.addStretch()

        self.d2c_check = QCheckBox("启用 D2C 对齐预览")
        self.d2c_check.setToolTip("勾选后实时预览使用 D2C 对齐；拍照时始终保存全部数据")
        self.d2c_check.stateChanged.connect(
            lambda s: self._worker.set_d2c(s == Qt.Checked))
        ctrl.addWidget(self.d2c_check)

        layout.addLayout(ctrl)

        # ── 预览区 ──
        preview = QFrame()
        preview.setFrameShape(QFrame.StyledPanel)
        preview.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding)
        pv_layout = QHBoxLayout(preview)
        pv_layout.setSpacing(6)
        pv_layout.setContentsMargins(6, 6, 6, 6)

        for attr, title in (("color_lbl", "彩色图"), ("depth_lbl", "深度伪彩图")):
            box = QVBoxLayout()
            hdr = QLabel(title)
            hdr.setAlignment(Qt.AlignCenter)
            hdr.setStyleSheet("font-weight: bold; font-size: 13px;")
            box.addWidget(hdr)
            lbl = QLabel("等待画面…")
            lbl.setAlignment(Qt.AlignCenter)
            lbl.setMinimumSize(560, 400)
            lbl.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding)
            lbl.setStyleSheet("background:#111; border:1px solid #555; color:#888;")
            setattr(self, attr, lbl)
            box.addWidget(lbl)
            pv_layout.addLayout(box)

        layout.addWidget(preview, stretch=1)

        # ── 拍照按钮 ──
        btn_row = QHBoxLayout()
        btn_row.addStretch()
        self.capture_btn = QPushButton("拍照")
        self.capture_btn.setFixedSize(160, 48)
        f = QFont()
        f.setPointSize(14)
        f.setBold(True)
        self.capture_btn.setFont(f)
        self.capture_btn.setStyleSheet(
            "QPushButton{"
            "  background:#2a6ebb; color:white; border-radius:8px;"
            "  border:none;"
            "}"
            "QPushButton:hover{ background:#3a80dd; }"
            "QPushButton:pressed{ background:#1a559e; }"
            "QPushButton:disabled{ background:#555; color:#999; }"
        )
        self.capture_btn.clicked.connect(self._on_capture)
        btn_row.addWidget(self.capture_btn)
        btn_row.addStretch()
        layout.addLayout(btn_row)

        return w

    def _build_d2c_tab(self) -> QWidget:
        w = QWidget()
        layout = QVBoxLayout(w)
        layout.setContentsMargins(12, 10, 12, 10)
        layout.setSpacing(8)

        # 深度图目录
        row1 = QHBoxLayout()
        row1.addWidget(QLabel("深度图目录："))
        self.d2c_dir_edit = QLineEdit()
        self.d2c_dir_edit.setPlaceholderText("选择包含 PNG 深度图的文件夹…")
        row1.addWidget(self.d2c_dir_edit)
        br1 = QPushButton("浏览…")
        br1.setFixedWidth(70)
        br1.clicked.connect(lambda: self._browse_dir(self.d2c_dir_edit))
        row1.addWidget(br1)
        layout.addLayout(row1)

        # 相机参数文件
        row2 = QHBoxLayout()
        row2.addWidget(QLabel("相机参数文件："))
        self.params_path_edit = QLineEdit()
        self.params_path_edit.setPlaceholderText(f"留空则自动查找 {PARAMS_FILE}")
        row2.addWidget(self.params_path_edit)
        br2 = QPushButton("浏览…")
        br2.setFixedWidth(70)
        br2.clicked.connect(lambda: self._browse_json(self.params_path_edit))
        row2.addWidget(br2)
        layout.addLayout(row2)

        # Splat 选项
        opt_row = QHBoxLayout()
        self.splat_check = QCheckBox("Splat 3×3 填充（减少正向投影空洞，默认开启）")
        self.splat_check.setChecked(True)
        self.splat_check.setToolTip(
            "对应 d2c_batch.py 的 splat 参数：将每个深度像素扩散到 3×3 邻域，"
            "有效减少软件 D2C 正向投影产生的空洞。")
        opt_row.addWidget(self.splat_check)
        opt_row.addStretch()
        layout.addLayout(opt_row)

        # 运行按钮
        run_row = QHBoxLayout()
        run_row.addStretch()
        self.run_d2c_btn = QPushButton("▶  开始批量 D2C 转换")
        self.run_d2c_btn.setFixedHeight(40)
        self.run_d2c_btn.setStyleSheet(
            "QPushButton{ background:#2a7a2a; color:white; border-radius:6px; border:none; }"
            "QPushButton:hover{ background:#3a9a3a; }"
            "QPushButton:disabled{ background:#555; color:#999; }"
        )
        self.run_d2c_btn.clicked.connect(self._run_batch_d2c)
        run_row.addWidget(self.run_d2c_btn)
        run_row.addStretch()
        layout.addLayout(run_row)

        # 日志区
        layout.addWidget(QLabel("输出日志："))
        self.d2c_log = QTextEdit()
        self.d2c_log.setReadOnly(True)
        self.d2c_log.setStyleSheet(
            "font-family: Consolas, 'Courier New', monospace; font-size: 12px;")
        layout.addWidget(self.d2c_log, stretch=1)

        return w

    # ── Slots ──

    def _update_preview(self, color_bgr: np.ndarray, depth_vis: np.ndarray):
        w, h = self.color_lbl.width(), self.color_lbl.height()
        self.color_lbl.setPixmap(ndarray_to_pixmap(color_bgr, w, h))
        w2, h2 = self.depth_lbl.width(), self.depth_lbl.height()
        self.depth_lbl.setPixmap(ndarray_to_pixmap(depth_vis, w2, h2))

    def _create_session(self):
        name = self.session_edit.text().strip()
        if not name:
            QMessageBox.warning(self, "提示", "请输入目录名称。")
            return
        session_dir = DATA_DIR / name
        session_dir.mkdir(parents=True, exist_ok=True)
        self._session_dir = session_dir

        # 如果已有相机参数，立即写入
        if self._worker._camera_params:
            self._save_params_to_session(self._worker._camera_params)

        self.session_status.setText(f"✓  {session_dir}")
        self.session_status.setStyleSheet("color: #2a9a2a; font-size: 12px;")

    def _on_params_ready(self, params: dict):
        # 根据 SN + 分辨率/帧率构建参数文件名
        sn   = params.get("device", {}).get("serial_number", "unknown")
        cw   = params["color"]["width"];  ch  = params["color"]["height"]
        cfps = params["color"]["fps"]
        dw   = params["depth"]["width"];  dh  = params["depth"]["height"]
        dfps = params["depth"]["fps"]
        self._params_filename = (
            f"SN{sn}_RGB{cw}x{ch}@{cfps}_D{dw}x{dh}@{dfps}.json"
        )
        if self._session_dir and self._session_dir.exists():
            self._save_params_to_session(params)

    def _save_params_to_session(self, params: dict):
        try:
            p = self._session_dir / self._params_filename
            p.write_text(json.dumps(params, indent=2, ensure_ascii=False),
                         encoding="utf-8")
        except Exception:
            pass

    def _on_capture(self):
        # 若尚未创建目录，自动创建
        if self._session_dir is None:
            self._create_session()
            if self._session_dir is None:
                return

        ts       = datetime.now().strftime("%Y%m%d%H%M%S")
        save_dir = str(self._session_dir / ts)

        self.capture_btn.setEnabled(False)
        self.capture_btn.setText("拍摄中…")
        self._worker.trigger_capture(save_dir)

    def _on_capture_done(self, save_dir: str):
        self.capture_btn.setEnabled(True)
        self.capture_btn.setText("拍照")
        self._show_toast("拍摄完成！", f"已保存至：\n{save_dir}")

    def _show_toast(self, title: str, message: str, ms: int = 3000):
        """短暂显示通知，ms 毫秒后自动关闭。"""
        dlg = QMessageBox(QMessageBox.Information, title, message,
                          QMessageBox.Ok, self)
        dlg.setWindowModality(Qt.NonModal)
        dlg.show()
        QTimer.singleShot(ms, dlg.accept)

    def _append_log(self, msg: str):
        self.d2c_log.append(msg)

    # ── Batch D2C tab ──

    def _browse_dir(self, edit: QLineEdit):
        d = QFileDialog.getExistingDirectory(self, "选择文件夹")
        if d:
            edit.setText(d)

    def _browse_json(self, edit: QLineEdit):
        f, _ = QFileDialog.getOpenFileName(self, "选择相机参数文件", "", "JSON (*.json)")
        if f:
            edit.setText(f)

    def _resolve_params_file(self) -> str | None:
        """返回有效的相机参数 JSON 路径，找不到返回 None。"""
        explicit = self.params_path_edit.text().strip()
        if explicit and Path(explicit).is_file():
            return explicit

        # 优先在 session 目录查找 SN 命名的文件
        if self._session_dir and self._session_dir.exists():
            sn_files = sorted(self._session_dir.glob("SN*.json"))
            if sn_files:
                return str(sn_files[0])
            fallback = self._session_dir / "camera_params.json"
            if fallback.is_file():
                return str(fallback)

        # 工作目录里的 SN 命名文件或旧版 camera_params.json
        sn_files = sorted(Path(".").glob("SN*.json"))
        if sn_files:
            return str(sn_files[0])
        if Path(PARAMS_FILE).is_file():
            return PARAMS_FILE
        return None

    def _run_batch_d2c(self):
        depth_dir = self.d2c_dir_edit.text().strip()
        if not depth_dir:
            QMessageBox.warning(self, "错误", "请先选择深度图目录。")
            return

        params_file = self._resolve_params_file()
        if params_file is None:
            QMessageBox.warning(
                self, "错误",
                "找不到相机参数文件。\n请先在「数据采集」页面连接相机以自动生成，"
                "或手动指定 JSON 文件。")
            return

        with open(params_file) as fh:
            params = json.load(fh)

        self.d2c_log.clear()
        self.d2c_log.append(f"[参数文件] {params_file}")
        self.d2c_log.append(f"[深度图目录] {depth_dir}\n")
        self.run_d2c_btn.setEnabled(False)

        splat = self.splat_check.isChecked()
        self._batch_worker = BatchD2CWorker(params, depth_dir, splat=splat)
        self._batch_worker.log.connect(self.d2c_log.append)
        self._batch_worker.done.connect(self._on_batch_done)
        self._batch_worker.start()

    def _on_batch_done(self):
        self.run_d2c_btn.setEnabled(True)
        self.d2c_log.append("\n[ 转换完成 ]")

    # ── Camera disconnect / reconnect ──

    def _on_camera_disconnected(self):
        self.capture_btn.setEnabled(False)

        dlg = QMessageBox(self)
        dlg.setWindowTitle("相机断开连接")
        dlg.setText("检测到相机已断开连接，请选择操作：")
        dlg.setIcon(QMessageBox.Warning)
        reconnect_btn = dlg.addButton("重新连接相机",   QMessageBox.AcceptRole)
        back_btn      = dlg.addButton("返回分辨率选择", QMessageBox.RejectRole)
        exit_btn      = dlg.addButton("退出程序",       QMessageBox.DestructiveRole)
        dlg.exec_()

        clicked = dlg.clickedButton()
        if clicked == reconnect_btn:
            self._do_reconnect()
        elif clicked == back_btn:
            self._restart_mode = True
            self.hide()
            self.back_to_profiles.emit()   # → _open_session → 新窗口
            self._worker.stop()
            self.close()
        else:
            QApplication.quit()

    def _do_reconnect(self):
        """用相同的分辨率配置重新连接相机。"""
        prog = QMessageBox(QMessageBox.Information, "重新连接",
                           "正在重新连接相机，请稍候…",
                           QMessageBox.NoButton, self)
        prog.setWindowModality(Qt.ApplicationModal)
        prog.show()
        QApplication.processEvents()
        time.sleep(2)

        try:
            from pyorbbecsdk import Pipeline, OBSensorType

            new_pipeline = Pipeline()

            def _best_profile(profile_list, target: dict):
                for i in range(profile_list.get_count()):
                    p  = profile_list.get_stream_profile_by_index(i)
                    vp = p.as_video_stream_profile()
                    if vp is None:
                        continue
                    if (vp.get_width()  == target["width"]
                            and vp.get_height() == target["height"]
                            and vp.get_fps()    == target["fps"]):
                        return vp
                # 没找到精确匹配时用第一个
                p = profile_list.get_stream_profile_by_index(0)
                return p.as_video_stream_profile()

            new_color = _best_profile(
                new_pipeline.get_stream_profile_list(OBSensorType.COLOR_SENSOR),
                self._saved_color_info)
            new_depth = _best_profile(
                new_pipeline.get_stream_profile_list(OBSensorType.DEPTH_SENSOR),
                self._saved_depth_info)

            prog.accept()

            old_worker = self._worker
            self._worker = CameraWorker(new_pipeline, new_color, new_depth)
            self._worker.preview_ready.connect(self._update_preview)
            self._worker.capture_done.connect(self._on_capture_done)
            self._worker.params_ready.connect(self._on_params_ready)
            self._worker.log_msg.connect(self._append_log)
            self._worker.disconnected.connect(self._on_camera_disconnected)
            self._worker.start()

            self.capture_btn.setEnabled(True)
            old_worker.deleteLater()

        except Exception as exc:
            prog.accept()
            QMessageBox.critical(self, "重新连接失败",
                                 f"无法重新连接相机：{exc}\n\n请检查连接后重试。")
            # 再次弹出选项
            self._on_camera_disconnected()

    # ── Window close ──

    def closeEvent(self, event):
        self._worker.stop()
        # _restart_mode 时由 _open_session 保持事件循环，不退出 app
        super().closeEvent(event)


# ─────────────────────────────────────────────────────────────────
#  Entry Point
# ─────────────────────────────────────────────────────────────────

def _open_session(app: QApplication) -> None:
    """显示分辨率选择对话框，成功后打开主窗口；取消则退出 app。"""
    dlg = ProfileDialog()
    if dlg.exec_() != QDialog.Accepted:
        app.quit()
        return

    if dlg.pipeline is None or dlg.selected_color is None or dlg.selected_depth is None:
        QMessageBox.critical(None, "错误", "相机初始化失败，请检查设备连接后重启程序。")
        app.quit()
        return

    win = MainWindow(dlg.pipeline, dlg.selected_color, dlg.selected_depth)
    # 断线→返回选择界面时，递归调用本函数
    win.back_to_profiles.connect(lambda: _open_session(app))
    win.show()


def main():
    app = QApplication(sys.argv)
    app.setStyle("Fusion")

    from PyQt5.QtGui import QPalette, QColor
    palette = QPalette()
    palette.setColor(QPalette.Window,          QColor(45,  45,  45))
    palette.setColor(QPalette.WindowText,      QColor(220, 220, 220))
    palette.setColor(QPalette.Base,            QColor(30,  30,  30))
    palette.setColor(QPalette.AlternateBase,   QColor(50,  50,  50))
    palette.setColor(QPalette.ToolTipBase,     QColor(255, 255, 220))
    palette.setColor(QPalette.ToolTipText,     QColor(0,   0,   0))
    palette.setColor(QPalette.Text,            QColor(220, 220, 220))
    palette.setColor(QPalette.Button,          QColor(60,  60,  60))
    palette.setColor(QPalette.ButtonText,      QColor(220, 220, 220))
    palette.setColor(QPalette.BrightText,      QColor(255, 0,   0))
    palette.setColor(QPalette.Link,            QColor(42,  130, 218))
    palette.setColor(QPalette.Highlight,       QColor(42,  130, 218))
    palette.setColor(QPalette.HighlightedText, QColor(0,   0,   0))
    app.setPalette(palette)

    _open_session(app)
    sys.exit(app.exec_())


if __name__ == "__main__":
    main()