Surn/DPTDepth3D

1import os2from pathlib import Path3 4import gradio as gr5import numpy as np6import open3d as o3d7import torch8from PIL import Image9from transformers import DPTForDepthEstimation, DPTImageProcessor10 11# Initialize the image processor and depth estimation model12image_processor = DPTImageProcessor.from_pretrained("Intel/dpt-large")13model = DPTForDepthEstimation.from_pretrained("Intel/dpt-large")14 15 16def process_image(image_path, resized_width=800, z_scale=208):17    """18    Processes the input image to generate a depth map and a 3D mesh reconstruction.19 20    Args:21        image_path (str): The file path to the input image.22 23    Returns:24        list: A list containing the depth image, 3D mesh reconstruction, and GLTF file path.25    """26    image_path = Path(image_path)27    if not image_path.exists():28        raise ValueError("Image file not found")29 30    # Load and resize the image31    image_raw = Image.open(image_path).convert("RGB")32    print(f"Original size: {image_raw.size}")33    resized_height = int(resized_width * image_raw.size[1] / image_raw.size[0])34    image = image_raw.resize((resized_width, resized_height), Image.Resampling.LANCZOS)35    print(f"Resized size: {image.size}")36 37    # Prepare image for the model38    encoding = image_processor(image, return_tensors="pt")39 40    # Perform depth estimation41    with torch.no_grad():42        outputs = model(**encoding)43        predicted_depth = outputs.predicted_depth44 45    # Interpolate depth to match the image size46    prediction = torch.nn.functional.interpolate(47        predicted_depth.unsqueeze(1),48        size=(image.height, image.width),49        mode="bicubic",50        align_corners=True,51    ).squeeze()52 53    # Normalize the depth image to 8-bit54    prediction = prediction.cpu().numpy()55    depth_min, depth_max = prediction.min(), prediction.max()56    depth_image = ((prediction - depth_min) / (depth_max - depth_min) * 255).astype("uint8")57 58    try:59        gltf_path = create_3d_obj(np.array(image), prediction, image_path, depth=10, z_scale=z_scale)60    except Exception:61        gltf_path = create_3d_obj(np.array(image), prediction, image_path, depth=8, z_scale=z_scale)62 63    img = Image.fromarray(depth_image)64    return [img, gltf_path, gltf_path]65 66 67def create_3d_obj(rgb_image, raw_depth, image_path, depth=10, z_scale=200):68    """69    Creates a 3D object from RGB and depth images.70 71    Args:72        rgb_image (np.ndarray): The RGB image as a NumPy array.73        raw_depth (np.ndarray): The raw depth data.74        image_path (Path): The path to the original image.75        depth (int, optional): Depth parameter for Poisson reconstruction. Defaults to 10.76        z_scale (float, optional): Scaling factor for the Z-axis. Defaults to 200.77 78    Returns:79        str: The file path to the saved GLTF model.80    """81    # Normalize the depth image82    depth_image = ((raw_depth - raw_depth.min()) / (raw_depth.max() - raw_depth.min()) * 255).astype("uint8")83    depth_o3d = o3d.geometry.Image(depth_image)84    image_o3d = o3d.geometry.Image(rgb_image)85 86    # Create RGBD image87    rgbd_image = o3d.geometry.RGBDImage.create_from_color_and_depth(88        image_o3d, depth_o3d, convert_rgb_to_intensity=False89    )90 91    height, width = depth_image.shape92 93    # Define camera intrinsics94    camera_intrinsic = o3d.camera.PinholeCameraIntrinsic(95        width,96        height,97        fx=1.0,98        fy=1.0,99        cx=width / 2.0,100        cy=height / 2.0,101    )102 103    # Generate point cloud from RGBD image104    pcd = o3d.geometry.PointCloud.create_from_rgbd_image(rgbd_image, camera_intrinsic)105 106    # Scale the Z dimension107    points = np.asarray(pcd.points)108    depth_scaled = ((raw_depth - raw_depth.min()) / (raw_depth.max() - raw_depth.min())) * z_scale109    z_values = depth_scaled.flatten()[:len(points)]110    points[:, 2] *= z_values111    pcd.points = o3d.utility.Vector3dVector(points)112 113    # Estimate and orient normals114    pcd.estimate_normals(115        search_param=o3d.geometry.KDTreeSearchParamHybrid(radius=0.01, max_nn=30)116    )117    pcd.orient_normals_towards_camera_location(camera_location=np.array([0.0, 0.0, 2.0 ]))118 119    # Apply transformations120    pcd.transform([[1, 0, 0, 0],121                   [0, -1, 0, 0],122                   [0, 0, -1, 0],123                   [0, 0, 0, 1]])124    pcd.transform([[-1, 0, 0, 0],125                   [0, 1, 0, 0],126                   [0, 0, 1, 0],127                   [0, 0, 0, 1]])128 129    # Perform Poisson surface reconstruction130    print(f"Running Poisson surface reconstruction with depth {depth}")131    mesh_raw, densities = o3d.geometry.TriangleMesh.create_from_point_cloud_poisson(132        pcd, depth=depth, width=0, scale=1.1, linear_fit=True133    )134    print(f"Raw mesh vertices: {len(mesh_raw.vertices)}, triangles: {len(mesh_raw.triangles)}")135 136    # Simplify the mesh using vertex clustering137    voxel_size = max(mesh_raw.get_max_bound() - mesh_raw.get_min_bound()) / (max(width, height) * 0.8)138    mesh = mesh_raw.simplify_vertex_clustering(139        voxel_size=voxel_size,140        contraction=o3d.geometry.SimplificationContraction.Average,141    )142    print(f"Simplified mesh vertices: {len(mesh.vertices)}, triangles: {len(mesh.triangles)}")143 144    # Crop the mesh to the bounding box of the point cloud145    bbox = pcd.get_axis_aligned_bounding_box()146    mesh_crop = mesh.crop(bbox)147 148    # Save the mesh as a GLTF file149    gltf_path = f"./models/{image_path.stem}.gltf"150    o3d.io.write_triangle_mesh(gltf_path, mesh_crop, write_triangle_uvs=True)151    return gltf_path152 153 154# Define Gradio interface components155title = "Demo: Zero-Shot Depth Estimation with DPT + 3D Point Cloud"156description = (157    "This demo is a variation from the original "158    "<a href='https://huggingface.co/spaces/nielsr/dpt-depth-estimation' target='_blank'>DPT Demo</a>. "159    "It uses the DPT model to predict the depth of an image and then uses 3D Point Cloud to create a 3D object."160)161# Create Gradio sliders for resized_width and z_scale162resized_width_slider = gr.Slider(163    minimum=400,164    maximum=1600,165    step=16,166    value=800,167    label="Resized Width",168    info="Adjust the width to which the input image is resized."169)170 171z_scale_slider = gr.Slider(172    minimum=160,173    maximum=1024,174    step=16,175    value=208,176    label="Z-Scale",177    info="Adjust the scaling factor for the Z-axis in the 3D model."178)179examples = [["examples/" + img] for img in os.listdir("examples/")]180 181iface = gr.Interface(182    fn=process_image,183        inputs=[184        gr.Image(type="filepath", label="Input Image"),185        resized_width_slider,186        z_scale_slider187    ],188    outputs=[189        gr.Image(label="Predicted Depth", type="pil"),190        gr.Model3D(label="3D Mesh Reconstruction", clear_color=[1.0, 1.0, 1.0, 1.0]),191        gr.File(label="3D GLTF"),192    ],193    title=title,194    description=description,195    examples=examples,196    allow_flagging="never",197    cache_examples=False,198    theme="Surn/Beeuty"199)200 201if __name__ == "__main__":202    iface.launch(debug=True, show_api=False, favicon_path="./favicon.ico")203