Add PyTorch BA-FPCC-OBB migration: datasets, models, losses, utils and train/test scripts

README_BA_FPCC.md

@@ -0,0 +1,217 @@
+# BA-FPCC-OBB
+
+本 README 是对当前仓库里 BA-FPCC-OBB 改造的**完整运行说明**，目标是：
+- 不破坏原始 FPCC（TensorFlow）流程；
+- 提供可独立运行的 PyTorch 迁移流程；
+- 支持 Boundary Label / Boundary Loss / Boundary Branch / Boundary-aware 后处理 / OBB / SYN-PBOX 单帧转换与可视化。
+
+---
+
+## 0. 分支与安全策略（强烈推荐）
+
+为了确保原始 FPCC 不受影响，请在独立分支开发与验证：
+
+```bash
+git checkout -b feature/ba-fpcc-pytorch-migration
+```
+
+- 原始 TF 入口：`fpcc_train.py`、`fpcc_test.py`
+- 新增 PyTorch 入口：`scripts/train_ba_fpcc.py`、`scripts/test_ba_fpcc.py`
+
+---
+
+## 1. 目录说明（新增模块）
+
+- `datasets/fpcc_dataset.py`：读取原 FPCC h5 列表文件，输出 PyTorch batch 字段。
+- `models/fpcc_pytorch.py`：PyTorch FPCC 主干 + FPCC 风格损失。
+- `models/boundary_branch.py`：边界预测分支。
+- `losses/boundary_loss.py`：boundary 加权损失与 BCE 损失。
+- `utils/boundary_utils.py`：基于 kNN 的 boundary label 自动生成。
+- `utils/obb_utils.py`：Open3D OBB 估计。
+- `utils/simple_visualize.py`：实例/边界/中心/OBB 可视化。
+- `utils/simple_metrics.py`：简化指标函数。
+- `scripts/train_ba_fpcc.py`：PyTorch 训练入口（含 BA 开关）。
+- `scripts/test_ba_fpcc.py`：PyTorch 测试入口（含后处理与 OBB 开关）。
+- `scripts/convert_syn_pbox_one_frame.py`：SYN-PBOX 单帧转换。
+- `scripts/visualize_syn_pbox_one_frame.py`：SYN-PBOX 单帧可视化。
+
+---
+
+## 2. 环境依赖
+
+> 你当前机器如果提示 `ModuleNotFoundError`，先执行这一节。
+
+### 2.1 PyTorch 分支依赖
+
+```bash
+pip install numpy scipy h5py scikit-learn imageio open3d
+pip install torch torchvision torchaudio
+```
+
+### 2.2 原始 FPCC（TensorFlow）依赖
+
+原仓库是 TensorFlow 1.x 风格代码，建议单独环境安装：
+
+```bash
+pip install numpy scipy h5py
+pip install tensorflow==1.15.5
+```
+
+> 若你的系统无法安装 TF1，请只跑 PyTorch 分支，原始分支作为结构参考。
+
+---
+
+## 3. 数据格式说明
+
+### 3.1 PyTorch 训练数据（与原 FPCC 一致）
+
+`--input_list` 指向一个 txt 文件（如 `datas/ring_train.txt`），每行是一个 h5 文件路径。h5 至少包含：
+- `data`：点云输入
+- `pid` 或 `gid`：实例 id
+- `score` 或 `center_score`：中心监督标签
+
+### 3.2 PyTorch 测试数据
+
+默认读取 `datas/ring_test/*.txt`，每行格式近似：
+
+```text
+x y z instance_id
+```
+
+---
+
+## 4. 运行命令
+
+## 4.1 原始 FPCC baseline（TF）
+
+```bash
+python fpcc_train.py --gpu 0 --input_list datas/ring_train.txt
+python fpcc_test.py --gpu 0 --restore_dir checkpoint/
+```
+
+## 4.2 PyTorch BA-FPCC 训练
+
+```bash
+python scripts/train_ba_fpcc.py \
+  --input_list datas/ring_train.txt \
+  --batch_size 4 \
+  --epochs 100 \
+  --point_dim 6 \
+  --num_groups 50 \
+  --use_boundary_loss \
+  --use_boundary_branch \
+  --lambda_boundary 1.0 \
+  --lambda_bce 0.5 \
+  --save_path checkpoint/ba_fpcc_torch.pt
+```
+
+## 4.3 PyTorch BA-FPCC 测试 + Boundary-aware + OBB
+
+```bash
+python scripts/test_ba_fpcc.py \
+  --ckpt checkpoint/ba_fpcc_torch.pt \
+  --test_glob 'datas/ring_test/*.txt' \
+  --point_dim 6 \
+  --point_num 4096 \
+  --center_score_th 0.6 \
+  --r_nms 0.1 \
+  --use_boundary_branch \
+  --boundary_th 0.5 \
+  --use_obb \
+  --save_obb \
+  --obb_json test_results/ba_fpcc_obb.json
+```
+
+## 4.4 SYN-PBOX 单帧转换
+
+```bash
+python scripts/convert_syn_pbox_one_frame.py \
+  --scene_camera path/to/scene_camera.json \
+  --depth path/to/depth.png \
+  --mask_dir path/to/mask_visib \
+  --frame_id 0 \
+  --output demo_syn_frame.npz \
+  --num_points 4096
+```
+
+## 4.5 SYN-PBOX 单帧可视化
+
+```bash
+python scripts/visualize_syn_pbox_one_frame.py --npz demo_syn_frame.npz --show_obb
+```
+
+---
+
+## 5. 开关与消融建议
+
+### 5.1 训练开关（`scripts/train_ba_fpcc.py`）
+- `--use_boundary_loss`：开启 boundary 加权项。
+- `--use_boundary_branch`：开启 boundary branch + BCE。
+- `--lambda_boundary`：boundary 加权强度。
+- `--lambda_bce`：boundary BCE 损失权重。
+
+建议消融顺序：
+1. 仅主干（两个开关都关）
+2. 主干 + boundary loss
+3. 主干 + boundary branch
+4. 全开（boundary loss + boundary branch）
+
+### 5.2 测试开关（`scripts/test_ba_fpcc.py`）
+- `--use_boundary_branch`：加载并使用边界分数。
+- `--boundary_th`：边界点阈值。
+- `--use_obb`：计算 OBB。
+- `--save_obb`：保存 OBB json。
+- `--visualize_obb`：Open3D 交互可视化。
+
+---
+
+## 6. 输出说明
+
+### 6.1 测试输出
+- `test_results/*_pred.txt`：每场景预测实例结果。
+- `test_results/ba_fpcc_obb.json`：OBB 列表（当 `--use_obb --save_obb` 开启时）。
+
+OBB 字段包含：
+- `scene_id`
+- `instance_id`
+- `center`
+- `extent`
+- `R`
+- `yaw`
+- `num_points`
+
+### 6.2 SYN-PBOX NPZ 输出
+`convert_syn_pbox_one_frame.py` 输出字段：
+- `points`
+- `instance_labels`
+- `center_labels`
+- `boundary_labels`
+
+---
+
+## 7. 常见错误与排查
+
+### 7.1 `ModuleNotFoundError: numpy / torch / tensorflow`
+先安装第 2 节依赖。
+
+### 7.2 `No test files matched`
+检查 `--test_glob` 路径与引号是否正确。
+
+### 7.3 OBB 可视化打不开
+检查 `open3d` 是否安装；无桌面环境时不要开 `--visualize_obb`，只保存 json。
+
+### 7.4 TF1 安装困难
+建议优先使用 PyTorch 分支；原 TF 分支可作为对照基线。
+
+---
+
+## 8. 当前实现边界（实话说明）
+
+- PyTorch 主干是“可读优先”的迁移版本，不是逐层 1:1 复刻 TF 主干。
+- 若你要严格复现实验数值，需要继续做：
+  1) backbone 细节对齐；
+  2) loss 细节对齐；
+  3) 数据预处理完全对齐。
+
+但从工程可运行性上，当前分支已经具备完整链路：
+**数据读取 -> 训练 -> 测试聚类 -> boundary 后处理 -> OBB 输出 -> 可视化**。

