i/Orbbec_D2C
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

设置目录选择

Browse Source
This commit is contained in:
i
2026-04-22 10:49:07 +08:00
parent 2996ba96e7
commit e355757b43
1 changed files with 237 additions and 75 deletions
Download Patch File Download Diff File Expand all files Collapse all files
main.py
+237 -75
View File
@@ -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)