radames/dpt-depth-estimation-3d-obj

1import gradio as gr2from transformers import DPTFeatureExtractor, DPTForDepthEstimation3import torch4import numpy as np5from PIL import Image6import open3d as o3d7from pathlib import Path8import os9 10feature_extractor = DPTFeatureExtractor.from_pretrained("Intel/dpt-large")11model = DPTForDepthEstimation.from_pretrained("Intel/dpt-large")12 13 14def process_image(image_path):15    image_path = Path(image_path)16    image_raw = Image.open(image_path)17    image = image_raw.resize(18        (800, int(800 * image_raw.size[1] / image_raw.size[0])),19        Image.Resampling.LANCZOS,20    )21 22    # prepare image for the model23    encoding = feature_extractor(image, return_tensors="pt")24 25    # forward pass26    with torch.no_grad():27        outputs = model(**encoding)28        predicted_depth = outputs.predicted_depth29 30    # interpolate to original size31    prediction = torch.nn.functional.interpolate(32        predicted_depth.unsqueeze(1),33        size=image.size[::-1],34        mode="bicubic",35        align_corners=False,36    ).squeeze()37    output = prediction.cpu().numpy()38    depth_image = (output * 255 / np.max(output)).astype("uint8")39    try:40        gltf_path = create_3d_obj(np.array(image), depth_image, image_path)41        img = Image.fromarray(depth_image)42        return [img, gltf_path, gltf_path]43    except Exception as e:44        gltf_path = create_3d_obj(np.array(image), depth_image, image_path, depth=8)45        img = Image.fromarray(depth_image)46        return [img, gltf_path, gltf_path]47    except:48        print("Error reconstructing 3D model")49        raise Exception("Error reconstructing 3D model")50 51 52def create_3d_obj(rgb_image, depth_image, image_path, depth=10):53    depth_o3d = o3d.geometry.Image(depth_image)54    image_o3d = o3d.geometry.Image(rgb_image)55    rgbd_image = o3d.geometry.RGBDImage.create_from_color_and_depth(56        image_o3d, depth_o3d, convert_rgb_to_intensity=False57    )58    w = int(depth_image.shape[1])59    h = int(depth_image.shape[0])60 61    camera_intrinsic = o3d.camera.PinholeCameraIntrinsic()62    camera_intrinsic.set_intrinsics(w, h, 500, 500, w / 2, h / 2)63 64    pcd = o3d.geometry.PointCloud.create_from_rgbd_image(rgbd_image, camera_intrinsic)65 66    print("normals")67    pcd.normals = o3d.utility.Vector3dVector(68        np.zeros((1, 3))69    )  # invalidate existing normals70    pcd.estimate_normals(71        search_param=o3d.geometry.KDTreeSearchParamHybrid(radius=0.01, max_nn=30)72    )73    pcd.orient_normals_towards_camera_location(74        camera_location=np.array([0.0, 0.0, 1000.0])75    )76    pcd.transform([[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]])77    pcd.transform([[-1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]])78 79    print("run Poisson surface reconstruction")80    with o3d.utility.VerbosityContextManager(o3d.utility.VerbosityLevel.Debug) as cm:81        mesh_raw, densities = o3d.geometry.TriangleMesh.create_from_point_cloud_poisson(82            pcd, depth=depth, width=0, scale=1.1, linear_fit=True83        )84 85    voxel_size = max(mesh_raw.get_max_bound() - mesh_raw.get_min_bound()) / 25686    print(f"voxel_size = {voxel_size:e}")87    mesh = mesh_raw.simplify_vertex_clustering(88        voxel_size=voxel_size,89        contraction=o3d.geometry.SimplificationContraction.Average,90    )91 92    # vertices_to_remove = densities < np.quantile(densities, 0.001)93    # mesh.remove_vertices_by_mask(vertices_to_remove)94    bbox = pcd.get_axis_aligned_bounding_box()95    mesh_crop = mesh.crop(bbox)96    gltf_path = f"./{image_path.stem}.gltf"97    o3d.io.write_triangle_mesh(gltf_path, mesh_crop, write_triangle_uvs=True)98    return gltf_path99 100 101title = "Demo: zero-shot depth estimation with DPT + 3D Point Cloud"102description = "This demo is a variation from the original <a href='https://huggingface.co/spaces/nielsr/dpt-depth-estimation' target='_blank'>DPT Demo</a>. It uses the DPT model to predict the depth of an image and then uses 3D Point Cloud to create a 3D object."103examples = [["examples/" + img] for img in os.listdir("examples/")]104 105iface = gr.Interface(106    fn=process_image,107    inputs=[gr.Image(type="filepath", label="Input Image")],108    outputs=[109        gr.Image(label="predicted depth", type="pil"),110        gr.Model3D(label="3d mesh reconstruction", clear_color=[1.0, 1.0, 1.0, 1.0]),111        gr.File(label="3d gLTF"),112    ],113    title=title,114    description=description,115    examples=examples,116    allow_flagging="never",117    cache_examples=False,118)119iface.launch(debug=True, show_api=False)120