Implement green screen matting pipeline

bgfilter/io.py

@@ -0,0 +1,39 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+import numpy as np
+from PIL import Image
+
+
+def load_rgb(path: str | Path) -> np.ndarray:
+    image = Image.open(path).convert("RGB")
+    return np.asarray(image, dtype=np.uint8)
+
+
+def save_rgb(path: str | Path, rgb: np.ndarray) -> None:
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    Image.fromarray(np.clip(rgb, 0, 255).astype(np.uint8), mode="RGB").save(path)
+
+
+def save_gray(path: str | Path, gray: np.ndarray) -> None:
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    if gray.dtype != np.uint8:
+        gray = np.clip(gray * 255.0, 0, 255).astype(np.uint8)
+    Image.fromarray(gray, mode="L").save(path)
+
+
+def save_rgba(path: str | Path, rgb: np.ndarray, alpha: np.ndarray) -> None:
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    alpha_u8 = np.clip(alpha * 255.0, 0, 255).astype(np.uint8)
+    rgba = np.dstack([np.clip(rgb, 0, 255).astype(np.uint8), alpha_u8])
+    Image.fromarray(rgba, mode="RGBA").save(path)
+
+
+def write_text(path: str | Path, text: str) -> None:
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(text, encoding="utf-8")