File size: 1,359 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
from .optimizer import Optimizer, OptimizerCfg
from .optimizer_learn2splat import Learn2SplatOptimizer
from .optimizer_knn_based import KnnBasedOptimizer, KnnBasedOptimizerCfg
from .optimizer_resplat import ResplatOptimizerV1, ResplatOptimizerV2
from .optimizer_adam import AdamOptimizerCfg, AdamOptimizer

SceneOptimizerCfg = KnnBasedOptimizerCfg | AdamOptimizerCfg


def extract_opt_params(cfg):
    opt_params = cfg.__dict__.copy()
    opt_params.pop("enabled", None)
    opt_params.pop("name", None)
    opt_params.pop("steps", None)
    opt_params.pop("scheduler", None)
    opt_params.pop("scheduler_warm_up_ratio", None)
    opt_params.pop("compute_metrics_every", None)
    return opt_params


SCENE_OPTIMIZERS = {
    "none": None,
    "depthsplat": KnnBasedOptimizer,
    "resplat_v1": ResplatOptimizerV1,
    "resplat_v2": ResplatOptimizerV2,
    "clogs": Learn2SplatOptimizer,  # TODO (release): remove
    "l2s": Learn2SplatOptimizer,
    "adam": AdamOptimizer,
}


def get_scene_optimizer(cfg: SceneOptimizerCfg | None) -> Optimizer | None:
    if cfg is None:
        print("Using scene optimizer: None")
        return None
    print(f"Using scene optimizer: {cfg.name}")
    scene_optimizer = SCENE_OPTIMIZERS[cfg.name]
    if scene_optimizer is None:
        return None
    scene_optimizer = scene_optimizer(cfg)
    return scene_optimizer