ffgs 0.1.0.dev0__py3-none-any.whl

ffgs/processor.py

@@ -0,0 +1,236 @@
+"""Processor base: a model's pre / post-processing, stored as `processor_config.json`.
+
+A processor owns everything between the standard `Views` / `Gaussians` and one
+model's tensors: resizing to the training rule, normalising cameras into the
+model's frame (the `ModelFrame` that also takes predictions back to the caller's
+world frame; fitted to the given poses for models that predict cameras), and
+building the model input. Its settings are plain JSON so they
+travel with the weights.
+
+Configs that subclass `ImageFitConfig` get `fit_views` / `fit_cameras`: resize +
+centre crop by default (the training rule), or pad / pass-through, and optionally a
+custom resize callable (`Processor(..., resize=fn)`, not serialised).
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import asdict, dataclass, field, fields
+from pathlib import Path
+from typing import Any, ClassVar
+
+import torch
+
+from .geometry import ModelFrame
+from .image import (
+    CROP_MODES,
+    RESIZE_MODES,
+    Resize,
+    fit_images,
+    fit_intrinsics,
+)
+from .types import Cameras, Gaussians, Views
+
+PROCESSOR_CONFIG_NAME = "processor_config.json"
+
+
+@dataclass
+class Prepared:
+    """A processor's output for one call: what the model eats, and the frame.
+
+    `frame` is None for models that predict their cameras: it is then fitted
+    after the model, from the predicted cameras (see `Processor.output_frame`).
+    """
+
+    model_input: Any
+    frame: ModelFrame | None
+
+
+@dataclass
+class ImageFitConfig:
+    """How inputs are brought to the model's input shape. Subclass per model and
+    re-declare fields to change their defaults."""
+
+    # [h, w] the model sees.
+    image_shape: list[int] = field(default_factory=lambda: [256, 256])
+    # Built-in resize kernel (see `ffgs.image.rescale`).
+    resize_mode: str = "bilinear"
+    # "crop": cover + centre crop (the training rule); "pad": fit inside + pad;
+    # "none": inputs are already `image_shape`.
+    crop_mode: str = "crop"
+    # Fill value of the padded border for crop_mode="pad".
+    pad_value: float = 0.0
+
+    def __post_init__(self) -> None:
+        self.image_shape = [int(x) for x in self.image_shape]
+        if len(self.image_shape) != 2:
+            raise ValueError(f"image_shape must be [h, w], got {self.image_shape}")
+        if self.resize_mode not in RESIZE_MODES:
+            raise ValueError(f"resize_mode must be one of {RESIZE_MODES}")
+        if self.crop_mode not in CROP_MODES:
+            raise ValueError(f"crop_mode must be one of {CROP_MODES}")
+
+
+class Processor:
+    config_cls: ClassVar[type]
+
+    def __init__(
+        self,
+        config: Any | None = None,
+        *,
+        resize: Resize | None = None,
+        **overrides: Any,
+    ) -> None:
+        if config is None:
+            config = self.config_cls()
+        elif isinstance(config, dict):
+            config = self.config_cls(**config)
+        if overrides:
+            config = self.config_cls(**{**asdict(config), **overrides})
+        self.config = config
+        # A custom resize callable for `fit_views`; replaces `resize_mode`.
+        self.resize = resize
+
+    # --- fitting inputs (configs subclassing ImageFitConfig) ----------------------
+
+    def fit_views(self, views: Views, image_shape: list[int] | None = None) -> Views:
+        """Bring `views` to `image_shape` (default: the model's) by the config's
+        `crop_mode`. Also the way to derive ground truth for target views. Returns
+        normalised intrinsics (None if `views` had none); c2w untouched."""
+        cfg = self._fit_config()
+        shape = tuple(image_shape or cfg.image_shape)
+        if views.intrinsics is None:
+            # The images are fitted the same way whatever K is; any valid one does.
+            source_k = torch.tensor(
+                [[1.0, 0, 0.5], [0, 1.0, 0.5], [0, 0, 1]], device=views.images.device
+            ).expand(*views.images.shape[:2], 3, 3)
+        else:
+            source_k = views.normalized_k
+        images, intrinsics = [], []
+        for i in range(views.images.shape[0]):
+            img, k = fit_images(
+                views.images[i],
+                source_k[i],
+                shape,
+                resize_mode=cfg.resize_mode,
+                crop_mode=cfg.crop_mode,
+                resize=self.resize,
+                pad_value=cfg.pad_value,
+            )
+            images.append(img)
+            intrinsics.append(k)
+        return Views(
+            torch.stack(images),
+            None if views.intrinsics is None else torch.stack(intrinsics),
+            views.c2w,
+            normalized_intrinsics=True,
+        )
+
+    def fit_cameras(
+        self, cameras: Cameras, image_shape: list[int] | None = None
+    ) -> Cameras:
+        """`fit_views` for cameras alone: where to render to compare against
+        `fit_views` images. `cameras.image_shape` is the source image size."""
+        cfg = self._fit_config()
+        shape = tuple(image_shape or cfg.image_shape)
+        k = fit_intrinsics(
+            cameras.normalized_k, cameras.image_shape, shape, cfg.crop_mode
+        )
+        return Cameras(cameras.c2w, k, shape, normalized_intrinsics=True)
+
+    def _fit_config(self) -> ImageFitConfig:
+        if not isinstance(self.config, ImageFitConfig):
+            raise TypeError(
+                f"{type(self.config).__name__} does not subclass ImageFitConfig"
+            )
+        return self.config
+
+    # --- serialisation ----------------------------------------------------------
+
+    def to_dict(self) -> dict[str, Any]:
+        return asdict(self.config)
+
+    def save_pretrained(self, save_directory: str | Path) -> Path:
+        path = Path(save_directory) / PROCESSOR_CONFIG_NAME
+        path.parent.mkdir(parents=True, exist_ok=True)
+        path.write_text(json.dumps(self.to_dict(), indent=2, sort_keys=True) + "\n")
+        return path
+
+    @classmethod
+    def from_pretrained(cls, pretrained: str | Path, **hub_kwargs: Any) -> Processor:
+        data = read_json(pretrained, PROCESSOR_CONFIG_NAME, **hub_kwargs)
+        return cls.from_dict(data, source=PROCESSOR_CONFIG_NAME)
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any], source: str = "config") -> Processor:
+        known = {f.name for f in fields(cls.config_cls)}
+        unknown = set(data) - known
+        if unknown:
+            raise ValueError(
+                f"{source} has keys {sorted(unknown)} unknown to {cls.__name__}"
+            )
+        return cls(data)
+
+    # --- the per-model contract -------------------------------------------------
+
+    def preprocess(
+        self, views: Views, device: torch.device, **options: Any
+    ) -> Prepared:
+        raise NotImplementedError
+
+    def postprocess(self, model_output: Any) -> Gaussians:
+        """Raw model output -> model-frame `Gaussians`."""
+        raise NotImplementedError
+
+    def predicted_cameras(self, model_output: Any) -> Cameras | None:
+        """The input cameras a model predicts, in its frame (normalised K for the
+        model's input images); None for models that take cameras as input."""
+        return None
+
+    def output_frame(
+        self, views: Views, prepared: Prepared, cameras: Cameras | None
+    ) -> ModelFrame | None:
+        """The model frame of this call. Posed models fix it in `preprocess`;
+        for models that predict cameras it is fitted to the given poses over all
+        views (`ModelFrame.fit`), and None when there are none."""
+        if prepared.frame is not None:
+            return prepared.frame
+        if cameras is None or views.c2w is None:
+            return None
+        return ModelFrame.fit(cameras.c2w, views.c2w)
+
+    def render_planes(self, frame: ModelFrame | None) -> tuple[float, float]:
+        """(near, far) for Gaussians predicted in `frame`: in world units, or in
+        model units when there is no frame (no input poses)."""
+        raise NotImplementedError
+
+    @property
+    def background(self) -> tuple[float, float, float] | None:
+        return None
+
+    @property
+    def render_settings(self) -> dict[str, Any]:
+        """The model's other rasteriser settings: keyword arguments of
+        `ffgs.render.render` (e.g. `radius_clip`, `rasterize_mode`, `clamp`)
+        that `GSPipeline.render` uses unless the caller passes them."""
+        return {}
+
+    @property
+    def autocast_dtype(self) -> torch.dtype | None:
+        return None
+
+
+def read_json(pretrained: str | Path, filename: str, **hub_kwargs: Any) -> dict:
+    """`filename` from a local directory or an HF Hub repo id."""
+    local = Path(pretrained)
+    if local.is_dir():
+        path = local / filename
+        if not path.is_file():
+            raise FileNotFoundError(f"{filename} not found in {local}")
+    else:
+        from huggingface_hub import hf_hub_download
+
+        path = Path(
+            hf_hub_download(repo_id=str(pretrained), filename=filename, **hub_kwargs)
+        )
+    return json.loads(path.read_text())

ffgs/registry.py

@@ -0,0 +1,132 @@
+"""model_type -> (model class, processor class).
+
+A `model_type` is looked up in order:
+
+1. models registered in this process (`register_model`), including built-ins,
+   which are imported only when their `model_type` is first asked for;
+2. installed plugins: the `ffgs.models` entry point group, entry point name =
+   model_type, value = a `ModelSpec` (or a module that calls `register_model`
+   on import). Plugins are loaded only when their `model_type` is asked for;
+3. code shipped in the model repo (`config.json:auto_map`), only with
+   `trust_remote_code=True` — see `ffgs.dynamic` and `GSPipeline.from_pretrained`.
+
+`import ffgs` imports no plugin and runs no repo code.
+"""
+
+from __future__ import annotations
+
+import importlib
+import warnings
+from collections.abc import Callable
+from dataclasses import dataclass
+from importlib.metadata import EntryPoint, entry_points
+from typing import Any
+
+from .processor import Processor
+
+ENTRY_POINT_GROUP = "ffgs.models"
+
+# model_type -> module that calls `register_model` for it on import.
+_BUILTIN: dict[str, str] = {
+    "anysplat": "ffgs.models.anysplat",
+    "tokengs": "ffgs.models.tokengs",
+}
+
+
+@dataclass(frozen=True)
+class ModelSpec:
+    model_type: str
+    model_cls: type
+    processor_cls: type[Processor]
+    # Upstream checkpoint (as loaded from a zoo entry's weights file) -> state dict
+    # of `model_cls`. Not applied to directories written by `save_pretrained`.
+    convert_state_dict: Callable[[Any], dict[str, Any]] | None = None
+
+
+_REGISTRY: dict[str, ModelSpec] = {}
+
+
+def register_model(spec: ModelSpec) -> ModelSpec:
+    existing = _REGISTRY.get(spec.model_type)
+    if existing is not None and existing != spec:
+        raise ValueError(f"model_type {spec.model_type!r} is already registered")
+    _REGISTRY[spec.model_type] = spec
+    return spec
+
+
+def get_model_spec(model_type: str) -> ModelSpec:
+    if model_type not in _REGISTRY and model_type in _BUILTIN:
+        importlib.import_module(_BUILTIN[model_type])
+    if model_type not in _REGISTRY:
+        plugins = _plugin_entry_points().get(model_type, [])
+        if len(plugins) > 1:
+            raise ValueError(
+                f"model_type {model_type!r} is provided by several installed "
+                f"plugins: {[ep.value for ep in plugins]}; uninstall all but one"
+            )
+        if plugins:
+            _load_plugin(plugins[0])
+    if model_type not in _REGISTRY:
+        raise KeyError(
+            f"unknown model_type {model_type!r}; known: {known_model_types()}. "
+            "Install the package that provides it (found through the "
+            f"{ENTRY_POINT_GROUP!r} entry point group) or call "
+            "ffgs.register_model(...) before loading. A model repo can also ship "
+            "its own code under config.json:auto_map, loaded with "
+            "GSPipeline.from_pretrained(..., trust_remote_code=True)."
+        )
+    return _REGISTRY[model_type]
+
+
+def known_model_types() -> list[str]:
+    """Registered, built-in and installed-plugin model types (nothing imported)."""
+    return sorted(set(_REGISTRY) | set(_BUILTIN) | set(_plugin_entry_points()))
+
+
+def spec_for_model(model: object) -> ModelSpec:
+    for model_type in _BUILTIN:
+        get_model_spec(model_type)
+    for model_type, plugins in _plugin_entry_points().items():
+        if model_type in _REGISTRY:
+            continue
+        try:
+            get_model_spec(model_type)
+        except Exception as err:  # a broken plugin must not hide the others
+            warnings.warn(
+                f"skipping ffgs plugin {[ep.value for ep in plugins]}: {err}",
+                stacklevel=2,
+            )
+    # The model's own class first: a plugin may subclass another model.
+    for spec in _REGISTRY.values():
+        if type(model) is spec.model_cls:
+            return spec
+    for spec in _REGISTRY.values():
+        if isinstance(model, spec.model_cls):
+            return spec
+    raise KeyError(f"no model_type registered for {type(model).__name__}")
+
+
+def _plugin_entry_points() -> dict[str, list[EntryPoint]]:
+    found: dict[str, list[EntryPoint]] = {}
+    for ep in entry_points(group=ENTRY_POINT_GROUP):
+        # The same distribution can be visible twice on sys.path.
+        if all(other.value != ep.value for other in found.get(ep.name, [])):
+            found.setdefault(ep.name, []).append(ep)
+    return found
+
+
+def _load_plugin(ep: EntryPoint) -> None:
+    loaded = ep.load()
+    if isinstance(loaded, ModelSpec):
+        if loaded.model_type != ep.name:
+            raise ValueError(
+                f"entry point {ep.name!r} = {ep.value!r} is a ModelSpec for "
+                f"model_type {loaded.model_type!r}; the entry point name must "
+                "be the model_type"
+            )
+        register_model(loaded)
+    elif ep.name not in _REGISTRY:
+        raise TypeError(
+            f"entry point {ep.name!r} = {ep.value!r} must be a ModelSpec, or a "
+            f"module that calls register_model for {ep.name!r} on import"
+        )

ffgs/render.py

@@ -0,0 +1,111 @@
+"""Rasterise `Gaussians` at `Cameras` with gsplat (imported on first use)."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+import torch
+
+from .types import Cameras, Gaussians
+
+
+def _rasterization():
+    try:
+        from gsplat import rasterization
+    except ImportError:
+        try:
+            from gsplat.rendering import rasterization
+        except ImportError as err:
+            raise ImportError(
+                'rendering needs gsplat: pip install "ffgs[render]"'
+            ) from err
+    return rasterization
+
+
+def render(
+    gaussians: Gaussians,
+    cameras: Cameras,
+    *,
+    near: float | torch.Tensor,
+    far: float | torch.Tensor,
+    background: Sequence[float] | None = (0.5, 0.5, 0.5),
+    render_depth: bool = False,
+    radius_clip: float = 0.0,
+    rasterize_mode: str = "classic",
+    clamp: bool = False,
+) -> dict[str, torch.Tensor]:
+    """-> {"images": [B, V, 3, H, W], "alphas": [B, V, 1, H, W], ["depths"]}.
+
+    `near` / `far` are in the units of the frame the Gaussians and cameras share:
+    one value, or one per sample ([B]). `radius_clip` (pixels) and
+    `rasterize_mode` ("classic" / "antialiased") are gsplat's; `clamp` clips the
+    images to [0, 1]. A model's own values come from its processor
+    (`Processor.render_settings`, applied by `GSPipeline.render`).
+    """
+    rasterization = _rasterization()
+    h, w = cameras.image_shape
+    device = gaussians.means.device
+    viewmats = torch.inverse(cameras.c2w.float()).to(device)
+    ks = cameras.pixel_k.to(device)
+    b, v = viewmats.shape[:2]
+
+    means = gaussians.means.contiguous().float()
+    quats = gaussians.quats.contiguous().float()
+    scales = gaussians.scales.contiguous().float()
+    opacities = gaussians.opacities.contiguous().float()
+    colors = gaussians.colors.contiguous().float()
+    kwargs = {}
+    if gaussians.sh_degree is not None:
+        kwargs["sh_degree"] = gaussians.sh_degree
+    backgrounds = None
+    if background is not None:
+        backgrounds = torch.tensor(background, dtype=torch.float32, device=device)
+        backgrounds = backgrounds[None]  # [1, 3]: one camera per call
+
+    images, alphas, depths = [], [], []
+    for i in range(b):
+        # One camera per call: for SH colours gsplat materialises the coefficients
+        # per camera ([C, N, K, 3]), so memory would grow with the number of
+        # cameras. Cameras are rasterised independently; the images are the same.
+        per_camera = [
+            rasterization(
+                means=means[i],
+                quats=quats[i],
+                scales=scales[i],
+                opacities=opacities[i],
+                colors=colors[i],
+                viewmats=viewmats[i, j : j + 1],
+                Ks=ks[i, j : j + 1],
+                width=w,
+                height=h,
+                near_plane=_per_sample(near, i),
+                far_plane=_per_sample(far, i),
+                packed=False,
+                backgrounds=backgrounds,
+                render_mode="RGB+ED" if render_depth else "RGB",
+                radius_clip=radius_clip,
+                rasterize_mode=rasterize_mode,
+                **kwargs,
+            )[:2]
+            for j in range(v)
+        ]
+        rendered = torch.cat([r for r, _ in per_camera])
+        alpha = torch.cat([a for _, a in per_camera])
+        image = rendered[..., :3]
+        if clamp:
+            image = image.clamp(0.0, 1.0)
+        images.append(image.permute(0, 3, 1, 2))
+        alphas.append(alpha.permute(0, 3, 1, 2))
+        if render_depth:
+            depths.append(rendered[..., 3:].permute(0, 3, 1, 2))
+
+    out = {"images": torch.stack(images), "alphas": torch.stack(alphas)}
+    if render_depth:
+        out["depths"] = torch.stack(depths)
+    return out
+
+
+def _per_sample(value: float | torch.Tensor, index: int) -> float:
+    if isinstance(value, torch.Tensor):
+        return float(value.reshape(-1)[index] if value.numel() > 1 else value)
+    return value