Try add depth cov loss to improve quality (#142)

Author: wanmeihuali

Dockerfile.aws

@@ -5,5 +5,6 @@ RUN pip install --upgrade pip && \
     pip install --no-cache-dir -U taichi==1.6.0 matplotlib numpy pytorch_msssim dataclass-wizard pillow pyyaml pandas[parquet]==2.0.0 scipy argparse tensorboard
 COPY . /opt/ml/code
 WORKDIR /opt/ml/code
+RUN pip install -i https://pypi.taichi.graphics/simple/ taichi-nightly
 RUN pip install -r requirements.txt
 RUN pip install -e .

requirements.txt

@@ -1,4 +1,3 @@
-taichi>=1.6.0
 matplotlib
 numpy
 pytorch_msssim
@@ -8,4 +7,4 @@ pyyaml
 pandas[parquet]>=2.0.0
 scipy
 argparse
-tensorboard
+tensorboard

scratch/playground.py

@@ -1,4 +1,9 @@
 # %%
+from scipy.stats import multivariate_normal
+import open3d as o3d
+from sympy import symbols, Matrix, diff, exp
+import matplotlib.pyplot as plt
+import pandas as pd
 from Camera import CameraInfo
 from utils import get_ray_origin_and_direction_from_camera, get_ray_origin_and_direction_from_camera_by_gpt
 from utils import torch_single_point_alpha_forward
@@ -88,7 +93,7 @@
                                   homogeneous_translation_camera[1, 0], homogeneous_translation_camera[2, 0]])
 
 D_translation_camrea_D_q = translation_camera.jacobian(q)
-D_translation_camrea_D_t = translation_camera.jacobian(translation)
+D_translation_camrea_D_t = translation_camera.jacobian(t)
 print(latex(D_translation_camrea_D_q))
 pprint(D_translation_camrea_D_q, use_unicode=True)
 print(latex(D_translation_camrea_D_t))
@@ -303,7 +308,6 @@ def rotation_matrix_from_quaternion(q: ti.math.vec4) -> ti.math.mat3:
 print(sympy.python(J))
 
 # %%
-import sympy
 xy = sympy.MatrixSymbol('xy', 2, 1)
 mu = sympy.Matrix(["mu_x", "mu_y"])
 cov = sympy.MatrixSymbol('cov', 2, 2)
@@ -325,7 +329,6 @@ def rotation_matrix_from_quaternion(q: ti.math.vec4) -> ti.math.mat3:
 print(J.shape)
 print(sympy.python(J))
 # %%
-import numpy as np
 xy = np.array([1, 2])
 x = xy[0]
 y = xy[1]
@@ -377,7 +380,6 @@ def gradient_cov(x, mean, cov):
 print(gradient_mean(xy, mu, cov))
 print(gradient_cov(xy, mu, cov))
 # %%
-import torch
 xy = torch.tensor([1., 2.])
 mu = torch.tensor([3., 1.], requires_grad=True)
 cov = torch.tensor([[0.8, 0.1], [0.1, 0.8]], requires_grad=True)
@@ -585,14 +587,11 @@ def quaternion_to_rotation_matrix_torch(q):
     T_pointcloud_camera=T_pointcloud_camera)
 
 # %%
-import pandas as pd
-import numpy as np
 path = "logs/sigmoid_on_image_fix_bug/scene_66000.parquet"
 df = pd.read_parquet(path)
 # %%
 df.head()
 # %%
-import matplotlib.pyplot as plt
 plt.hist(np.exp(df.cov_s0), bins=100)
 # %%
 np.exp(df.cov_s0).argmax()
@@ -608,11 +607,11 @@ def quaternion_to_rotation_matrix_torch(q):
     print(col, df[col].isnull().sum())
 
 # %%
-df = df.dropna() 
+df = df.dropna()
 
 # %%
-import numpy as np
-from sympy import symbols, Matrix, diff, exp
+
+
 def compute_derivatives(mu, Sigma, x):
     # 定义符号变量
     mu1, mu2, x1, x2 = symbols('mu1 mu2 x1 x2')
