Add pcd2grid

app.py

@@ -13,6 +13,7 @@ import mcc_model
 import util.misc as misc
 from engine_mcc import prepare_data
 from plyfile import PlyData, PlyElement
+import trimesh
 
 def run_inference(model, samples, device, temperature, args):
     model.eval()
@@ -122,6 +123,25 @@ def normalize(seen_xyz):
     seen_xyz = seen_xyz - seen_xyz[torch.isfinite(seen_xyz.sum(dim=-1))].mean(axis=0)
     return seen_xyz
 
+def voxel_grid_downsample(points, colors, voxel_size):
+    # Compute voxel indices
+    voxel_indices = np.floor(points / voxel_size).astype(int)
+
+    # Remove duplicate voxel indices
+    unique_voxel_indices, inverse_indices = np.unique(voxel_indices, axis=0, return_inverse=True)
+
+    # Compute the centroid of the points and the average color in each voxel
+    centroids = np.empty_like(unique_voxel_indices, dtype=float)
+    avg_colors = np.empty((len(unique_voxel_indices), colors.shape[1]), dtype=colors.dtype)
+    for i in range(len(unique_voxel_indices)):
+        centroids[i] = points[inverse_indices == i].mean(axis=0)
+        avg_colors[i] = colors[inverse_indices == i].mean(axis=0)
+
+    # Convert colors from RGB to BGR
+    avg_colors = avg_colors[:, ::-1]
+    
+    return centroids, avg_colors
+
 def infer(
           image,
           depth_image,
@@ -203,10 +223,33 @@ def infer(
         PlyData([element], text=True).write(f)
         temp_file_name = f.name
 
-    return temp_file_name
+    # Perform voxel grid downsampling
+    voxel_size = 0.2  # Change this to the size of your cubes
+    downsampled_xyz, downsampled_colors = voxel_grid_downsample(unseen_xyz, pred_colors, voxel_size)
+    
+    meshes = []
+    for point, color in zip(downsampled_xyz, downsampled_colors):
+        # Create a cube mesh at the given point
+        cube = trimesh.creation.box(extents=[voxel_size]*3)
+        cube.apply_translation(point)
+
+        # Assign the average color to the vertices
+        cube.visual.vertex_colors = np.hstack([color, 255])  # Set alpha to 255
+        meshes.append(cube)
+
+    # Save point cloud data to a temporary file
+    with tempfile.NamedTemporaryFile(suffix=".obj", delete=False) as f:
+        temp_obj_file = f.name
+        print(temp_obj_file)
+        # Combine all the cubes into a single mesh
+        combined = trimesh.util.concatenate(meshes)
+        # Save the combined mesh to a file
+        combined.export(temp_obj_file)
+    return temp_file_name, temp_obj_file
 
 if __name__ == '__main__':
-    device = "cuda" if torch.cuda.is_available() else "cpu"
+    device = "cpu"
+    # device = "cuda" if torch.cuda.is_available() else "cpu"
     
     parser = main_mcc.get_args_parser()
     parser.set_defaults(eval=True)
@@ -231,7 +274,8 @@ if __name__ == '__main__':
                                 gr.Slider(minimum=0.05, maximum=0.5, step=0.05, value=0.2, label="Grain Size"),
                                 gr.Slider(minimum=0, maximum=1.0, step=0.1, value=0.1, label="Color Temperature")
                                 ], 
-                        outputs=[gr.outputs.File(label="Point Cloud")],
+                        outputs=[gr.outputs.File(label="Point Cloud"),
+                                 gr.Model3D( clear_color=[0.0, 0.0, 0.0, 0.0], label="3D Model")],
                         examples=[["demo/quest2.jpg", "demo/quest2_depth.png", "demo/quest2_seg.png", 0.2, 0.1]],
                         cache_examples=True)
     demo.launch(server_name="0.0.0.0", server_port=7860)