Spaces:
Sleeping
Sleeping
| from dataclasses import dataclass | |
| from pathlib import Path | |
| from typing import Literal, Optional | |
| from optgs.dataset.data_types import BatchedViews | |
| import torch | |
| import torch.nn.functional as F | |
| # from optgs.dataset.colmap.utils import Parser | |
| # from optgs.experimental.initializers_utils import knn, points_to_gaussians | |
| from optgs.scene_trainer.common.gaussian_adapter import build_covariance | |
| from optgs.model.types import Gaussians | |
| from optgs.model.ply_export import load_gaussians_ply | |
| from optgs.scene_trainer.initializer.initializer import NonlearnedInitializer, InitializerOutput, NonlearnedInitializerCfg | |
| class InitializerPlyCfg(NonlearnedInitializerCfg): | |
| name: Literal["ply"] | |
| path: Path | |
| # normalize_world_space: bool | |
| # scaling_factor: float | |
| # init_opacity: float | |
| sh_degree: int | |
| # dl3dv_settings: bool | |
| ply_filename: str = "gaussians.ply" # relative path under the scene dir; can include subdirs, e.g. "iteration_20000/point_cloud.ply" | |
| def get_gaussian_param_num(self): | |
| # calculate the number of parameters per Gaussian | |
| sh_d = self.get_sh_d() | |
| init_gaussian_param_num = 3 + 4 + 3 * sh_d + 2 + 1 | |
| return init_gaussian_param_num | |
| def get_sh_d(self): | |
| sh_d = (self.sh_degree + 1) ** 2 | |
| return sh_d | |
| class InitializerPly(NonlearnedInitializer[InitializerPlyCfg]): | |
| def __init__(self, cfg: InitializerPlyCfg) -> None: | |
| super().__init__(cfg) | |
| def forward( | |
| self, | |
| context: BatchedViews, | |
| visualization_dump: Optional[dict] = None, | |
| **kwargs | |
| ) -> InitializerOutput: | |
| device = context["extrinsics"].device | |
| verbose = False | |
| # assert COLMAP dir exists | |
| if not self.cfg.path.exists(): | |
| raise ValueError(f"COLMAP dir {self.cfg.path} does not exist.") | |
| if "scene" in kwargs: | |
| scene_name = kwargs["scene"] | |
| assert len(scene_name) == 1, f"Only single scene initialization supported. {scene_name}" | |
| scene_name = scene_name[0] | |
| if verbose: | |
| print(f"Initializing scene '{scene_name}' from COLMAP at {self.cfg.path}.") | |
| datadir = self.cfg.path / scene_name | |
| if not datadir.exists(): | |
| raise ValueError(f"COLMAP scene dir {datadir} does not exist.") | |
| else: | |
| scene_name = None | |
| datadir = self.cfg.path | |
| # ply_filename supports {scene} substitution and glob patterns. The glob | |
| # is matched relative to datadir; exactly one match is required. | |
| rel = self.cfg.ply_filename | |
| if scene_name is not None and "{scene}" in rel: | |
| rel = rel.replace("{scene}", scene_name) | |
| if any(c in rel for c in "*?["): | |
| matches = sorted(datadir.glob(rel)) | |
| if not matches: | |
| raise FileNotFoundError(f"No PLY matched pattern {rel!r} under {datadir}.") | |
| if len(matches) > 1: | |
| raise ValueError(f"PLY pattern {rel!r} matched {len(matches)} files under {datadir}; expected one. Matches: {matches}") | |
| ply_path = matches[0] | |
| else: | |
| ply_path = datadir / rel | |
| # pre-activation values on device | |
| gaussians = load_gaussians_ply(ply_path, max_sh_degree=self.cfg.sh_degree) | |
| # move to device | |
| gaussians = gaussians.to(device) | |
| # build covariances | |
| covariances = build_covariance(scale=gaussians.scales, rotation_xyzw=gaussians.rotations) | |
| gaussians.covariances = covariances | |
| return InitializerOutput( | |
| gaussians=gaussians, | |
| features=None, | |
| depths=None | |
| ) | |