└── README.md /README.md: -------------------------------------------------------------------------------- 1 | # GaussianFlow: Splatting Gaussian Dynamics for 4D Content Creation 2 | [ArXiv](https://arxiv.org/pdf/2403.12365.pdf) | [Project Page](https://zerg-overmind.github.io/GaussianFlow.github.io/) 3 | 4 | Please refer to [this repo](https://github.com/Zerg-Overmind/diff-gaussian-rasterization) for our cuda implementations of variables for calculating Gaussian flow. 5 | 6 | For now, we refer the code below to calculate Gaussian flow with variables from above. 7 | 8 | ```python 9 | ### We detach the variables related to t_1 in calculation of GaussianFlow such that the gradient backward 10 | ### only works for variables at t_2 while keeping variables at t_1 unchanged because 11 | ### variables at t_1 have been updated at t_1 - 1 with the same logic. 12 | 13 | ### This can accelerate the training process since less variables needed to be updated. BTW, not detach 14 | #### variables at t_1 will not decrase the performance but slow down the training. 15 | 16 | # Gaussian parameters at t_1 17 | proj_2D_t_1 = render_t_1["proj_2D"] 18 | gs_per_pixel = render_t_1["gs_per_pixel"].long() 19 | weight_per_gs_pixel = render_t_1["weight_per_gs_pixel"] 20 | x_mu = render_t_1["x_mu"] 21 | cov2D_inv_t_1 = render_t_1["conic_2D"].detach() 22 | 23 | # Gaussian parameters at t_2 24 | proj_2D_t_2 = render_t_2["proj_2D"] 25 | cov2D_inv_t_2 = render_t_2["conic_2D"] 26 | cov2D_t_2 = render_t_2["conic_2D_inv"] 27 | 28 | 29 | cov2D_t_2_mtx = torch.zeros([cov2D_t_2.shape[0], 2, 2]).cuda() 30 | cov2D_t_2_mtx[:, 0, 0] = cov2D_t_2[:, 0] 31 | cov2D_t_2_mtx[:, 0, 1] = cov2D_t_2[:, 1] 32 | cov2D_t_2_mtx[:, 1, 0] = cov2D_t_2[:, 1] 33 | cov2D_t_2_mtx[:, 1, 1] = cov2D_t_2[:, 2] 34 | 35 | cov2D_inv_t_1_mtx = torch.zeros([cov2D_inv_t_1.shape[0], 2, 2]).cuda() 36 | cov2D_inv_t_1_mtx[:, 0, 0] = cov2D_inv_t_1[:, 0] 37 | cov2D_inv_t_1_mtx[:, 0, 1] = cov2D_inv_t_1[:, 1] 38 | cov2D_inv_t_1_mtx[:, 1, 0] = cov2D_inv_t_1[:, 1] 39 | cov2D_inv_t_1_mtx[:, 1, 1] = cov2D_inv_t_1[:, 2] 40 | 41 | # B_t_2 42 | U_t_2 = torch.svd(cov2D_t_2_mtx)[0] 43 | S_t_2 = torch.svd(cov2D_t_2_mtx)[1] 44 | V_t_2 = torch.svd(cov2D_t_2_mtx)[2] 45 | B_t_2 = torch.bmm(torch.bmm(U_t_2, torch.diag_embed(S_t_2)**(1/2)), V_t_2.transpose(1,2)) 46 | 47 | # B_t_1 ^(-1) 48 | U_inv_t_1 = torch.svd(cov2D_inv_t_1_mtx)[0] 49 | S_inv_t_1 = torch.svd(cov2D_inv_t_1_mtx)[1] 50 | V_inv_t_1 = torch.svd(cov2D_inv_t_1_mtx)[2] 51 | B_inv_t_1 = torch.bmm(torch.bmm(U_inv_t_1, torch.diag_embed(S_inv_t_1)**(1/2)), V_inv_t_1.transpose(1,2)) 52 | 53 | # calculate B_t_2*B_inv_t_1 54 | B_t_2_B_inv_t_1 = torch.bmm(B_t_2, B_inv_t_1) 55 | 56 | # calculate cov2D_t_2*cov2D_inv_t_1 57 | # cov2D_t_2cov2D_inv_t_1 = torch.zeros([cov2D_inv_t_2.shape[0],2,2]).cuda() 58 | # cov2D_t_2cov2D_inv_t_1[:, 0, 0] = cov2D_t_2[:, 0] * cov2D_inv_t_1[:, 0] + cov2D_t_2[:, 1] * cov2D_inv_t_1[:, 1] 59 | # cov2D_t_2cov2D_inv_t_1[:, 0, 1] = cov2D_t_2[:, 0] * cov2D_inv_t_1[:, 1] + cov2D_t_2[:, 1] * cov2D_inv_t_1[:, 2] 60 | # cov2D_t_2cov2D_inv_t_1[:, 1, 0] = cov2D_t_2[:, 1] * cov2D_inv_t_1[:, 0] + cov2D_t_2[:, 2] * cov2D_inv_t_1[:, 1] 61 | # cov2D_t_2cov2D_inv_t_1[:, 1, 1] = cov2D_t_2[:, 1] * cov2D_inv_t_1[:, 1] + cov2D_t_2[:, 2] * cov2D_inv_t_1[:, 2] 62 | 63 | # isotropic version of GaussianFlow 64 | #predicted_flow_by_gs = (proj_2D_next[gs_per_pixel] - proj_2D[gs_per_pixel].detach()) * weights.detach() 65 | 66 | # full formulation of GaussianFlow 67 | cov_multi = (B_t_2_B_inv_t_1[gs_per_pixel] @ x_mu.permute(0,2,3,1).unsqueeze(-1).detach()).squeeze() 68 | predicted_flow_by_gs = (cov_multi + proj_2D_next[gs_per_pixel] - proj_2D[gs_per_pixel].detach() - x_mu.permute(0,2,3,1).detach()) * weights.detach() 69 | 70 | # flow supervision loss 71 | large_motion_msk = torch.norm(optical_flow, p=2, dim=-1) >= flow_thresh # flow_thresh = 0.1 or other value to filter out noise, here we assume that we have already loaded pre-computed optical flow somewhere as pseudo GT 72 | Lflow = torch.norm((optical_flow - predicted_flow_by_gs.sum(0))[large_motion_msk], p=2, dim=-1).mean() 73 | loss = loss + flow_weight * Lflow # flow_weight could be 1, 0.1, ... whatever you want. 74 | 75 | ``` 76 | --------------------------------------------------------------------------------