add open3d and let user set parameters for 3D model

app.py

@@ -54,7 +54,7 @@ def resize_image(image_path, max_size=1024):
             img.save(temp_file, format="PNG")
             return temp_file.name
 
-def generate_3d_model(depth, image_path, focallength_px):
+def generate_3d_model(depth, image_path, focallength_px, simplification_factor=0.8, smoothing_iterations=1, thin_threshold=0.01):
     """
     Generate a textured 3D mesh from the depth map and the original image.
     """
@@ -112,7 +112,7 @@ def generate_3d_model(depth, image_path, focallength_px):
     print("Original mesh - vertices: {}, faces: {}".format(len(mesh.vertices), len(mesh.faces)))
 
     # 1. Mesh simplification
-    target_faces = int(len(mesh.faces) * 0.8)  # Simplify to 80% of original faces
+    target_faces = int(len(mesh.faces) * simplification_factor)
     mesh = mesh.simplify_quadric_decimation(face_count=target_faces)
     print("After simplification - vertices: {}, faces: {}".format(len(mesh.vertices), len(mesh.faces)))
 
@@ -124,11 +124,12 @@ def generate_3d_model(depth, image_path, focallength_px):
         print("After removing small components - vertices: {}, faces: {}".format(len(mesh.vertices), len(mesh.faces)))
 
     # 3. Smooth the mesh
-    mesh = mesh.smoothed()
+    for _ in range(smoothing_iterations):
+        mesh = mesh.smoothed()
     print("After smoothing - vertices: {}, faces: {}".format(len(mesh.vertices), len(mesh.faces)))
 
     # 4. Remove thin features
-    mesh = remove_thin_features(mesh)
+    mesh = remove_thin_features(mesh, thickness_threshold=thin_threshold)
     print("After removing thin features - vertices: {}, faces: {}".format(len(mesh.vertices), len(mesh.faces)))
 
     # Export the mesh to OBJ files with unique filenames
@@ -163,6 +164,18 @@ def remove_thin_features(mesh, thickness_threshold=0.01):
     
     return mesh
 
+def regenerate_3d_model(depth_csv, image_path, focallength_px, simplification_factor, smoothing_iterations, thin_threshold):
+    # Load depth from CSV
+    depth = np.loadtxt(depth_csv, delimiter=',')
+    
+    # Generate new 3D model with updated parameters
+    view_model_path, download_model_path = generate_3d_model(
+        depth, image_path, focallength_px, 
+        simplification_factor, smoothing_iterations, thin_threshold
+    )
+    
+    return view_model_path, download_model_path
+
 @spaces.GPU(duration=20)
 def predict_depth(input_image):
     temp_file = None
@@ -222,13 +235,13 @@ def predict_depth(input_image):
         # Generate the 3D model from the depth map and resized image
         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, view_model_path, download_model_path
+        return output_path, f"Focal length: {focallength_px:.2f} pixels", raw_depth_path, view_model_path, download_model_path, temp_file, focallength_px
     except Exception as e:
         # Return error messages in case of failures
         import traceback
         error_message = f"An error occurred: {str(e)}\n\nTraceback:\n{traceback.format_exc()}"
         print(error_message)  # Print the full error message to the console
-        return None, error_message, None, None, None
+        return None, error_message, None, None, None, None, None
     finally:
         # Clean up by removing the temporary resized image file
         if temp_file and os.path.exists(temp_file):
@@ -245,18 +258,9 @@ def get_last_commit_timestamp():
 # Create the Gradio interface with appropriate input and output components. 
 last_updated = get_last_commit_timestamp()
 
-iface = gr.Interface(
-    fn=predict_depth,
-    inputs=gr.Image(type="filepath"),
-    outputs=[
-        gr.Image(type="filepath", label="Depth Map"),
-        gr.Textbox(label="Focal Length or Error Message"),
-        gr.File(label="Download Raw Depth Map (CSV)"),
-        gr.Model3D(label="View 3D Model"),
-        gr.File(label="Download 3D Model (OBJ)")
-    ],
-    title="DepthPro Demo with 3D Visualization",
-    description=(
+with gr.Blocks() as iface:
+    gr.Markdown("# DepthPro Demo with 3D Visualization")
+    gr.Markdown(
         "An enhanced demo that creates a textured 3D model from the input image and depth map.\n\n"
         "Forked from https://huggingface.co/spaces/akhaliq/depth-pro and model from https://huggingface.co/apple/DepthPro\n"
         "**Instructions:**\n"
@@ -264,10 +268,45 @@ iface = gr.Interface(
         "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.\n"
-        "5. Download the 3D model as an OBJ file if desired.\n\n"
+        "5. Adjust parameters and click 'Regenerate 3D Model' to update the model.\n"
+        "6. Download the 3D model as an OBJ file if desired.\n\n"
         f"Last updated: {last_updated}"
-    ),
-)
+    )
+    
+    with gr.Row():
+        input_image = gr.Image(type="filepath", label="Input Image")
+        depth_map = gr.Image(type="filepath", label="Depth Map")
+    
+    focal_length = gr.Textbox(label="Focal Length")
+    raw_depth_csv = gr.File(label="Download Raw Depth Map (CSV)")
+    
+    with gr.Row():
+        view_3d_model = gr.Model3D(label="View 3D Model")
+        download_3d_model = gr.File(label="Download 3D Model (OBJ)")
+    
+    with gr.Row():
+        simplification_factor = gr.Slider(minimum=0.1, maximum=1.0, value=0.8, step=0.1, label="Simplification Factor")
+        smoothing_iterations = gr.Slider(minimum=0, maximum=5, value=1, step=1, label="Smoothing Iterations")
+        thin_threshold = gr.Slider(minimum=0.001, maximum=0.1, value=0.01, step=0.001, label="Thin Feature Threshold")
+    
+    regenerate_button = gr.Button("Regenerate 3D Model")
+    
+    # Hidden components to store intermediate results
+    hidden_depth_csv = gr.State()
+    hidden_image_path = gr.State()
+    hidden_focal_length = gr.State()
+    
+    input_image.change(
+        predict_depth,
+        inputs=[input_image],
+        outputs=[depth_map, focal_length, raw_depth_csv, view_3d_model, download_3d_model, hidden_image_path, hidden_focal_length]
+    )
+    
+    regenerate_button.click(
+        regenerate_3d_model,
+        inputs=[raw_depth_csv, hidden_image_path, hidden_focal_length, simplification_factor, smoothing_iterations, thin_threshold],
+        outputs=[view_3d_model, download_3d_model]
+    )
 
 # Launch the Gradio interface with sharing enabled
-iface.launch(share=True)  # share=True allows you to share the interface with others.
+iface.launch(share=True)

requirements.txt

@@ -7,4 +7,5 @@ pillow_heif==0.8.0
 timm
 trimesh==3.22.1
 gradio
-opencv-python==4.5.5.64
+opencv-python==4.5.5.64
+open3d