"""
Batch-processing example: estimate depth for a folder of images and export
point clouds with one call.
"""

from pathlib import Path
from depthpro_wrapper import DepthProEstimator, rgbd_to_point_cloud, save_point_cloud

image_dir = Path("photos/")
output_dir = Path("pointclouds/")
output_dir.mkdir(exist_ok=True)

estimator = DepthProEstimator(device="cuda:0")

# Batch size = 4 is a sweet spot for most GPUs
image_paths = sorted(image_dir.glob("*.jpg"))
for i in range(0, len(image_paths), 4):
    batch_paths = image_paths[i : i + 4]
    results = estimator.estimate_batch(batch_paths)

    for path, result in zip(batch_paths, results):
        points, colors = rgbd_to_point_cloud(
            result.depth,
            result.image,
            result.focal_length,
        )
        out_path = output_dir / f"{path.stem}.ply"
        save_point_cloud(out_path, points, colors=colors)
        print(f"  {path.name} → {out_path} ({len(points):,} pts)")