datasets/__init__.py

@@ -0,0 +1 @@
+"""Dataset package for BA-FPCC-OBB."""

datasets/fpcc_dataset.py

@@ -0,0 +1,54 @@
+"""PyTorch dataset wrappers for original FPCC h5 files."""
+
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+
+import provider
+
+
+class FPCCDataset(Dataset):
+    def __init__(self, file_list_txt, point_dim=6, num_groups=50):
+        self.files = provider.getDataFiles(file_list_txt)
+        self.point_dim = point_dim
+        self.num_groups = num_groups
+
+        self.data = []
+        self.group = []
+        self.center = []
+        for f in self.files:
+            cur_data, cur_group, _, _, cur_score = provider.loadDataFile_with_groupseglabel_stanfordindoor(f)
+            self.data.append(cur_data)
+            self.group.append(cur_group)
+            self.center.append(cur_score)
+
+        self.data = np.concatenate(self.data, axis=0).astype(np.float32)
+        self.group = np.concatenate(self.group, axis=0).astype(np.int64)
+        self.center = np.concatenate(self.center, axis=0).astype(np.float32)
+
+    def __len__(self):
+        return self.data.shape[0]
+
+    def __getitem__(self, idx):
+        points = self.data[idx, :, : self.point_dim]
+        instance_labels = np.asarray(self.group[idx]).reshape(-1)
+        center_labels = np.asarray(self.center[idx]).reshape(-1)
+        one_hot = group_to_one_hot(instance_labels, self.num_groups)
+        return {
+            "points": torch.from_numpy(points),
+            "instance_labels": torch.from_numpy(instance_labels),
+            "center_labels": torch.from_numpy(center_labels),
+            "group_one_hot": torch.from_numpy(one_hot),
+        }
+
+
+def group_to_one_hot(group_labels, num_groups=50):
+    """group_labels: [N], output [N,num_groups]."""
+    group_labels = np.asarray(group_labels)
+    one_hot = np.zeros((group_labels.shape[0], num_groups), dtype=np.float32)
+    uniq = [u for u in np.unique(group_labels) if u >= 0]
+    mapping = {u: i for i, u in enumerate(uniq[:num_groups])}
+    for i, gid in enumerate(group_labels):
+        if gid in mapping:
+            one_hot[i, mapping[gid]] = 1.0
+    return one_hot

