File size: 4,839 Bytes
1cb8231 a3d35f1 1cb8231 a3d35f1 1cb8231 a3d35f1 1cb8231 a3d35f1 1cb8231 a3d35f1 1cb8231 a3d35f1 1cb8231 a3d35f1 1cb8231 c825b73 1cb8231 c825b73 1cb8231 c825b73 1cb8231 c825b73 1cb8231 a3d35f1 c825b73 8c61c71 c825b73 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 |
from email.policy import default
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
import os
feature_extractor = DPTFeatureExtractor.from_pretrained("Intel/dpt-large")
model = DPTForDepthEstimation.from_pretrained("Intel/dpt-large")
def process_image(image_path, voxel_s):
voxel_s = max(voxel_s/500, 0.0001)
image_path = Path(image_path)
image_raw = Image.open(image_path)
image = image_raw.resize(
(800, int(800 * image_raw.size[1] / image_raw.size[0])),
Image.Resampling.LANCZOS)
# prepare image for the model
encoding = feature_extractor(image, return_tensors="pt")
# forward pass
with torch.no_grad():
outputs = model(**encoding)
predicted_depth = outputs.predicted_depth
# interpolate to original size
prediction = torch.nn.functional.interpolate(
predicted_depth.unsqueeze(1),
size=image.size[::-1],
mode="bicubic",
align_corners=False,
).squeeze()
output = prediction.cpu().numpy()
depth_image = (output * 255 / np.max(output)).astype('uint8')
try:
gltf_path = create_3d_voxels_obj(
np.array(image), depth_image, image_path, voxel_s)
img = Image.fromarray(depth_image)
return [img, gltf_path, gltf_path]
except Exception as e:
print("Error reconstructing 3D model")
raise Exception("Error reconstructing 3D model")
def create_3d_voxels_obj(rgb_image, depth_image, image_path, voxel_s):
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 = int(depth_image.shape[1])
h = int(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)
print('normals')
pcd.normals = o3d.utility.Vector3dVector(
np.zeros((1, 3))) # invalidate existing normals
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.]))
pcd.transform([[1, 0, 0, 0],
[0, -1, 0, 0],
[0, 0, -1, 0],
[0, 0, 0, 1]])
pcd.transform([[-1, 0, 0, 0],
[0, 1, 0, 0],
[0, 0, 1, 0],
[0, 0, 0, 1]])
print('voxels')
# ref https://towardsdatascience.com/how-to-automate-voxel-modelling-of-3d-point-cloud-with-python-459f4d43a227
voxel_size = round(
max(pcd.get_max_bound()-pcd.get_min_bound())*voxel_s, 10)
print("Voxel size", voxel_size, "voxel_s", voxel_s)
voxel_grid = o3d.geometry.VoxelGrid.create_from_point_cloud(
pcd, voxel_size=voxel_size)
voxels = voxel_grid.get_voxels()
vox_mesh = o3d.geometry.TriangleMesh()
for v in voxels:
cube = o3d.geometry.TriangleMesh.create_box(width=1, height=1, depth=1)
cube.paint_uniform_color(v.color)
cube.translate(v.grid_index, relative=False)
vox_mesh += cube
print(voxel_grid, vox_mesh)
gltf_path = f'./{image_path.stem}.gltf'
o3d.io.write_triangle_mesh(gltf_path, vox_mesh, write_triangle_uvs=True)
return gltf_path
title = "Demo: zero-shot depth estimation with DPT + 3D Voxels reconstruction"
description = "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 reconstruct the 3D model as voxels."
examples = [["examples/" + img, 10] for img in os.listdir("examples/")]
iface = gr.Interface(fn=process_image,
inputs=[
gr.Image(
type="filepath", label="Input Image"),
gr.Slider(value=10, minimum=5, maximum=100, step=1, label="Voxel Size",)
],
outputs=[
gr.Image(label="predicted depth", type="pil"),
gr.Model3D(label="3d mesh reconstruction", clear_color=[
1.0, 1.0, 1.0, 1.0]),
gr.File(label="3d gLTF")
],
title=title,
description=description,
examples=examples,
allow_flagging="never",
cache_examples=False)
iface.launch(debug=True)
|