Update app.py

app.py

@@ -4,6 +4,8 @@ from PIL import Image
 import torch
 import requests
 from io import BytesIO
+import trimesh
+import plotly.graph_objects as go
 
 # Load model and processor from Hugging Face
 def load_model_and_processor():
@@ -36,38 +38,42 @@ def image_to_3d(image):
         return "Model or processor not loaded."
     
     try:
-        # Preprocess the input image
+        # Preprocess input image
         inputs = processor(images=image, return_tensors="pt")
         
         # Run inference
         with torch.no_grad():
             outputs = model(**inputs)
         
-        # Placeholder return, replace this with actual 3D visualization logic
-        return "3D model generated from input image!"  
+        # Convert outputs to a 3D mesh (replace with actual logic based on model output)
+        # Assuming 'vertices' and 'faces' are returned by the model (adjust as needed)
+        vertices = outputs['vertices'].numpy()  # Placeholder for vertex output
+        faces = outputs['faces'].numpy()  # Placeholder for face output
+        
+        # Create a mesh using trimesh
+        mesh = trimesh.Trimesh(vertices=vertices, faces=faces)
+        
+        # Visualize the mesh using Plotly
+        fig = go.Figure(data=[go.Mesh3d(x=vertices[:,0], y=vertices[:,1], z=vertices[:,2], 
+                                         i=faces[:,0], j=faces[:,1], k=faces[:,2])])
+        
+        return fig  # return the figure for display
     except Exception as e:
         return f"Error during inference: {str(e)}"
 
 # Load the example image for the Gradio interface
 example_image = load_example_image()
 
-# Description of image types to upload
-image_type_description = """
-Upload a clear image of a single object with minimal background distractions for best results. Example image types:
-- Objects such as cars, furniture, geometric shapes, or architectural structures.
-"""
-
 # Gradio interface setup
 interface = gr.Interface(
     fn=image_to_3d,
     inputs=gr.Image(type="pil", label="Upload an Image"),
-    outputs="text",  # Placeholder output (replace with 3D rendering if needed)
+    outputs=gr.Plot(label="3D Model"),
     title="OpenLRM Mix-Large 1.1 - Image to 3D",
     description="Upload an image to generate a 3D model using OpenLRM Mix-Large 1.1.",
-    examples=[[example_image]] if example_image else None,  # Include the example image if loaded
-    theme="compact",  # Compact layout for better UI experience
-    allow_flagging="never"
+    examples=[[example_image]] if example_image else None,
+    theme="compact"
 )
 
 # Display a suggestion below the upload widget
-interface.launch(description=image_type_description)
+interface.launch()