File size: 2,824 Bytes
ab4b137 | 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 | """
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")
|