@@ -637,13 +636,14 @@ def compute_derivatives(mu, Sigma, x):
     """
 
     # 用实际值替换符号变量
-    subs = {mu1: mu[0], mu2: mu[1], x1: x[0], x2: x[1], 
+    subs = {mu1: mu[0], mu2: mu[1], x1: x[0], x2: x[1],
             s11: Sigma[0, 0], s12: Sigma[0, 1], s21: Sigma[1, 0], s22: Sigma[1, 1]}
     dp_dmu_val = dp_dmu.subs(subs)
     dp_dSigma_val = dp_dSigma.subs(subs)
 
     return dp_dmu_val, dp_dSigma_val
 
+
 # 测试函数
 mu = np.array([1, 2])
 Sigma = np.array([[100., 0], [0, 100]])
@@ -653,7 +653,7 @@ def compute_derivatives(mu, Sigma, x):
 print("dp/dmu:", dp_dmu)
 print("dp/dSigma:", dp_dSigma)
 # %%
-import torch
+
 
 def compute_derivatives_torch(mu, Sigma, x):
     # 将输入转换为PyTorch张量，并设置requires_grad=True以启用自动微分
@@ -671,6 +671,7 @@ def compute_derivatives_torch(mu, Sigma, x):
 
     return mu_torch.grad, Sigma_torch.grad
 
+
 def my_compute(mu, Sigma, x):
     gaussian_mean = mu
     xy = x
@@ -680,7 +681,7 @@ def my_compute(mu, Sigma, x):
     det_cov = Sigma[0, 0] * Sigma[1, 1] - Sigma[0, 1] * Sigma[1, 0]
     inv_cov = (1. / det_cov) * \
         np.array([[gaussian_covariance[1, 1], -gaussian_covariance[0, 1]],
-                      [-gaussian_covariance[1, 0], gaussian_covariance[0, 0]]])
+                  [-gaussian_covariance[1, 0], gaussian_covariance[0, 0]]])
     cov_inv_xy_mean = inv_cov @ xy_mean
     xy_mean_T_cov_inv_xy_mean = xy_mean @ cov_inv_xy_mean
     exponent = -0.5 * xy_mean_T_cov_inv_xy_mean
@@ -689,9 +690,10 @@ def my_compute(mu, Sigma, x):
     xy_mean_outer_xy_mean = np.array([[xy_mean[0] * xy_mean[0], xy_mean[0] * xy_mean[1]],
                                       [xy_mean[1] * xy_mean[0], xy_mean[1] * xy_mean[1]]])
     d_p_d_cov = 0.5 * p * (inv_cov @
-                            xy_mean_outer_xy_mean @ inv_cov)
+                           xy_mean_outer_xy_mean @ inv_cov)
     return d_p_d_mean, d_p_d_cov
 
+
 # 测试函数
 mu = np.array([1.0, 2.0])
 Sigma = np.array([[1.0, 0.0], [0.0, 1.0]])
@@ -711,7 +713,8 @@ def my_compute(mu, Sigma, x):
     tmp = np.random.rand(2, 2)
     Sigma = tmp @ tmp.T
     dp_dmu, dp_dSigma = compute_derivatives(mu, Sigma, x)
-    dp_dmu, dp_dSigma = np.array(dp_dmu, dtype=np.float32), np.array(dp_dSigma, dtype=np.float32)
+    dp_dmu, dp_dSigma = np.array(dp_dmu, dtype=np.float32), np.array(
+        dp_dSigma, dtype=np.float32)
     dp_dmu = dp_dmu.reshape(-1)
     dp_dSigma = dp_dSigma.reshape(2, 2)
     dp_dmu_torch, dp_dSigma_torch = compute_derivatives_torch(mu, Sigma, x)
@@ -723,25 +726,24 @@ def my_compute(mu, Sigma, x):
     print("dp/dSigma:", dp_dSigma)
     print("dp/dSigma (my):", dp_dSigma_my)
     print("dp/dSigma (torch):", dp_dSigma_torch)
-    
+
     # assert np.allclose(dp_dmu, dp_dmu_torch.detach().numpy(), rtol=1e-3), f"dp_dmu: {dp_dmu}, dp_dmu_torch: {dp_dmu_torch}"
     # assert np.allclose(dp_dSigma, dp_dSigma_torch.detach().numpy(), rtol=1e-3), f"dp_dSigma: {dp_dSigma}, dp_dSigma_torch: {dp_dSigma_torch}"
-    assert np.allclose(dp_dmu, dp_dmu_my, rtol=1e-3), f"dp_dmu: {dp_dmu}, dp_dmu_my: {dp_dmu_my}"
-    assert np.allclose(dp_dSigma, dp_dSigma_my, rtol=1e-3), f"dp_dSigma: {dp_dSigma}, dp_dSigma_my: {dp_dSigma_my}"
-    
-    
+    assert np.allclose(dp_dmu, dp_dmu_my,
+                       rtol=1e-3), f"dp_dmu: {dp_dmu}, dp_dmu_my: {dp_dmu_my}"
+    assert np.allclose(dp_dSigma, dp_dSigma_my,
+                       rtol=1e-3), f"dp_dSigma: {dp_dSigma}, dp_dSigma_my: {dp_dSigma_my}"
+
+
 # %%
-import pandas as pd
-import numpy as np
-import open3d as o3d
 parquet_path = "/home/kuangyuan/hdd/Development/taichi_3d_gaussian_splatting/logs/tat_truck_experiment_more_val/scene_13750.parquet"
 df = pd.read_parquet(parquet_path)
 # %%
 df.head()
 # %%
 point_cloud = df[["x", "y", "z"]].values
 point_cloud_rgb = df[["r_sh0", "g_sh0", "b_sh0"]].values
-# here rgb are actually sh coefficients (-inf, inf), 
+# here rgb are actually sh coefficients (-inf, inf),
 # need to apply sigmoid to get (0, 1) rgb
 point_cloud_rgb = 1.0 / (1.0 + np.exp(-point_cloud_rgb))
 # %%
@@ -789,6 +791,7 @@ def rotation_matrix_from_quaternion(q: ti.math.vec4) -> ti.math.mat3:
     ])
 """
 
