Update app.py

app.py

@@ -6,39 +6,49 @@ from PIL import Image
 import open3d as o3d
 from pathlib import Path
 import os
+import cv2
+from rembg import remove  # Import the rembg library for background removal
 
+# Initialize model and feature extractor for depth estimation
 feature_extractor = DPTFeatureExtractor.from_pretrained("Intel/dpt-large")
 model = DPTForDepthEstimation.from_pretrained("Intel/dpt-large")
+
 def process_image(image_path, depth_map_path=None):
     image_path = Path(image_path)
-    image_raw = Image.open(image_path)
+    image_raw = Image.open(image_path).convert("RGB")  # Ensure the image is in RGB format
     image = image_raw.resize(
         (1600, int(1600 * image_raw.size[1] / image_raw.size[0])),
         Image.Resampling.LANCZOS,
     )
 
-    if depth_map_path:
-        # Load and resize the user-provided depth map to match the RGB image size
-        depth_image_raw = Image.open(depth_map_path)
-        depth_image = depth_image_raw.resize(image.size, Image.Resampling.NEAREST)
-        depth_image = np.array(depth_image)
-
-        # Adjust the depth map based on the Depth-Anything-V2 format
-        # Normalize depth if needed (assuming values are not scaled to 0-255)
-        depth_image = (depth_image - np.min(depth_image)) / (np.max(depth_image) - np.min(depth_image))
-        depth_image = (depth_image * 255).astype('uint8')
+    # Remove background using rembg
+    foreground = remove(image_raw)  # Remove background
+    foreground = Image.fromarray(np.array(foreground))  # Convert back to PIL Image
+    foreground = foreground.convert("RGB")  # Ensure the foreground is in RGB format
 
-        # Optionally invert the depth map if it's inverse (like MiDaS)
-        # depth_image = np.max(depth_image) - depth_image
+    # Check if user-provided depth map is available
+    if depth_map_path:
+        depth_map_path = Path(depth_map_path)
+        if depth_map_path.suffix == '.npy':
+            # Load depth map from .npy file
+            depth_image = np.load(depth_map_path)
+        else:
+            # Load depth map from image file
+            depth_image_raw = Image.open(depth_map_path).convert("L")  # Convert to grayscale
+            depth_image = depth_image_raw.resize(image.size, Image.Resampling.NEAREST)
+            depth_image = np.array(depth_image)
+
+            # Normalize depth image to match expected format
+            depth_image = (depth_image - np.min(depth_image)) / (np.max(depth_image) - np.min(depth_image))
+            depth_image = (depth_image * 255).astype('uint8')
 
     else:
-        # Automatically generate the depth map using the model (if no depth map provided)
-        encoding = feature_extractor(image, return_tensors="pt")
+        # Generate depth map using DPT model
+        encoding = feature_extractor(foreground, return_tensors="pt")
         with torch.no_grad():
             outputs = model(**encoding)
             predicted_depth = outputs.predicted_depth
 
-        # Interpolate to original size
         prediction = torch.nn.functional.interpolate(
             predicted_depth.unsqueeze(1),
             size=image.size[::-1],
@@ -47,13 +57,17 @@ def process_image(image_path, depth_map_path=None):
         ).squeeze()
         depth_image = (prediction.cpu().numpy() * 255 / np.max(prediction.cpu().numpy())).astype("uint8")
 
+    # Step 1: Apply Gaussian smoothing on the depth map
+    smoothed_depth_map = cv2.GaussianBlur(depth_image, (5, 5), 0)
+
     try:
-        gltf_path = create_3d_obj(np.array(image), depth_image, image_path)
-        img = Image.fromarray(depth_image)
+        gltf_path = create_3d_obj(np.array(image), smoothed_depth_map, image_path)
+        img = Image.fromarray(smoothed_depth_map)
         return [img, gltf_path, gltf_path]
     except Exception as e:
-        gltf_path = create_3d_obj(np.array(image), depth_image, image_path, depth=8)
-        img = Image.fromarray(depth_image)
+        print("Error with default depth. Retrying with a shallower depth.")
+        gltf_path = create_3d_obj(np.array(image), smoothed_depth_map, image_path, depth=8)
+        img = Image.fromarray(smoothed_depth_map)
         return [img, gltf_path, gltf_path]
     except:
         print("Error reconstructing 3D model")
