Spaces:
Sleeping
Sleeping
File size: 3,707 Bytes
78d2329 | 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 | 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
@dataclass
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
)
|