Added trimesh install and allow user file download

app.py

@@ -9,6 +9,7 @@ import torch
 import tempfile
 import os
 import trimesh
+import time  # Add this import at the top of the file
 
 # Run the script to download pretrained models
 subprocess.run(["bash", "get_pretrained_models.sh"])
@@ -56,7 +57,7 @@ def generate_3d_model(depth, image_path, focallength_px):
         focallength_px (float): Focal length in pixels.
 
     Returns:
-        str: Path to the exported 3D model file in OBJ format.
+        tuple: Paths to the exported 3D model files for viewing and downloading.
     """
     # Load the RGB image and convert to a NumPy array
     image = np.array(Image.open(image_path))
@@ -96,10 +97,13 @@ def generate_3d_model(depth, image_path, focallength_px):
     # Create the mesh using Trimesh with vertex colors
     mesh = trimesh.Trimesh(vertices=vertices, faces=faces, vertex_colors=colors)
 
-    # Export the mesh to an OBJ file .
-    model_path = 'output_model.obj'
-    mesh.export(model_path)
-    return model_path
+    # Export the mesh to OBJ files with unique filenames
+    timestamp = int(time.time())
+    view_model_path = f'view_model_{timestamp}.obj'
+    download_model_path = f'download_model_{timestamp}.obj'
+    mesh.export(view_model_path)
+    mesh.export(download_model_path)
+    return view_model_path, download_model_path
 
 @spaces.GPU(duration=20)
 def predict_depth(input_image):
@@ -114,7 +118,8 @@ def predict_depth(input_image):
             - str: Path to the depth map image.
             - str: Focal length in pixels or an error message.
             - str: Path to the raw depth data CSV file.
-            - str: Path to the generated 3D model file.
+            - str: Path to the generated 3D model file for viewing.
+            - str: Path to the downloadable 3D model file.
     """
     temp_file = None
     try:
@@ -176,12 +181,12 @@ def predict_depth(input_image):
         np.savetxt(raw_depth_path, depth, delimiter=',')
 
         # Generate the 3D model from the depth map and resized image
-        model_path = generate_3d_model(depth, temp_file, focallength_px)
+        view_model_path, download_model_path = generate_3d_model(depth, temp_file, focallength_px)
 
-        return output_path, f"Focal length: {focallength_px:.2f} pixels", raw_depth_path, model_path
+        return output_path, f"Focal length: {focallength_px:.2f} pixels", raw_depth_path, view_model_path, download_model_path
     except Exception as e:
         # Return error messages in case of failures
-        return None, f"An error occurred: {str(e)}", None, None
+        return None, f"An error occurred: {str(e)}", None, None, None
     finally:
         # Clean up by removing the temporary resized image file
         if temp_file and os.path.exists(temp_file):
@@ -195,7 +200,8 @@ iface = gr.Interface(
         gr.Image(type="filepath", label="Depth Map"),  # Displays the depth map image
         gr.Textbox(label="Focal Length or Error Message"),  # Shows focal length or error messages
         gr.File(label="Download Raw Depth Map (CSV)"),  # Allows downloading the raw depth data
-        gr.Model3D(label="3D Model")  # Displays the generated 3D model
+        gr.Model3D(label="View 3D Model"),  # For viewing the 3D model
+        gr.File(label="Download 3D Model (OBJ)")  # For downloading the 3D model
     ],
     title="DepthPro Demo with 3D Visualization",
     description=(
@@ -204,7 +210,8 @@ iface = gr.Interface(
         "1. Upload an image.\n"
         "2. The app will predict the depth map, display it, and provide the focal length.\n"
         "3. Download the raw depth data as a CSV file.\n"
-        "4. View the generated 3D model textured with the original image."
+        "4. View the generated 3D model textured with the original image.\n"
+        "5. Download the 3D model as an OBJ file if desired."
     ),
 )
 

requirements.txt

@@ -5,4 +5,6 @@ torch
 torchvision
 numpy
 pillow_heif
-timm
+timm
+trimesh
+time