Implement 3D mesh visualization with trimesh and marching_cubes

app.py

@@ -26,6 +26,8 @@ from monai.transforms import (
 )
 from monai.inferers import sliding_window_inference
 from labels import LABEL_NAMES, get_color_map, get_label_name, get_organ_categories
+import trimesh
+from skimage import measure
 
 # Constants
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -116,6 +118,9 @@ def run_inference(image_path: str, progress=gr.Progress()):
     postprocess = get_postprocessing_transforms()
     
     # Load and preprocess
+    image_nib = nib.load(image_path)
+    original_data = image_nib.get_fdata() # Keep original data for visualization
+    
     image = preprocess(image_path)
     image = image.unsqueeze(0).to(DEVICE)  # Add batch dimension
     
@@ -133,17 +138,41 @@ def run_inference(image_path: str, progress=gr.Progress()):
     
     progress(0.8, desc="Post-processing...")
     
-    # Post-process
-    outputs = postprocess(outputs)
-    segmentation = outputs.squeeze().cpu().numpy().astype(np.uint8)
-    
-    # Load original image for visualization
-    original_nib = nib.load(image_path)
-    original_data = original_nib.get_fdata()
+    # Post-processing
+    seg_data = postprocess(outputs).squeeze().cpu().numpy().astype(np.uint8)
     
     progress(1.0, desc="Complete!")
     
-    return original_data, segmentation
+    return original_data, seg_data
+
+
+def generate_3d_mesh(seg_data, step_size=2):
+    """Generate a 3D mesh from segmentation data using Marching Cubes"""
+    if seg_data is None or np.max(seg_data) == 0:
+        return None
+        
+    try:
+        # Create a boolean mask of all structures (excluding background 0)
+        # Using a step_size > 1 reduces resolution but speeds up generation significantly
+        # This is crucial for CPU performance on Hugging Face Spaces
+        mask = seg_data > 0
+        
+        # Marching cubes to get vertices and faces
+        # level=0.5 because boolean mask is 0 or 1
+        verts, faces, normals, values = measure.marching_cubes(mask, level=0.5, step_size=step_size)
+        
+        # Create trimesh object
+        mesh = trimesh.Trimesh(vertices=verts, faces=faces, vertex_normals=normals)
+        
+        # Export to a temporary GLB file (efficient binary format)
+        temp_file = tempfile.NamedTemporaryFile(suffix=".glb", delete=False)
+        mesh.export(temp_file.name)
+        temp_file.close()
+        
+        return temp_file.name
+    except Exception as e:
+        print(f"Error generating 3D mesh: {e}")
+        return None
 
 
 def create_slice_visualization(ct_data, seg_data, axis, slice_idx, alpha=0.5, show_overlay=True):
@@ -231,15 +260,19 @@ def process_upload(file_path, progress=gr.Progress()):
         fig = create_slice_visualization(ct_data, seg_data, "Axial", mid_axial)
         structures = get_detected_structures(seg_data)
         
+        # Generate 3D mesh (this might take a few seconds)
+        mesh_path = generate_3d_mesh(seg_data)
+        
         return (
             fig,
             structures,
+            mesh_path,
             gr.update(maximum=ct_data.shape[2] - 1, value=mid_axial),
             gr.update(maximum=ct_data.shape[1] - 1, value=mid_coronal),
             gr.update(maximum=ct_data.shape[0] - 1, value=mid_sagittal),
         )
     except Exception as e:
-        return None, f"Error processing file: {str(e)}", gr.update(), gr.update(), gr.update()
+        return None, f"Error processing file: {str(e)}", None, gr.update(), gr.update(), gr.update()
 
 
 def update_visualization(axis, slice_idx, alpha, show_overlay):
@@ -342,6 +375,14 @@ with gr.Blocks(
                 lines=10,
                 max_lines=20
             )
+            
+            # 3D Model Output
+            gr.Markdown("### 🧊 3D View")
+            model_3d_output = gr.Model3D(
+                label="3D Segmentation Mesh",
+                clear_color=[0.0, 0.0, 0.0, 0.0],
+                camera_position=(90, 90, 3) 
+            )
     
     # Model info section
     with gr.Accordion("ℹ️ Model Information", open=False):
@@ -372,7 +413,7 @@ with gr.Blocks(
     process_btn.click(
         fn=process_upload,
         inputs=[file_input],
-        outputs=[output_image, structures_output, axial_slider, coronal_slider, sagittal_slider]
+        outputs=[output_image, structures_output, model_3d_output, axial_slider, coronal_slider, sagittal_slider]
     )
     
     # Update visualization when controls change

requirements.txt

@@ -7,3 +7,4 @@ nibabel>=5.0.0
 numpy
 matplotlib
 scikit-image
+trimesh