From e355757b43b28cc66828430468a0c69d80744ef8 Mon Sep 17 00:00:00 2001 From: Arkylin Date: Wed, 22 Apr 2026 10:49:07 +0800 Subject: [PATCH] =?UTF-8?q?=E8=AE=BE=E7=BD=AE=E7=9B=AE=E5=BD=95=E9=80=89?= =?UTF-8?q?=E6=8B=A9?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- main.py | 312 ++++++++++++++++++++++++++++++++++++++++++-------------- 1 file changed, 237 insertions(+), 75 deletions(-) diff --git a/main.py b/main.py index 6f5f5ef..bfb4a51 100644 --- a/main.py +++ b/main.py @@ -53,7 +53,6 @@ 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, MAX_DEPTH_MM = 200, 5000 _DISCONNECT_TIMEOUTS = 6 # 6 × 500 ms = 3 s 无帧 → 断线 @@ -81,16 +80,32 @@ class AppSettings: FONT_SCALE_AUTO = "auto" def __init__(self): - self.font_scale: str = self.FONT_SCALE_AUTO - self.splat_radius: int = 1 # 0=关闭, 1=3×3, 2=5×5, 3=7×7 + self.font_scale: str = self.FONT_SCALE_AUTO + self.splat_radius: int = 1 # 0=关闭, 1=3×3, 2=5×5, 3=7×7 + self.save_base_dir: str = "" # 空 = 当前目录 + self.capture_rgb: bool = True + self.capture_depth: bool = True + self.capture_d2c: bool = True + self.capture_irl: bool = False + self.capture_irr: bool = False + self.capture_pc: bool = False + self.capture_pc_color:bool = False self._load() def _load(self): if self._CONFIG_PATH.is_file(): try: - data = json.loads(self._CONFIG_PATH.read_text(encoding="utf-8")) - self.font_scale = data.get("font_scale", self.FONT_SCALE_AUTO) - self.splat_radius = int(data.get("splat_radius", 1)) + d = json.loads(self._CONFIG_PATH.read_text(encoding="utf-8")) + self.font_scale = d.get("font_scale", self.FONT_SCALE_AUTO) + self.splat_radius = int(d.get("splat_radius", 1)) + self.save_base_dir = d.get("save_base_dir", "") + self.capture_rgb = bool(d.get("capture_rgb", True)) + self.capture_depth = bool(d.get("capture_depth", True)) + self.capture_d2c = bool(d.get("capture_d2c", True)) + self.capture_irl = bool(d.get("capture_irl", False)) + self.capture_irr = bool(d.get("capture_irr", False)) + self.capture_pc = bool(d.get("capture_pc", False)) + self.capture_pc_color= bool(d.get("capture_pc_color", False)) except Exception: pass @@ -98,7 +113,18 @@ class AppSettings: self._APP_DIR.mkdir(parents=True, exist_ok=True) self._CONFIG_PATH.write_text( json.dumps( - {"font_scale": self.font_scale, "splat_radius": self.splat_radius}, + { + "font_scale": self.font_scale, + "splat_radius": self.splat_radius, + "save_base_dir": self.save_base_dir, + "capture_rgb": self.capture_rgb, + "capture_depth": self.capture_depth, + "capture_d2c": self.capture_d2c, + "capture_irl": self.capture_irl, + "capture_irr": self.capture_irr, + "capture_pc": self.capture_pc, + "capture_pc_color": self.capture_pc_color, + }, ensure_ascii=False, indent=2, ), encoding="utf-8", @@ -129,7 +155,7 @@ class SettingsDialog(QDialog): def __init__(self, settings: "AppSettings", parent=None): super().__init__(parent) self.setWindowTitle("程序设置") - self.setFixedSize(_sz(440), _sz(210)) + self.setFixedSize(_sz(520), _sz(360)) self.setWindowFlags(self.windowFlags() & ~Qt.WindowContextHelpButtonHint) self._settings = settings self._build_ui() @@ -187,6 +213,60 @@ class SettingsDialog(QDialog): splat_info.setStyleSheet(f"color:#888; font-size:{_fs(11)}px;") layout.addWidget(splat_info) + # ── 分隔 ── + sep = QLabel("── 采集设置 ──") + sep.setAlignment(Qt.AlignCenter) + sep.setStyleSheet(f"color:#666; font-size:{_fs(11)}px; margin-top:{_sz(4)}px;") + layout.addWidget(sep) + + # 保存根目录 + dir_row = QHBoxLayout() + dir_lbl = QLabel("根目录：") + dir_lbl.setMinimumWidth(_sz(80)) + dir_row.addWidget(dir_lbl) + self._base_dir_edit = QLineEdit(self._settings.save_base_dir) + self._base_dir_edit.setPlaceholderText("留空 = 当前工作目录") + dir_row.addWidget(self._base_dir_edit) + browse_btn = QPushButton("浏览…") + browse_btn.setFixedWidth(_sz(60)) + browse_btn.clicked.connect(self._browse_save_dir) + dir_row.addWidget(browse_btn) + layout.addLayout(dir_row) + + # 默认采集内容 + cap_lbl = QLabel("默认采集内容：") + layout.addWidget(cap_lbl) + + self._cap_checks: dict = {} + cap_items = [ + ("capture_rgb", "RGB"), + ("capture_depth", "Depth"), + ("capture_d2c", "Depth D2C"), + ("capture_irl", "IR-L"), + ("capture_irr", "IR-R"), + ("capture_pc", "点云"), + ("capture_pc_color", "颜色点云"), + ] + cap_row1 = QHBoxLayout() + cap_row1.setSpacing(_sz(8)) + for attr, label in cap_items[:5]: + cb = QCheckBox(label) + cb.setChecked(getattr(self._settings, attr, True)) + cap_row1.addWidget(cb) + self._cap_checks[attr] = cb + cap_row1.addStretch() + layout.addLayout(cap_row1) + + cap_row2 = QHBoxLayout() + cap_row2.setSpacing(_sz(8)) + for attr, label in cap_items[5:]: + cb = QCheckBox(label) + cb.setChecked(getattr(self._settings, attr, False)) + cap_row2.addWidget(cb) + self._cap_checks[attr] = cb + cap_row2.addStretch() + layout.addLayout(cap_row2) + layout.addStretch() btn_row = QHBoxLayout() @@ -201,9 +281,19 @@ class SettingsDialog(QDialog): btn_row.addWidget(cancel_btn) layout.addLayout(btn_row) + def _browse_save_dir(self): + d = QFileDialog.getExistingDirectory( + self, "选择数据保存根目录", + self._base_dir_edit.text() or str(Path.cwd())) + if d: + self._base_dir_edit.setText(d) + def _save_and_close(self): - self._settings.font_scale = self._scale_combo.currentData() - self._settings.splat_radius = self._splat_combo.currentData() + self._settings.font_scale = self._scale_combo.currentData() + self._settings.splat_radius = self._splat_combo.currentData() + self._settings.save_base_dir = self._base_dir_edit.text().strip() + for attr, cb in self._cap_checks.items(): + setattr(self._settings, attr, cb.isChecked()) self._settings.save() self.accept() @@ -345,6 +435,7 @@ class CameraWorker(QThread): self._disconnect_emitted = False self._sync_wait_count = 0 # 等待同步帧的计数器 self._splat_radius = 0 + self._capture_opts: dict = {} def set_d2c(self, enabled: bool): with self._lock: @@ -354,10 +445,11 @@ class CameraWorker(QThread): with self._lock: self._splat_radius = radius - def trigger_capture(self, save_dir: str): + def trigger_capture(self, save_dir: str, capture_opts: dict | None = None): with self._lock: - self._capture_dir = save_dir - self._do_capture = True + self._capture_dir = save_dir + self._capture_opts = capture_opts or {} + self._do_capture = True def stop(self): self._running = False @@ -498,6 +590,7 @@ class CameraWorker(QThread): do_cap = self._do_capture cap_dir = self._capture_dir splat_r = self._splat_radius + cap_opts = self._capture_opts # 首帧提取相机参数 if self._camera_params is None: @@ -530,7 +623,7 @@ class CameraWorker(QThread): self.log_msg.emit(f"帧同步良好（RGB-D 差 {ts_diff}ms），开始拍摄") try: self._run_capture( - frames, cap_dir, + frames, cap_dir, cap_opts, align_filter, point_cloud_filter, frame_to_bgr_image, save_point_cloud_to_ply, OBFrameType, OBFormat, @@ -586,7 +679,7 @@ class CameraWorker(QThread): # ── Capture ── - def _run_capture(self, frames, save_dir, + def _run_capture(self, frames, save_dir, capture_opts, align_filter, point_cloud_filter, frame_to_bgr_image, save_point_cloud_to_ply, OBFrameType, OBFormat): @@ -596,57 +689,66 @@ class CameraWorker(QThread): cf = frames.get_color_frame() df = frames.get_depth_frame() - if cf: + if capture_opts.get("rgb", True) and cf: bgr = frame_to_bgr_image(cf) if bgr is not None: imwrite_utf8(str(base / "color.png"), bgr) if self._ir_enabled: if self._has_dual_ir: - img = process_ir_frame(frames.get_frame(OBFrameType.LEFT_IR_FRAME)) - if img is not None: - imwrite_utf8(str(base / "ir_left.png"), img) - img = process_ir_frame(frames.get_frame(OBFrameType.RIGHT_IR_FRAME)) - if img is not None: - imwrite_utf8(str(base / "ir_right.png"), img) + if capture_opts.get("irl", True): + img = process_ir_frame(frames.get_frame(OBFrameType.LEFT_IR_FRAME)) + if img is not None: + imwrite_utf8(str(base / "ir_left.png"), img) + if capture_opts.get("irr", True): + img = process_ir_frame(frames.get_frame(OBFrameType.RIGHT_IR_FRAME)) + if img is not None: + imwrite_utf8(str(base / "ir_right.png"), img) else: - img = process_ir_frame(frames.get_frame(OBFrameType.IR_FRAME)) - if img is not None: - imwrite_utf8(str(base / "ir.png"), img) + if capture_opts.get("irl", True): + img = process_ir_frame(frames.get_frame(OBFrameType.IR_FRAME)) + if img is not None: + imwrite_utf8(str(base / "ir.png"), img) - if df: + if capture_opts.get("depth", True) and df: raw = np.frombuffer(df.get_data(), dtype=np.uint16).reshape( df.get_height(), df.get_width()) imwrite_utf8(str(base / "depth.png"), raw) - 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) - imwrite_utf8(str(base / "depth_d2c_vis.png"), - depth_to_vis(ad, adf.get_depth_scale())) + need_aligned = (capture_opts.get("d2c", True) or + capture_opts.get("pc", False) or + capture_opts.get("pc_color", False)) + if need_aligned: + aligned = align_filter.process(frames) + if aligned: + aset = aligned.as_frame_set() + adf = aset.get_depth_frame() + if adf and capture_opts.get("d2c", True): + ad = np.frombuffer(adf.get_data(), dtype=np.uint16).reshape( + adf.get_height(), adf.get_width()) + imwrite_utf8(str(base / "depth_d2c.png"), ad) + imwrite_utf8(str(base / "depth_d2c_vis.png"), + depth_to_vis(ad, adf.get_depth_scale())) - 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}") + if capture_opts.get("pc", False): + 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}") + if capture_opts.get("pc_color", False): + 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) @@ -932,12 +1034,27 @@ class MainWindow(QMainWindow): tabs.addTab(self._build_d2c_tab(), " Orbbec Viewer 深度图 D2C ") root.addWidget(tabs) + def _browse_base_dir(self): + d = QFileDialog.getExistingDirectory( + self, "选择数据保存根目录", + self.base_dir_edit.text() or str(Path.cwd())) + if d: + self.base_dir_edit.setText(d) + self._app_settings.save_base_dir = d + self._app_settings.save() + def _show_settings(self): dlg = SettingsDialog(self._app_settings, self) if dlg.exec_() == QDialog.Accepted: # 立即更新当前流的 splat 半径 if self._worker and self.capture_splat_check.isChecked(): self._worker.set_splat(max(1, self._app_settings.splat_radius)) + # 同步根目录显示 + self.base_dir_edit.setText( + self._app_settings.save_base_dir or str(Path.cwd())) + # 同步采集内容复选框为新默认值 + for key, cb in self._capture_checks.items(): + cb.setChecked(getattr(self._app_settings, f"capture_{key}", cb.isChecked())) def _on_capture_splat_changed(self, state: int): radius = max(1, self._app_settings.splat_radius) if state == Qt.Checked else 0 @@ -1023,40 +1140,83 @@ class MainWindow(QMainWindow): g2.addWidget(self.stream_btn) layout.addWidget(grp2) - # ── 第 3 步：保存目录 ── - grp3 = self._step_group(3, "保存目录") - g3 = QHBoxLayout(grp3) - g3.setContentsMargins(8, 4, 8, 6) - g3.setSpacing(8) + # ── 第 3 步：保存目录 & 采集内容 ── + grp3 = self._step_group(3, "保存目录 & 采集内容") + vg3 = QVBoxLayout(grp3) + vg3.setContentsMargins(8, 4, 8, 8) + vg3.setSpacing(5) + # 行 1：根目录选择 + row_base = QHBoxLayout() + row_base.setSpacing(6) + lbl_base = QLabel("根目录：") + lbl_base.setFixedWidth(_sz(58)) + row_base.addWidget(lbl_base) + self.base_dir_edit = QLineEdit( + self._app_settings.save_base_dir or str(Path.cwd())) + row_base.addWidget(self.base_dir_edit) + browse_base_btn = QPushButton("浏览…") + browse_base_btn.setFixedWidth(_sz(58)) + browse_base_btn.setFixedHeight(_sz(26)) + browse_base_btn.clicked.connect(self._browse_base_dir) + row_base.addWidget(browse_base_btn) + vg3.addLayout(row_base) + + # 行 2：子目录 + 状态 + D2C + Splat + row_ses = QHBoxLayout() + row_ses.setSpacing(8) + lbl_ses = QLabel("子目录：") + lbl_ses.setFixedWidth(_sz(58)) + row_ses.addWidget(lbl_ses) today = datetime.now() self.session_edit = QLineEdit(f"{today.year}年{today.month}月{today.day}日") - self.session_edit.setMinimumWidth(_sz(160)) - g3.addWidget(self.session_edit) - + self.session_edit.setMinimumWidth(_sz(140)) + row_ses.addWidget(self.session_edit) create_btn = QPushButton("创建 / 载入") - create_btn.setFixedHeight(_sz(28)) + create_btn.setFixedHeight(_sz(26)) create_btn.clicked.connect(self._create_session) - g3.addWidget(create_btn) - + row_ses.addWidget(create_btn) self.session_status = QLabel("（未创建）") self.session_status.setStyleSheet(f"color:gray; font-size:{_fs(12)}px;") - g3.addWidget(self.session_status) - g3.addStretch() - - self.d2c_check = QCheckBox("启用 D2C 对齐预览") - self.d2c_check.setToolTip("勾选后预览使用 D2C 对齐；拍照时始终保存全部数据") + row_ses.addWidget(self.session_status) + row_ses.addStretch() + self.d2c_check = QCheckBox("D2C 预览") + self.d2c_check.setToolTip("勾选后预览使用 D2C 对齐；拍照时始终保存全部选中的数据") self.d2c_check.stateChanged.connect( lambda s: self._worker and self._worker.set_d2c(s == Qt.Checked)) - g3.addWidget(self.d2c_check) - - g3.addSpacing(8) + row_ses.addWidget(self.d2c_check) + row_ses.addSpacing(6) self.capture_splat_check = QCheckBox("Splat 填充") self.capture_splat_check.setToolTip( "对深度图应用 Splat 扩散填充，减少 D2C 空洞（大小在「设置→程序设置」中配置）") self.capture_splat_check.setChecked(self._app_settings.splat_radius > 0) self.capture_splat_check.stateChanged.connect(self._on_capture_splat_changed) - g3.addWidget(self.capture_splat_check) + row_ses.addWidget(self.capture_splat_check) + vg3.addLayout(row_ses) + + # 行 3：采集内容复选框 + row_cap = QHBoxLayout() + row_cap.setSpacing(_sz(8)) + lbl_cap = QLabel("采集内容：") + lbl_cap.setFixedWidth(_sz(58)) + row_cap.addWidget(lbl_cap) + self._capture_checks: dict = {} + for key, label, default in [ + ("rgb", "RGB", self._app_settings.capture_rgb), + ("depth", "Depth", self._app_settings.capture_depth), + ("d2c", "Depth D2C", self._app_settings.capture_d2c), + ("irl", "IR-L", self._app_settings.capture_irl), + ("irr", "IR-R", self._app_settings.capture_irr), + ("pc", "点云", self._app_settings.capture_pc), + ("pc_color", "颜色点云", self._app_settings.capture_pc_color), + ]: + cb = QCheckBox(label) + cb.setChecked(default) + row_cap.addWidget(cb) + self._capture_checks[key] = cb + row_cap.addStretch() + vg3.addLayout(row_cap) + layout.addWidget(grp3) # ── 预览区 ── @@ -1406,9 +1566,10 @@ class MainWindow(QMainWindow): def _create_session(self): name = self.session_edit.text().strip() if not name: - QMessageBox.warning(self, "提示", "请输入目录名称。") + QMessageBox.warning(self, "提示", "请输入子目录名称。") return - session_dir = DATA_DIR / name + base = Path(self.base_dir_edit.text().strip() or ".") + session_dir = base / name session_dir.mkdir(parents=True, exist_ok=True) self._session_dir = session_dir if self._worker and self._worker._camera_params: @@ -1441,7 +1602,8 @@ class MainWindow(QMainWindow): ts = datetime.now().strftime("%Y%m%d%H%M%S") self.capture_btn.setEnabled(False) self.capture_btn.setText("拍摄中…") - self._worker.trigger_capture(str(self._session_dir / ts)) + capture_opts = {k: cb.isChecked() for k, cb in self._capture_checks.items()} + self._worker.trigger_capture(str(self._session_dir / ts), capture_opts) def _on_capture_done(self, save_dir: str): self.capture_btn.setEnabled(True)