modify ccpd

neka-nat · May 10, 2024 · c553922 · c553922
1 parent 6da9474
commit c553922
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 7 deletions.
diff --git a/examples/color_cpd_rigid.py b/examples/color_cpd_rigid.py
@@ -7,16 +7,18 @@
 log = logging.getLogger('probreg')
 log.setLevel(logging.DEBUG)
 
-voxel_size = 0.05
+voxel_size = 0.03
 source = o3.io.read_point_cloud("frag_115.ply")
 source = source.voxel_down_sample(voxel_size=voxel_size)
 
 target = o3.io.read_point_cloud("frag_116.ply")
 target = target.voxel_down_sample(voxel_size=voxel_size)
 
-cbs = [callbacks.Open3dVisualizerCallback(source, target)]
+cbs = [callbacks.Open3dVisualizerCallback(source, target, save=True)]
 tf_param, _, _ = cpd.registration_cpd(source, target,
-                                      callbacks=cbs, use_color=True)
+                                      callbacks=cbs,
+                                      use_color=True,
+                                      update_scale=False)
 
 print("result: ", np.rad2deg(t3d.euler.mat2euler(tf_param.rot)),
       tf_param.scale, tf_param.t)
diff --git a/probreg/cpd.py b/probreg/cpd.py
@@ -89,14 +89,13 @@ def expectation_step(
         """Expectation step for CPD"""
         assert t_source.ndim == 2 and target.ndim == 2, "source and target must have 2 dimensions."
         pmat = self._compute_pmat_numerator(t_source[:, : self._N_DIM], target[:, : self._N_DIM], sigma2)
-        if self._use_color:
-            pmat_c = self._compute_pmat_numerator(t_source[:, 3:], target[:, 3:], sigma2_c)
 
         c = (2.0 * np.pi * sigma2) ** (self._N_DIM * 0.5)
         c *= w / (1.0 - w) * t_source.shape[0] / target.shape[0]
         den = self.xp.sum(pmat, axis=0)
         den[den == 0] = self.xp.finfo(np.float32).eps
         if self._use_color:
+            pmat_c = self._compute_pmat_numerator(t_source[:, self._N_DIM :], target[:, self._N_DIM :], sigma2_c)
             den_c = self.xp.sum(pmat_c, axis=0)
             den_c[den_c == 0] = self.xp.finfo(np.float32).eps
             den = np.multiply(den, den_c)
@@ -106,10 +105,9 @@ def expectation_step(
             )
             den += o_c
             c *= (2.0 * np.pi * sigma2_c) ** (self._N_COLOR * 0.5)
+            pmat = self.xp.multiply(pmat, pmat_c)
         den += c
 
-        if self._use_color:
-            pmat = self.xp.multiply(pmat, pmat_c)
         pmat = self.xp.divide(pmat, den)
         pt1 = self.xp.sum(pmat, axis=0)
         p1 = self.xp.sum(pmat, axis=1)