File size: 4,611 Bytes
808837f | 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 | #!/usr/bin/env python3
"""
CLI script: image → metric depth → 3D point cloud using DepthPro.
Usage
-----
python image_to_pointcloud.py photo.jpg output.ply --device cuda:0 --sample-step 2
python image_to_pointcloud.py photo.jpg output.ply --colored --save-depth depth.png
"""
import argparse
import sys
from pathlib import Path
import numpy as np
sys.path.insert(0, str(Path(__file__).parent.parent))
from depthpro_wrapper import (
DepthProEstimator,
depth_to_point_cloud,
rgbd_to_point_cloud,
normals_from_depth,
load_image,
save_point_cloud,
)
def main() -> None:
parser = argparse.ArgumentParser(
description="Apple DepthPro: image → metric depth → 3D point cloud"
)
parser.add_argument("image", type=Path, help="Input RGB image")
parser.add_argument("output", type=Path, help="Output point cloud (.ply)")
parser.add_argument("--device", default="cuda:0", help="PyTorch device")
parser.add_argument("--dtype", choices=["float16", "float32"], default="float16", help="Inference dtype")
parser.add_argument("--colored", action="store_true", help="Include per-point RGB colours")
parser.add_argument("--normals", action="store_true", help="Include per-point normals")
parser.add_argument("--sample-step", type=int, default=1, help="Spatial sub-sampling (1 = full res, 2 = 1/4 points)")
parser.add_argument("--save-depth", type=Path, default=None, help="Also save depth map as .npy")
parser.add_argument("--save-confidence", type=Path, default=None, help="Also save confidence map as .npy")
args = parser.parse_args()
if not args.image.exists():
parser.error(f"Input image not found: {args.image}")
# ---- depth estimation -----------------------------------------------
print(f"Loading DepthPro on {args.device} (dtype={args.dtype}) ...")
dtype = {"float16": "float16", "float32": "float32"}[args.dtype]
import torch
torch_dtype = torch.float16 if dtype == "float16" else torch.float32
estimator = DepthProEstimator(device=args.device, dtype=torch_dtype)
print(f"Estimating depth for {args.image} ...")
result = estimator.estimate(
args.image,
return_confidence=args.save_confidence is not None,
)
print(f" Image size: {result.width}×{result.height}")
print(f" Estimated focal length: {result.focal_length:.1f} px")
print(f" Estimated FOV: {result.field_of_view:.1f}°")
print(f" Depth range: {result.depth.min():.2f} m – {result.depth.max():.2f} m")
# ---- optional saves -------------------------------------------------
if args.save_depth:
np.save(args.save_depth, result.depth)
print(f" Saved depth map → {args.save_depth}")
if args.save_confidence and result.confidence is not None:
np.save(args.save_confidence, result.confidence)
print(f" Saved confidence map → {args.save_confidence}")
# ---- back-projection ------------------------------------------------
print("\nBack-projecting to 3D point cloud ...")
normals = None
if args.normals:
normals = normals_from_depth(result.depth, result.focal_length)
if args.colored:
points, colors = rgbd_to_point_cloud(
result.depth,
result.image,
result.focal_length,
sample_step=args.sample_step,
)
if args.normals:
# Sample normals at same grid
H, W = result.depth.shape
v_idx = np.arange(0, H, args.sample_step)
u_idx = np.arange(0, W, args.sample_step)
valid = result.depth[v_idx[:, None], u_idx[None, :]] > 0
normals = normals[v_idx[:, None], u_idx[None, :]]
normals = normals[valid]
print(f" Colored point cloud: {len(points):,} points")
save_point_cloud(args.output, points, colors=colors, normals=normals)
else:
points = depth_to_point_cloud(
result.depth,
result.focal_length,
sample_step=args.sample_step,
)
if args.normals:
H, W = result.depth.shape
v_idx = np.arange(0, H, args.sample_step)
u_idx = np.arange(0, W, args.sample_step)
valid = result.depth[v_idx[:, None], u_idx[None, :]] > 0
normals = normals[v_idx[:, None], u_idx[None, :]]
normals = normals[valid]
print(f" Point cloud: {len(points):,} points")
save_point_cloud(args.output, points, normals=normals)
print(f"\nDone — saved to {args.output}")
if __name__ == "__main__":
main()
|