Spaces:
Running
on
Zero
Running
on
Zero
File size: 3,908 Bytes
6ef117e |
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 |
# utils/depth_estimation.py
import torch
import numpy as np
from PIL import Image
import open3d as o3d
from transformers import DPTImageProcessor, DPTForDepthEstimation
from pathlib import Path
import logging
logging.getLogger("transformers.modeling_utils").setLevel(logging.ERROR)
from utils.image_utils import (
change_color,
open_image,
build_prerendered_images,
upscale_image,
crop_and_resize_image,
resize_image_with_aspect_ratio,
show_lut,
apply_lut_to_image_path
)
# Load models once during module import
image_processor = DPTImageProcessor.from_pretrained("Intel/dpt-large")
depth_model = DPTForDepthEstimation.from_pretrained("Intel/dpt-large", ignore_mismatched_sizes=True)
def estimate_depth(image):
# Ensure image is in RGB mode
if image.mode != "RGB":
image = image.convert("RGB")
# Resize the image for the model
image_resized = image.resize(
(image.width, image.height),
Image.Resampling.LANCZOS
)
# Prepare image for the model
encoding = image_processor(image_resized, return_tensors="pt")
# Forward pass
with torch.no_grad():
outputs = depth_model(**encoding)
predicted_depth = outputs.predicted_depth
# Interpolate to original size
prediction = torch.nn.functional.interpolate(
predicted_depth.unsqueeze(1),
size=(image.height, image.width),
mode="bicubic",
align_corners=False,
).squeeze()
# Convert to depth image
output = prediction.cpu().numpy()
depth_min = output.min()
depth_max = output.max()
max_val = (2**8) - 1
# Normalize and convert to 8-bit image
depth_image = max_val * (output - depth_min) / (depth_max - depth_min)
depth_image = depth_image.astype("uint8")
depth_pil = Image.fromarray(depth_image)
return depth_pil, output
def create_3d_model(rgb_image, depth_array, voxel_size_factor=0.01):
depth_o3d = o3d.geometry.Image(depth_array.astype(np.float32))
rgb_o3d = o3d.geometry.Image(np.array(rgb_image))
rgbd_image = o3d.geometry.RGBDImage.create_from_color_and_depth(
rgb_o3d,
depth_o3d,
convert_rgb_to_intensity=False
)
# Create a point cloud from the RGBD image
camera_intrinsic = o3d.camera.PinholeCameraIntrinsic(
rgb_image.width,
rgb_image.height,
fx=1.0,
fy=1.0,
cx=rgb_image.width / 2.0,
cy=rgb_image.height / 2.0,
)
pcd = o3d.geometry.PointCloud.create_from_rgbd_image(
rgbd_image,
camera_intrinsic
)
# Voxel downsample
voxel_size = max(pcd.get_max_bound() - pcd.get_min_bound()) * voxel_size_factor
voxel_grid = o3d.geometry.VoxelGrid.create_from_point_cloud(pcd, voxel_size=voxel_size)
# Save the 3D model to a temporary file
temp_dir = Path.cwd() / "temp_models"
temp_dir.mkdir(exist_ok=True)
model_path = temp_dir / "model.ply"
o3d.io.write_voxel_grid(str(model_path), voxel_grid)
return str(model_path)
def generate_depth_and_3d(input_image_path, voxel_size_factor):
image = Image.open(input_image_path).convert("RGB")
resized_image = resize_image_with_aspect_ratio(image, 2688, 1680)
depth_image, depth_array = estimate_depth(resized_image)
model_path = create_3d_model(resized_image, depth_array, voxel_size_factor=voxel_size_factor)
return depth_image, model_path
def generate_depth_button_click(depth_image_source, voxel_size_factor, input_image, output_image, overlay_image, bordered_image_output):
if depth_image_source == "Input Image":
image_path = input_image
elif depth_image_source == "Output Image":
image_path = output_image
elif depth_image_source == "Image with Margins":
image_path = bordered_image_output
else:
image_path = overlay_image
return generate_depth_and_3d(image_path, voxel_size_factor) |