File size: 3,077 Bytes
3941fab | 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 | """Example: Convert a single image to a metric point cloud with UniDepth.
Usage:
python image_to_pointcloud.py room.jpg --output room.ply --checkpoint lpiccinelli/unidepth-v2-vits14
"""
import argparse
from pathlib import Path
from PIL import Image
from unidepth.inference import UniDepth, save_pointcloud_ply
def main():
parser = argparse.ArgumentParser(description="Image → metric depth → point cloud")
parser.add_argument("image", type=str, help="Input image path")
parser.add_argument("--output", "-o", type=str, default="output.ply", help="Output PLY file")
parser.add_argument("--checkpoint", type=str, default="lpiccinelli/unidepth-v2-vits14",
help="HF checkpoint: lpiccinelli/unidepth-v2-vits14 or lpiccinelli/unidepth-v2-vitl14")
parser.add_argument("--device", type=str, default="cuda", choices=["cuda", "cpu"])
parser.add_argument("--confidence-threshold", type=float, default=None,
help="Filter points with confidence > threshold (V2 only)")
args = parser.parse_args()
# ------------------------------------------------------------------ #
# 1. Load image
# ------------------------------------------------------------------ #
image = Image.open(args.image).convert("RGB")
print(f"Loaded image: {image.size[0]}×{image.size[1]}")
# ------------------------------------------------------------------ #
# 2. Load model (downloads weights on first run)
# ------------------------------------------------------------------ #
model = UniDepth.from_pretrained(args.checkpoint, device=args.device)
# ------------------------------------------------------------------ #
# 3. Inference
# ------------------------------------------------------------------ #
results = model(image, confidence_threshold=args.confidence_threshold)
depth = results["depth"]
points = results["points"]
colors = results["colors"]
intrinsics = results["intrinsics"]
confidence = results["confidence"]
print(f"Depth range: [{depth.min():.3f}, {depth.max():.3f}] meters")
print(f"Predicted intrinsics K:\n{intrinsics}")
if confidence is not None:
print(f"Confidence range: [{confidence.min():.3f}, {confidence.max():.3f}]")
print(f"Generated {len(points)} 3D points")
# ------------------------------------------------------------------ #
# 4. Save outputs
# ------------------------------------------------------------------ #
out_path = Path(args.output)
out_path.parent.mkdir(parents=True, exist_ok=True)
# Save point cloud
save_pointcloud_ply(str(out_path), points, colors)
print(f"Saved point cloud to {out_path}")
# Optionally save depth map as 16-bit PNG (mm precision)
depth_png = out_path.with_suffix(".depth.png")
import numpy as np
from PIL import Image as PILImage
depth_mm = (depth * 1000).astype(np.uint16)
PILImage.fromarray(depth_mm).save(depth_png)
print(f"Saved depth map to {depth_png}")
if __name__ == "__main__":
main()
|