+
 def rotation_matrix_from_quaternion(q: np.ndarray) -> np.ndarray:
     xx = q[0] * q[0]
     yy = q[1] * q[1]
@@ -805,6 +808,7 @@ def rotation_matrix_from_quaternion(q: np.ndarray) -> np.ndarray:
         [2 * (xz - wy), 2 * (yz + wx), 1 - 2 * (xx + yy)]
     ])
 
+
 S = np.exp(s)
 
 rotated_S = np.zeros((len(q), 3))
@@ -814,8 +818,7 @@ def rotation_matrix_from_quaternion(q: np.ndarray) -> np.ndarray:
     normal[i] = rotation_matrix_from_quaternion(q[i]) @ base_vector[i]
     normal[i] *= np.linalg.norm(rotated_S[i])
 
- 
-    
+
 # %%
 point_cloud_o3d = o3d.geometry.PointCloud()
 point_cloud_o3d.points = o3d.utility.Vector3dVector(point_cloud[mask])
@@ -824,8 +827,9 @@ def rotation_matrix_from_quaternion(q: np.ndarray) -> np.ndarray:
 o3d.visualization.draw_geometries([point_cloud_o3d])
 # %%
 
-import taichi as ti
 ti.init(arch=ti.cpu)
+
+
 @ti.kernel
 def test():
     Cov = ti.Matrix([
@@ -835,11 +839,10 @@ def test():
     eig, V = ti.sym_eig(Cov)
     print(eig)
     print(V)
+
+
 test()
 # %%
-import numpy as np
-import matplotlib.pyplot as plt
-from scipy.stats import multivariate_normal
 
 # Define the mean and covariance matrix
 mean = np.array([0, 0])
@@ -875,8 +878,10 @@ def test():
 plt.imshow(mask, extent=(-50, 50, -50, 50), origin='lower')
 
 # plt eigenvectors
-plt.quiver(mean[0], mean[1], eigen_vectors[0, 0], eigen_vectors[1, 0], color='r', scale=10 / np.sqrt(eigen_values[0]))
-plt.quiver(mean[0], mean[1], eigen_vectors[0, 1], eigen_vectors[1, 1], color='r', scale=10 / np.sqrt(eigen_values[1]))
+plt.quiver(mean[0], mean[1], eigen_vectors[0, 0], eigen_vectors[1,
+           0], color='r', scale=10 / np.sqrt(eigen_values[0]))
+plt.quiver(mean[0], mean[1], eigen_vectors[0, 1], eigen_vectors[1,
+           1], color='r', scale=10 / np.sqrt(eigen_values[1]))
 
 plt.colorbar()
 plt.show()
@@ -885,13 +890,14 @@ def test():
 print(np.sqrt(eigen_values[1]) * 4)
 
 # %%
-import pandas as pd
 "/home/kuangyuan/hdd/Development/other/taichi_3d_gaussian_splatting/logs/tat_truck_every_8_experiment/camera_poses_6000.parquet"
-df = pd.read_parquet("/home/kuangyuan/hdd/Development/other/taichi_3d_gaussian_splatting/logs/tat_truck_every_8_with_pose_noise_optimization/camera_poses_10000.parquet")
+df = pd.read_parquet(
+    "/home/kuangyuan/hdd/Development/other/taichi_3d_gaussian_splatting/logs/tat_truck_every_8_with_pose_noise_optimization/camera_poses_10000.parquet")
 # %%
 df.head()
 # %%
-df1 = pd.read_parquet("/home/kuangyuan/hdd/Development/other/taichi_3d_gaussian_splatting/logs/tat_truck_every_8_baseline/camera_poses_30000.parquet")
+df1 = pd.read_parquet(
+    "/home/kuangyuan/hdd/Development/other/taichi_3d_gaussian_splatting/logs/tat_truck_every_8_baseline/camera_poses_30000.parquet")
 # %%
 
 df1.head()

taichi_3d_gaussian_splatting/GaussianPoint3D.py

@@ -211,7 +211,14 @@ def project_to_camera_position_by_q_t_jacobian(
         d_uv_d_q = d_uv_d_translation_camera @ d_translation_camera_d_q
         # d_uv_d_t = d_uv_d_translation_camera @ d_translation_camera_d_t
         d_uv_d_t = d_uv_d_translation_camera
-        return d_uv_d_translation, d_uv_d_q, d_uv_d_t
+        d_depth_dq = ti.math.vec4([
+            d_translation_camera_d_q[2, 0],
+            d_translation_camera_d_q[2, 1],
+            d_translation_camera_d_q[2, 2],
+            d_translation_camera_d_q[2, 3]
+        ])
+        d_depth_dt = ti.math.vec3([0, 0, 1])
+        return d_uv_d_translation, d_uv_d_q, d_uv_d_t, d_depth_dq, d_depth_dt
 
     @ti.func
     def project_to_camera_covariance(