docs/fpcc_simple_analysis.md

@@ -0,0 +1,52 @@
+# FPCC 简要结构分析（第一步，仅阅读）
+
+> 说明：原始项目是 **TensorFlow 1.x** 实现，不是 PyTorch。若要在 BA-FPCC-OBB 中统一到 PyTorch，需要新增脚本做兼容/迁移。
+
+1. **训练入口文件**
+   - `fpcc_train.py`。
+
+2. **测试入口文件**
+   - `fpcc_test.py`。
+
+3. **模型 forward 输入/输出**
+   - 输入来自 `models/model.py::get_model(backbone, point_cloud, is_training, ...)`。
+   - `point_cloud` shape 约为 `[B, N, POINT_DIM]`（训练中一般取前 6 维）。
+   - 输出是 dict：
+     - `center_score`: 每点中心置信度，shape 约 `[B, N]`
+     - `point_features`: 每点特征（Fsim），shape 约 `[B, N, 128]`
+     - `simmat`: 相似度矩阵，shape `[B, N, N]`（train=True 时）
+     - `3d_distance`: 点间 3D 距离矩阵，shape `[B, N, N]`
+
+4. **dataloader 返回字段**
+   - `provider.loadDataFile_with_groupseglabel_stanfordindoor(...)` 在训练里被调用，返回：
+     - `cur_data`
+     - `cur_group`
+     - `_`（未使用）
+     - `_`（未使用）
+     - `cur_score`
+
+5. **每个 batch 点云 shape**
+   - 训练 feed 到模型的是 `input_data[..., :POINT_DIM]`，shape `[B, 4096, POINT_DIM]`。
+
+6. **instance label 与 center label shape**
+   - instance label（group label）在 one-hot 后为 `ptsgroup_label_ph`: `[B, N, NUM_GROUPS]`。
+   - center label 为 `pts_score_ph`: `[B, N]`。
+
+7. **loss 在哪里计算**
+   - 在 `models/model.py::get_loss(...)` 计算。
+   - 在 `fpcc_train.py` 中通过 `loss, score_loss, grouperr = model.get_loss(...)` 调用。
+
+8. **聚类在哪里做**
+   - 在 `fpcc_test.py` 中调用 `utils.test_utils` 里的 `GroupMerging_fpcc(...)` 完成实例聚类。
+
+9. **加 boundary loss 建议改哪里**
+   - 最小侵入：新增训练脚本 `scripts/train_ba_fpcc.py`，在原始 loss 基础上额外做逐点加权。
+   - 逐点 loss 可以优先作用在 center 回归项（`ptscenter_loss`）或新增点级监督项。
+
+10. **加 boundary branch 建议改哪里**
+   - 最小侵入：不要改原 `fpcc_train.py`，在新脚本里拿 `net_output['point_features']` 后接一个轻量 BoundaryBranch。
+   - 若 TF 图里插分支过于复杂，可先在 PyTorch 训练流程里演示并保留 TODO。
+
+11. **加 OBB 后处理从哪里拿预测实例点云**
+   - 在 `fpcc_test.py` 聚类后得到 `ins_pre`（每点实例 id）。
+   - 用 `pts[:,0:3]`（或对齐后的 `pts[:,3:6]`）按实例 id 取点，再做 OBB。

