Update app.py

app.py

@@ -1,49 +1,83 @@
-import gradio as gr
 import trimesh
+import gradio as gr
+import matplotlib.pyplot as plt
 import io
 import base64
 from PIL import Image
-from matplotlib import pyplot as plt
 
 def generate_cad_model(base_length, base_width, base_height):
     """
-    Generate a CAD model as a 3D mesh and return a 2D projection image for demo purposes.
+    Generate a CAD model as a 3D mesh.
     """
     # Create a 3D box (CAD model representation)
     mesh = trimesh.creation.box(extents=(base_length, base_width, base_height))
+    return mesh
+
+def render_3d_cad_model(base_length, base_width, base_height):
+    """
+    Render a 3D CAD model for interactive viewing.
+    """
+    # Generate the 3D CAD model
+    mesh = generate_cad_model(base_length, base_width, base_height)
     
-    # Create an image buffer to save a rendered view
-    image_buffer = io.BytesIO()
-    # Render the model to an image
-    mesh.show(smooth=False)
+    # Convert the mesh to GLB format for 3D visualization
+    glb_data = io.BytesIO()
+    mesh.export(glb_data, file_type="glb")
+    glb_data.seek(0)
+    return glb_data
+
+def render_2d_cad_projection(base_length, base_width, base_height):
+    """
+    Generate a 2D projection of the CAD model and return it as an image.
+    """
+    # Generate the 3D CAD model
+    mesh = generate_cad_model(base_length, base_width, base_height)
     
-    # Generate a 2D projection image (top view for demo)
+    # Create a 2D projection of the model
     projection = mesh.section(plane_origin=(0, 0, 0), plane_normal=(0, 0, 1))
-    projection_img = trimesh.util.grid_2d(projection.to_image(), tile=(1, 1))
-    plt.imshow(projection_img, cmap="gray")
-    plt.axis("off")
-    plt.savefig(image_buffer, format="png")
-    image_buffer.seek(0)
-    
-    # Convert to Base64 for Gradio to display as an image
-    encoded_image = base64.b64encode(image_buffer.getvalue()).decode("utf-8")
-    return f"data:image/png;base64,{encoded_image}"
+    if projection is not None:
+        projection_img = trimesh.util.grid_2d(projection.to_image(), tile=(1, 1))
+        
+        # Save the projection as an image
+        buffer = io.BytesIO()
+        plt.imshow(projection_img, cmap="gray")
+        plt.axis("off")
+        plt.savefig(buffer, format="png", bbox_inches="tight")
+        buffer.seek(0)
+        
+        # Convert to Base64 for Gradio display
+        encoded_image = base64.b64encode(buffer.getvalue()).decode("utf-8")
+        return f"data:image/png;base64,{encoded_image}"
+    else:
+        return None
 
 # Gradio Interface
-def display_cad_model(base_length, base_width, base_height):
-    cad_image = generate_cad_model(base_length, base_width, base_height)
-    return cad_image
+def display_cad(base_length, base_width, base_height):
+    """
+    Render both 3D and 2D CAD outputs.
+    """
+    # Render 3D model
+    glb_data = render_3d_cad_model(base_length, base_width, base_height)
+    
+    # Render 2D projection
+    cad_image = render_2d_cad_projection(base_length, base_width, base_height)
+    
+    return glb_data, cad_image
 
-# Hosting the interface with Gradio
+# Gradio interface for 3D and 2D CAD visualization
 demo = gr.Interface(
-    fn=display_cad_model,
+    fn=display_cad,
     inputs=[
         gr.Number(label="Base Length (mm)", value=10),
         gr.Number(label="Base Width (mm)", value=10),
         gr.Number(label="Base Height (mm)", value=5),
     ],
-    outputs=gr.Image(label="CAD Model Projection"),
+    outputs=[
+        gr.Model3D(label="3D CAD Model"),  # Interactive 3D viewer
+        gr.Image(label="2D CAD Projection"),  # Static 2D projection
+    ],
     live=True,
 )
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()