Remove halo and add rotation

utils/app_utils.py

@@ -8,7 +8,8 @@ import os
 import torch
 from .camera_utils import get_loop_cameras
 from .graphics_utils import getProjectionMatrix
-from .general_utils import matrix_to_quaternion
+from .general_utils import matrix_to_quaternion, quaternion_raw_multiply
+import math
 
 def remove_background(image, rembg_session):
     do_remove = True
@@ -158,14 +159,44 @@ def export_to_obj(reconstruction, ply_out_path):
         assert v.shape[0] == 1, "Expected batch size to be 0"
         reconstruction[k] = v[0]
 
-    xyz = reconstruction["xyz"].detach().cpu().numpy()
+   
+    valid_gaussians = torch.where(reconstruction["opacity"] > -2.5)[0]
+
+    # transforms for visualisation in Gradio
+    # ============= Transform locations =============
+    xyz = reconstruction["xyz"][valid_gaussians].detach().cpu().clone()
+    t1 = torch.tensor([[1,  0, 0],
+                        [0,  0, 1],
+                        [0, -1, 0]], dtype=torch.float32)
+    angle1 = 30 * math.pi * 2 / 360
+    t2 = torch.tensor([[math.cos(angle1),  -math.sin(angle1), 0],
+                       [math.sin(angle1),  math.cos(angle1), 0],
+                       [0, 0, 1]], dtype=torch.float32)
+    angle2 = -60 * math.pi * 2 / 360
+    t3 = torch.tensor([[math.cos(angle2), 0, math.sin(angle2)],
+                       [0,  1, 0],
+                       [-math.sin(angle2), 0, math.cos(angle2)]], dtype=torch.float32)
+
+    overall_transform_matrix = (t1 @ t2)@ t3
+
+    xyz = torch.matmul(xyz, overall_transform_matrix).numpy()
     normals = np.zeros_like(xyz)
-    f_dc = reconstruction["features_dc"].detach().transpose(1, 2).flatten(start_dim=1).contiguous().cpu().numpy()
-    f_rest = reconstruction["features_rest"].detach().transpose(1, 2).flatten(start_dim=1).contiguous().cpu().numpy()
-    opacities = reconstruction["opacity"].detach().cpu().numpy()
+
+    # ============= Transform rotations =============
+    camera_transformation_matrix = overall_transform_matrix.inverse()
+    camera_quaternions = matrix_to_quaternion(camera_transformation_matrix).to(reconstruction["rotation"].device)
+    rotation = reconstruction["rotation"].clone()
+    rotation = rotation.unsqueeze(0)
+    rotation = quaternion_raw_multiply(camera_quaternions.unsqueeze(0).unsqueeze(0).expand(*rotation.shape), 
+                rotation).squeeze(0)
+
+    f_dc = reconstruction["features_dc"][valid_gaussians].detach().transpose(1, 2).flatten(start_dim=1).contiguous().cpu().numpy()
+    f_rest = reconstruction["features_rest"][valid_gaussians].detach().transpose(1, 2).flatten(start_dim=1).contiguous().cpu().numpy()
+    opacities = reconstruction["opacity"][valid_gaussians].detach().contiguous().cpu().numpy()
+
     # enlarge Gaussians - otherwise transforming them to .ply results in artefacts
-    scale = reconstruction["scaling"].detach().cpu().numpy() + np.abs(reconstruction["scaling"].detach().cpu().numpy() * 0.1)
-    rotation = reconstruction["rotation"].detach().cpu().numpy()
+    scale = (reconstruction["scaling"][valid_gaussians] + torch.abs(reconstruction["scaling"][valid_gaussians] * 0.1)).detach().contiguous().cpu().numpy()
+    rotation = rotation[valid_gaussians].detach().contiguous().cpu().numpy()
 
     dtype_full = [(attribute, 'f4') for attribute in construct_list_of_attributes()]