losses/__init__.py

@@ -0,0 +1 @@
+"""Loss package for BA-FPCC-OBB."""

losses/boundary_loss.py

@@ -0,0 +1,40 @@
+"""Simple boundary-related loss functions for BA-FPCC-OBB."""
+
+import torch
+import torch.nn.functional as F
+
+
+def boundary_weighted_loss(point_loss, boundary_labels, lambda_boundary=1.0):
+    """Weighted mean point loss.
+
+    point_loss: [B,N] or [N]
+    boundary_labels: same shape, {0,1}
+    """
+    if point_loss.shape != boundary_labels.shape:
+        raise ValueError(f"shape mismatch: point_loss {point_loss.shape} vs boundary_labels {boundary_labels.shape}")
+    weights = 1.0 + lambda_boundary * boundary_labels.float()
+    return (point_loss * weights).mean()
+
+
+def boundary_bce_loss(boundary_logits, boundary_labels):
+    """BCEWithLogits loss for boundary branch.
+
+    boundary_logits: [B,N] or [N]
+    boundary_labels: same shape
+    """
+    if boundary_logits.shape != boundary_labels.shape:
+        raise ValueError(f"shape mismatch: boundary_logits {boundary_logits.shape} vs boundary_labels {boundary_labels.shape}")
+    return F.binary_cross_entropy_with_logits(boundary_logits, boundary_labels.float())
+
+
+if __name__ == "__main__":
+    pl = torch.rand(2, 16)
+    bl = torch.randint(0, 2, (2, 16))
+    lw = boundary_weighted_loss(pl, bl, lambda_boundary=1.0)
+    lw0 = boundary_weighted_loss(pl, bl, lambda_boundary=0.0)
+    print("weighted loss:", lw.item())
+    print("plain mean loss:", lw0.item(), "raw mean:", pl.mean().item())
+
+    logits = torch.randn(2, 16)
+    bce = boundary_bce_loss(logits, bl)
+    print("boundary bce:", bce.item())

models/__init__.py

@@ -0,0 +1 @@
+"""Model package for BA-FPCC-OBB and original FPCC modules."""