@@ -65,60 +79,51 @@ def create_3d_obj(rgb_image, depth_image, image_path, depth=12):
     rgbd_image = o3d.geometry.RGBDImage.create_from_color_and_depth(
         image_o3d, depth_o3d, convert_rgb_to_intensity=False
     )
-    w = int(depth_image.shape[1])
-    h = int(depth_image.shape[0])
+    w, h = depth_image.shape[1], depth_image.shape[0]
 
+    # Camera intrinsic setup for 3D point cloud
     camera_intrinsic = o3d.camera.PinholeCameraIntrinsic()
     camera_intrinsic.set_intrinsics(w, h, 1000, 1000, w / 2, h / 2)
 
+    # Point Cloud and Normals
     pcd = o3d.geometry.PointCloud.create_from_rgbd_image(rgbd_image, camera_intrinsic)
-
-    print("normals")
-    pcd.normals = o3d.utility.Vector3dVector(
-        np.zeros((1, 3))
-    )  # invalidate existing normals
-    pcd.estimate_normals(
-        search_param=o3d.geometry.KDTreeSearchParamHybrid(radius=0.05, max_nn=30)
-    )
-    pcd.orient_normals_towards_camera_location(
-        camera_location=np.array([0.0, 0.0, 1000.0])
-    )
+    pcd.normals = o3d.utility.Vector3dVector(np.zeros((1, 3)))
+    pcd.estimate_normals(search_param=o3d.geometry.KDTreeSearchParamHybrid(radius=0.05, max_nn=30))
+    pcd.orient_normals_towards_camera_location(camera_location=np.array([0.0, 0.0, 1000.0]))
     pcd.transform([[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]])
     pcd.transform([[-1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]])
 
-    print("run Poisson surface reconstruction")
-    with o3d.utility.VerbosityContextManager(o3d.utility.VerbosityLevel.Debug) as cm:
-        mesh_raw, densities = o3d.geometry.TriangleMesh.create_from_point_cloud_poisson(
-            pcd, depth=depth, width=0, scale=1.1, linear_fit=True
-        )
+    # Poisson Surface Reconstruction
+    mesh_raw, densities = o3d.geometry.TriangleMesh.create_from_point_cloud_poisson(
+        pcd, depth=depth, width=0, scale=1.1, linear_fit=True
+    )
 
+    # Step 3: Mesh simplification and smoothing
     voxel_size = max(mesh_raw.get_max_bound() - mesh_raw.get_min_bound()) / 512
-    print(f"voxel_size = {voxel_size:e}")
     mesh = mesh_raw.simplify_vertex_clustering(
         voxel_size=voxel_size,
         contraction=o3d.geometry.SimplificationContraction.Average,
     )
-
-    # Apply smoothing
     mesh = mesh.filter_smooth_simple(number_of_iterations=3)
-
     bbox = pcd.get_axis_aligned_bounding_box()
     mesh_crop = mesh.crop(bbox)
+
+    # Save GLTF
     gltf_path = f"./{image_path.stem}.gltf"
     o3d.io.write_triangle_mesh(gltf_path, mesh_crop, write_triangle_uvs=True)
     return gltf_path
 
+# Gradio Interface
+title = "Depth Estimation & 3D Reconstruction Demo"
+description = "Upload an image and optionally a depth map (in .npy or image format) to generate a 3D model. If no depth map is provided, the DPT model will generate it."
 
-
-title = "Demo: Depth Estimation & 3D Reconstruction with DPT + Point Cloud"
-description = "This demo allows users to provide an RGB image and optionally a depth map to create a 3D object. If no depth map is provided, the DPT model will generate it."
 examples = [["examples/" + img] for img in os.listdir("examples/")]
 
 iface = gr.Interface(
     fn=process_image,
     inputs=[
         gr.Image(type="filepath", label="Input Image"),
-        gr.Image(type="filepath", label="Input Depth Map (optional)"),
+        gr.File(type="filepath", label="Input Depth Map (optional)"),  # Changed input type to allow .npy
     ],
     outputs=[
         gr.Image(label="Predicted Depth", type="pil"),
@@ -131,4 +136,4 @@ iface = gr.Interface(
     allow_flagging="never",
     cache_examples=False,
 )
-iface.launch(debug=True, show_api=True, share=True)
+iface.launch(debug=True, show_api=True, share=True)