Upload depthpro_wrapper/io.py

depthpro_wrapper/io.py

@@ -0,0 +1,102 @@
+"""
+I/O helpers for images and point-cloud files.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Optional, Tuple, Union
+
+import numpy as np
+from PIL import Image
+
+
+def load_image(path: Union[str, Path, Image.Image, np.ndarray]) -> Image.Image:
+    """
+    Normalise any image input to a PIL RGB image.
+
+    Parameters
+    ----------
+    path : str, Path, PIL.Image, or np.ndarray
+        If a string/Path, loaded from disk.  If an ndarray, converted.
+
+    Returns
+    -------
+    PIL.Image.Image
+        RGB image ready for DepthPro.
+    """
+    if isinstance(path, (str, Path)):
+        return Image.open(str(path)).convert("RGB")
+    if isinstance(path, np.ndarray):
+        if path.dtype != np.uint8:
+            path = (path * 255).clip(0, 255).astype(np.uint8)
+        return Image.fromarray(path).convert("RGB")
+    if isinstance(path, Image.Image):
+        return path.convert("RGB")
+    raise TypeError(f"Unsupported image type: {type(path)}")
+
+
+def save_point_cloud(
+    path: Union[str, Path],
+    points: np.ndarray,
+    colors: Optional[np.ndarray] = None,
+    normals: Optional[np.ndarray] = None,
+) -> None:
+    """
+    Save a point cloud to a PLY file.
+
+    Parameters
+    ----------
+    path : str or Path
+        Output file path.  Should end with ``.ply``.
+    points : np.ndarray
+        (N, 3) float array of 3D positions.
+    colors : np.ndarray, optional
+        (N, 3) uint8 RGB colours.
+    normals : np.ndarray, optional
+        (N, 3) float array of normals.
+    """
+    path = Path(path)
+    points = np.asarray(points, dtype=np.float32)
+    has_colors = colors is not None
+    has_normals = normals is not None
+
+    header_lines = [
+        "ply",
+        "format ascii 1.0",
+        f"element vertex {len(points)}",
+        "property float x",
+        "property float y",
+        "property float z",
+    ]
+
+    if has_normals:
+        header_lines += [
+            "property float nx",
+            "property float ny",
+            "property float nz",
+        ]
+
+    if has_colors:
+        header_lines += [
+            "property uchar red",
+            "property uchar green",
+            "property uchar blue",
+        ]
+
+    header_lines += ["end_header"]
+
+    with open(path, "w") as f:
+        f.write("\n".join(header_lines) + "\n")
+        for i in range(len(points)):
+            row = [f"{points[i, 0]:.6f}", f"{points[i, 1]:.6f}", f"{points[i, 2]:.6f}"]
+            if has_normals:
+                row += [
+                    f"{normals[i, 0]:.6f}",
+                    f"{normals[i, 1]:.6f}",
+                    f"{normals[i, 2]:.6f}",
+                ]
+            if has_colors:
+                c = np.clip(colors[i], 0, 255).astype(np.uint8)
+                row += [str(c[0]), str(c[1]), str(c[2])]
+            f.write(" ".join(row) + "\n")