Update LHM/models/rendering/gs_renderer.py

LHM/models/rendering/gs_renderer.py

@@ -210,19 +210,79 @@ class GaussianModel:
             l.append("rot_{}".format(i))
         return l
 
+    # def save_ply(self, path: str):
+
+    #     xyz = self.xyz.detach().cpu().numpy()
+    #     normals = np.zeros_like(xyz)
+
+    #     if self.use_rgb:
+    #         shs = RGB2SH(self.shs)
+    #     else:
+    #         shs = self.shs
+
+    #     features_dc = shs[:, :1]
+    #     features_rest = shs[:, 1:]
+
+    #     f_dc = (
+    #         features_dc.float().detach().flatten(start_dim=1).contiguous().cpu().numpy()
+    #     )
+    #     f_rest = (
+    #         features_rest.float()
+    #         .detach()
+    #         .flatten(start_dim=1)
+    #         .contiguous()
+    #         .cpu()
+    #         .numpy()
+    #     )
+    #     opacities = (
+    #         inverse_sigmoid(torch.clamp(self.opacity, 1e-3, 1 - 1e-3))
+    #         .detach()
+    #         .cpu()
+    #         .numpy()
+    #     )
+
+    #     scale = np.log(self.scaling.detach().cpu().numpy())
+    #     rotation = self.rotation.detach().cpu().numpy()
+
+    #     dtype_full = [
+    #         (attribute, "f4") for attribute in self.construct_list_of_attributes()
+    #     ]
+
+    #     elements = np.empty(xyz.shape[0], dtype=dtype_full)
+    #     attributes = np.concatenate(
+    #         (xyz, normals, f_dc, f_rest, opacities, scale, rotation), axis=1
+    #     )
+    #     elements[:] = list(map(tuple, attributes))
+    #     el = PlyElement.describe(elements, "vertex")
+    #     PlyData([el]).write(path)
     def save_ply(self, path: str):
-
+        """
+        Save the Gaussian Splatting model as a PLY file in a location accessible for download.
+        
+        Args:
+            path (str): Path where to save the PLY file
+            
+        Returns:
+            str: Path to the saved PLY file
+        """
+        import os
+        
+        # Ensure output directory exists
+        output_dir = os.path.dirname(path)
+        os.makedirs(output_dir, exist_ok=True)
+        
+        # Save original PLY data
         xyz = self.xyz.detach().cpu().numpy()
         normals = np.zeros_like(xyz)
-
+    
         if self.use_rgb:
             shs = RGB2SH(self.shs)
         else:
             shs = self.shs
-
+    
         features_dc = shs[:, :1]
         features_rest = shs[:, 1:]
-
+    
         f_dc = (
             features_dc.float().detach().flatten(start_dim=1).contiguous().cpu().numpy()
         )
@@ -240,14 +300,14 @@ class GaussianModel:
             .cpu()
             .numpy()
         )
-
+    
         scale = np.log(self.scaling.detach().cpu().numpy())
         rotation = self.rotation.detach().cpu().numpy()
-
+    
         dtype_full = [
             (attribute, "f4") for attribute in self.construct_list_of_attributes()
         ]
-
+    
         elements = np.empty(xyz.shape[0], dtype=dtype_full)
         attributes = np.concatenate(
             (xyz, normals, f_dc, f_rest, opacities, scale, rotation), axis=1
@@ -255,6 +315,10 @@ class GaussianModel:
         elements[:] = list(map(tuple, attributes))
         el = PlyElement.describe(elements, "vertex")
         PlyData([el]).write(path)
+        
+        print(f"Model saved to {path}")
+        return path
+
 
     def load_ply(self, path):