Upload examples/image_to_pointcloud.py
Browse files
examples/image_to_pointcloud.py
ADDED
|
@@ -0,0 +1,73 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Example: Convert a single image to a metric point cloud with UniDepth.
|
| 2 |
+
|
| 3 |
+
Usage:
|
| 4 |
+
python image_to_pointcloud.py room.jpg --output room.ply --checkpoint lpiccinelli/unidepth-v2-vits14
|
| 5 |
+
"""
|
| 6 |
+
|
| 7 |
+
import argparse
|
| 8 |
+
from pathlib import Path
|
| 9 |
+
|
| 10 |
+
from PIL import Image
|
| 11 |
+
from unidepth.inference import UniDepth, save_pointcloud_ply
|
| 12 |
+
|
| 13 |
+
|
| 14 |
+
def main():
|
| 15 |
+
parser = argparse.ArgumentParser(description="Image → metric depth → point cloud")
|
| 16 |
+
parser.add_argument("image", type=str, help="Input image path")
|
| 17 |
+
parser.add_argument("--output", "-o", type=str, default="output.ply", help="Output PLY file")
|
| 18 |
+
parser.add_argument("--checkpoint", type=str, default="lpiccinelli/unidepth-v2-vits14",
|
| 19 |
+
help="HF checkpoint: lpiccinelli/unidepth-v2-vits14 or lpiccinelli/unidepth-v2-vitl14")
|
| 20 |
+
parser.add_argument("--device", type=str, default="cuda", choices=["cuda", "cpu"])
|
| 21 |
+
parser.add_argument("--confidence-threshold", type=float, default=None,
|
| 22 |
+
help="Filter points with confidence > threshold (V2 only)")
|
| 23 |
+
args = parser.parse_args()
|
| 24 |
+
|
| 25 |
+
# ------------------------------------------------------------------ #
|
| 26 |
+
# 1. Load image
|
| 27 |
+
# ------------------------------------------------------------------ #
|
| 28 |
+
image = Image.open(args.image).convert("RGB")
|
| 29 |
+
print(f"Loaded image: {image.size[0]}×{image.size[1]}")
|
| 30 |
+
|
| 31 |
+
# ------------------------------------------------------------------ #
|
| 32 |
+
# 2. Load model (downloads weights on first run)
|
| 33 |
+
# ------------------------------------------------------------------ #
|
| 34 |
+
model = UniDepth.from_pretrained(args.checkpoint, device=args.device)
|
| 35 |
+
|
| 36 |
+
# ------------------------------------------------------------------ #
|
| 37 |
+
# 3. Inference
|
| 38 |
+
# ------------------------------------------------------------------ #
|
| 39 |
+
results = model(image, confidence_threshold=args.confidence_threshold)
|
| 40 |
+
|
| 41 |
+
depth = results["depth"]
|
| 42 |
+
points = results["points"]
|
| 43 |
+
colors = results["colors"]
|
| 44 |
+
intrinsics = results["intrinsics"]
|
| 45 |
+
confidence = results["confidence"]
|
| 46 |
+
|
| 47 |
+
print(f"Depth range: [{depth.min():.3f}, {depth.max():.3f}] meters")
|
| 48 |
+
print(f"Predicted intrinsics K:\n{intrinsics}")
|
| 49 |
+
if confidence is not None:
|
| 50 |
+
print(f"Confidence range: [{confidence.min():.3f}, {confidence.max():.3f}]")
|
| 51 |
+
print(f"Generated {len(points)} 3D points")
|
| 52 |
+
|
| 53 |
+
# ------------------------------------------------------------------ #
|
| 54 |
+
# 4. Save outputs
|
| 55 |
+
# ------------------------------------------------------------------ #
|
| 56 |
+
out_path = Path(args.output)
|
| 57 |
+
out_path.parent.mkdir(parents=True, exist_ok=True)
|
| 58 |
+
|
| 59 |
+
# Save point cloud
|
| 60 |
+
save_pointcloud_ply(str(out_path), points, colors)
|
| 61 |
+
print(f"Saved point cloud to {out_path}")
|
| 62 |
+
|
| 63 |
+
# Optionally save depth map as 16-bit PNG (mm precision)
|
| 64 |
+
depth_png = out_path.with_suffix(".depth.png")
|
| 65 |
+
import numpy as np
|
| 66 |
+
from PIL import Image as PILImage
|
| 67 |
+
depth_mm = (depth * 1000).astype(np.uint16)
|
| 68 |
+
PILImage.fromarray(depth_mm).save(depth_png)
|
| 69 |
+
print(f"Saved depth map to {depth_png}")
|
| 70 |
+
|
| 71 |
+
|
| 72 |
+
if __name__ == "__main__":
|
| 73 |
+
main()
|