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

# تحميل النموذج المدرب مسبقًا لتقدير العمق
depth_extractor = DPTFeatureExtractor.from_pretrained("Intel/dpt-large")
depth_model = DPTForDepthEstimation.from_pretrained("Intel/dpt-large")

def estimate_depth_and_generate_3d(image_file, voxel_scale):
    voxel_scale = max(voxel_scale / 500.0, 0.0001)
    image_path = Path(image_file)
    image_raw = Image.open(image_path)
    resized = image_raw.resize((800, int(800 * image_raw.height / image_raw.width)), Image.Resampling.LANCZOS)

    inputs = depth_extractor(resized, return_tensors="pt")

    with torch.no_grad():
        result = depth_model(**inputs)
        depth_tensor = result.predicted_depth

    depth_resized = torch.nn.functional.interpolate(
        depth_tensor.unsqueeze(1),
        size=resized.size[::-1],
        mode="bicubic",
        align_corners=False
    ).squeeze().cpu().numpy()

    depth_image = (depth_resized * 255 / np.max(depth_resized)).astype(np.uint8)

    try:
        gltf_file = create_voxel_model(np.array(resized), depth_image, image_path, voxel_scale)
        return [Image.fromarray(depth_image), gltf_file, gltf_file]
    except Exception as err:
        print("3D creation error:", err)
        raise gr.Error("Failed to generate 3D model.")

def create_voxel_model(rgb_data, depth_data, image_path, voxel_scale):
    depth_o3d = o3d.geometry.Image(depth_data)
    color_o3d = o3d.geometry.Image(rgb_data)
    rgbd = o3d.geometry.RGBDImage.create_from_color_and_depth(color_o3d, depth_o3d, convert_rgb_to_intensity=False)

    h, w = depth_data.shape
    intrinsics = o3d.camera.PinholeCameraIntrinsic()
    intrinsics.set_intrinsics(w, h, 500, 500, w / 2, h / 2)

    pcd = o3d.geometry.PointCloud.create_from_rgbd_image(rgbd, intrinsics)

    pcd.estimate_normals(search_param=o3d.geometry.KDTreeSearchParamHybrid(radius=0.01, max_nn=30))
    pcd.orient_normals_towards_camera_location(np.array([0, 0, 1000]))
    pcd.transform([[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]])

    voxel_size = round(max(pcd.get_max_bound() - pcd.get_min_bound()) * voxel_scale, 10)
    vox_grid = o3d.geometry.VoxelGrid.create_from_point_cloud(pcd, voxel_size=voxel_size)

    mesh = o3d.geometry.TriangleMesh()
    for voxel in vox_grid.get_voxels():
        cube = o3d.geometry.TriangleMesh.create_box(1, 1, 1)
        cube.paint_uniform_color(voxel.color)
        cube.translate(voxel.grid_index, relative=False)
        mesh += cube

    output_file = f"./{image_path.stem}_3d.gltf"
    o3d.io.write_triangle_mesh(output_file, mesh, write_triangle_uvs=True)
    return output_file

# واجهة الاستخدام
gr.Interface(
    fn=estimate_depth_and_generate_3d,
    inputs=[
        gr.Image(type="filepath", label="Upload Image (PNG or JPG)"),
        gr.Slider(5, 100, value=10, step=1, label="Voxel Density")
    ],
    outputs=[
        gr.Image(label="Estimated Depth Map"),
        gr.Model3D(label="3D Voxel Mesh"),
        gr.File(label="Download GLTF")
    ],
    title="3D Reconstruction from Depth Estimation",
    description="Upload an image to estimate its depth and reconstruct a voxel-based 3D mesh.",
    allow_flagging="never"
).launch(debug=True)