app.py

import gradio as gr
from transformers import DPTFeatureExtractor, DPTForDepthEstimation
import torch
import numpy as np
from PIL import Image
import open3d as o3d
from pathlib import Path

# Load model and feature extractor
feature_extractor = DPTFeatureExtractor.from_pretrained("Intel/dpt-large")
model = DPTForDepthEstimation.from_pretrained("Intel/dpt-large")

def process_image(image_path):
    image_path = Path(image_path) if isinstance(image_path, str) else image_path
    try:
        image_raw = Image.open(image_path).convert("RGB")
    except Exception as e:
        return f"Error loading image: {e}"

    # Resize while maintaining aspect ratio
    image = image_raw.resize(
        (800, int(800 * image_raw.size[1] / image_raw.size[0])),
        Image.Resampling.LANCZOS
    )

    encoding = feature_extractor(image, return_tensors="pt")
    
    with torch.no_grad():
        outputs = model(**encoding)
        predicted_depth = outputs.predicted_depth

    # Normalize depth image
    prediction = torch.nn.functional.interpolate(
        predicted_depth.unsqueeze(1),
        size=image.size[::-1],
        mode="bicubic",
        align_corners=False,
    ).squeeze()
    output = prediction.cpu().numpy()
    
    if np.max(output) > 0:
        depth_image = (output * 255 / np.max(output)).astype('uint8')
    else:
        depth_image = np.zeros_like(output, dtype='uint8')  # Handle empty output
    
    gltf_path = create_3d_obj(np.array(image), depth_image, image_path)
    
    if gltf_path and Path(gltf_path).exists():
        return Image.fromarray(depth_image), gltf_path, gltf_path
    else:
        return Image.fromarray(depth_image), None, "3D model generation failed"

def create_3d_obj(rgb_image, depth_image, image_path):
    try:
        depth_o3d = o3d.geometry.Image(depth_image)
        image_o3d = o3d.geometry.Image(rgb_image)
        rgbd_image = o3d.geometry.RGBDImage.create_from_color_and_depth(
            image_o3d, depth_o3d, convert_rgb_to_intensity=False)

        w, h = depth_image.shape[1], depth_image.shape[0]
        camera_intrinsic = o3d.camera.PinholeCameraIntrinsic()
        camera_intrinsic.set_intrinsics(w, h, 500, 500, w / 2, h / 2)

        pcd = o3d.geometry.PointCloud.create_from_rgbd_image(rgbd_image, camera_intrinsic)
        pcd.estimate_normals(
            search_param=o3d.geometry.KDTreeSearchParamHybrid(radius=0.01, max_nn=30))
        pcd.orient_normals_towards_camera_location(camera_location=np.array([0., 0., 1000.]))

        mesh_raw, _ = o3d.geometry.TriangleMesh.create_from_point_cloud_poisson(
            pcd, depth=10, width=0, scale=1.1, linear_fit=True)

        if not mesh_raw.has_triangles():
            return None  # Mesh generation failed
        
        # Center the mesh for better preview
        bbox = pcd.get_axis_aligned_bounding_box()
        mesh_raw.translate(-bbox.get_center())

        # Save the 3D model
        gltf_path = f'./{image_path.stem}.gltf'
        o3d.io.write_triangle_mesh(gltf_path, mesh_raw, write_triangle_uvs=True)
        return gltf_path

    except Exception as e:
        print(f"3D model generation failed: {e}")
        return None

title = "Zero-shot Depth Estimation with DPT + 3D Model Preview"
description = "Upload an image to generate a depth map and reconstruct a 3D model in .gltf format."

with gr.Blocks() as demo:
    gr.Markdown(f"## {title}")
    gr.Markdown(description)

    with gr.Row():
        with gr.Column(scale=1):
            image_input = gr.Image(type="filepath", label="Upload Image")
            generate_button = gr.Button("Generate 3D Model")
        
        with gr.Column(scale=2):
            depth_output = gr.Image(label="Predicted Depth", type="pil")
            model_output = gr.Model3D(label="3D Model Preview (GLTF)")
            file_output = gr.File(label="Download 3D GLTF File")

    generate_button.click(fn=process_image, inputs=[image_input], outputs=[depth_output, model_output, file_output])

if __name__ == "__main__":
    demo.launch()