sudo-soldier commited on
Commit
ccfb544
·
verified ·
1 Parent(s): 854f993

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +19 -115
app.py CHANGED
@@ -1,119 +1,23 @@
1
  import gradio as gr
2
- from transformers import DPTFeatureExtractor, DPTForDepthEstimation
3
- import torch
4
- import numpy as np
5
- from PIL import Image
6
- import open3d as o3d
7
- from pathlib import Path
8
- import subprocess
9
-
10
- # Load model and feature extractor
11
- feature_extractor = DPTFeatureExtractor.from_pretrained("Intel/dpt-large")
12
- model = DPTForDepthEstimation.from_pretrained("Intel/dpt-large")
13
-
14
- def process_image(image_path):
15
- image_path = Path(image_path) if isinstance(image_path, str) else image_path
16
- try:
17
- image_raw = Image.open(image_path).convert("RGB")
18
- except Exception as e:
19
- return f"Error loading image: {e}"
20
-
21
- # Resize while maintaining aspect ratio
22
- image = image_raw.resize(
23
- (800, int(800 * image_raw.size[1] / image_raw.size[0])),
24
- Image.Resampling.LANCZOS
25
- )
26
-
27
- encoding = feature_extractor(image, return_tensors="pt")
28
-
29
- with torch.no_grad():
30
- outputs = model(**encoding)
31
- predicted_depth = outputs.predicted_depth
32
-
33
- # Normalize depth image
34
- prediction = torch.nn.functional.interpolate(
35
- predicted_depth.unsqueeze(1),
36
- size=image.size[::-1],
37
- mode="bicubic",
38
- align_corners=False,
39
- ).squeeze()
40
- output = prediction.cpu().numpy()
41
-
42
- if np.max(output) > 0:
43
- depth_image = (output * 255 / np.max(output)).astype('uint8')
44
- else:
45
- depth_image = np.zeros_like(output, dtype='uint8') # Handle empty output
46
-
47
- glb_path = create_3d_obj(np.array(image), depth_image, image_path)
48
-
49
- if glb_path and Path(glb_path).exists():
50
- return Image.fromarray(depth_image), glb_path, glb_path
51
- else:
52
- return Image.fromarray(depth_image), None, "3D model generation failed"
53
-
54
- def create_3d_obj(rgb_image, depth_image, image_path):
55
- try:
56
- depth_o3d = o3d.geometry.Image(depth_image)
57
- image_o3d = o3d.geometry.Image(rgb_image)
58
- rgbd_image = o3d.geometry.RGBDImage.create_from_color_and_depth(
59
- image_o3d, depth_o3d, convert_rgb_to_intensity=False)
60
-
61
- w, h = depth_image.shape[1], depth_image.shape[0]
62
- camera_intrinsic = o3d.camera.PinholeCameraIntrinsic()
63
- camera_intrinsic.set_intrinsics(w, h, 500, 500, w / 2, h / 2)
64
-
65
- pcd = o3d.geometry.PointCloud.create_from_rgbd_image(rgbd_image, camera_intrinsic)
66
- pcd.estimate_normals(
67
- search_param=o3d.geometry.KDTreeSearchParamHybrid(radius=0.01, max_nn=30))
68
- pcd.orient_normals_towards_camera_location(camera_location=np.array([0., 0., 1000.]))
69
-
70
- mesh_raw, _ = o3d.geometry.TriangleMesh.create_from_point_cloud_poisson(
71
- pcd, depth=10, width=0, scale=1.1, linear_fit=True)
72
-
73
- if not mesh_raw.has_triangles():
74
- print("Mesh generation failed: No triangles in mesh")
75
- return None # Mesh generation failed
76
-
77
- # Center the mesh for better preview
78
- bbox = pcd.get_axis_aligned_bounding_box()
79
- mesh_raw.translate(-bbox.get_center())
80
-
81
- # Save the 3D model as .gltf
82
- gltf_path = str(Path.cwd() / f"{image_path.stem}.gltf")
83
- o3d.io.write_triangle_mesh(gltf_path, mesh_raw, write_triangle_uvs=True)
84
-
85
- # Convert .gltf to .glb
86
- glb_path = gltf_path.replace(".gltf", ".glb")
87
- subprocess.run(["npx", "gltf-pipeline", "-i", gltf_path, "-o", glb_path])
88
-
89
- if Path(glb_path).exists():
90
- return glb_path
91
- else:
92
- print("GLB conversion failed.")
93
- return None
94
-
95
- except Exception as e:
96
- print(f"3D model generation failed: {e}")
97
- return None
98
-
99
- title = "Zero-shot Depth Estimation with DPT + 3D Model Preview"
100
- description = "Upload an image to generate a depth map and reconstruct a 3D model in .glb format."
101
-
102
- with gr.Blocks() as demo:
103
- gr.Markdown(f"## {title}")
104
- gr.Markdown(description)
105
-
106
- with gr.Row():
107
- with gr.Column(scale=1):
108
- image_input = gr.Image(type="filepath", label="Upload Image")
109
- generate_button = gr.Button("Generate 3D Model")
110
-
111
- with gr.Column(scale=2):
112
- depth_output = gr.Image(label="Predicted Depth", type="pil")
113
- model_output = gr.Model3D(label="3D Model Preview (GLB)")
114
- file_output = gr.File(label="Download 3D GLB File")
115
-
116
- generate_button.click(fn=process_image, inputs=[image_input], outputs=[depth_output, model_output, file_output])
117
 
118
  if __name__ == "__main__":
119
  demo.launch()
 
1
  import gradio as gr
2
+ import os
3
+
4
+ def load_mesh(mesh_file_name):
5
+ return mesh_file_name
6
+
7
+ demo = gr.Interface(
8
+ fn=load_mesh,
9
+ inputs=gr.Model3D(),
10
+ outputs=gr.Model3D(
11
+ clear_color=(0.0, 0.0, 0.0, 0.0), label="3D Model", display_mode="wireframe"),
12
+ examples=[
13
+ [os.path.join(os.path.dirname(__file__), "files/model1.glb")],
14
+ [os.path.join(os.path.dirname(__file__), "files/model2.glb")],
15
+ [os.path.join(os.path.dirname(__file__), "files/model3.glb")],
16
+ [os.path.join(os.path.dirname(__file__), "files/model4.glb")],
17
+ ["https://huggingface.co/datasets/dylanebert/3dgs/resolve/main/bonsai/bonsai-7k-mini.splat"],
18
+ ],
19
+ cache_examples=True
20
+ )
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
21
 
22
  if __name__ == "__main__":
23
  demo.launch()