File size: 1,273 Bytes
8bc0782 |
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 |
from enum import Enum
from .optimizer import PointCloudOptimizer
from .modular_optimizer import ModularPointCloudOptimizer
from .pair_viewer import PairViewer
from mini_dust3r.inference import Dust3rResult
from typing import Literal
class GlobalAlignerMode(Enum):
PointCloudOptimizer = "PointCloudOptimizer"
ModularPointCloudOptimizer = "ModularPointCloudOptimizer"
PairViewer = "PairViewer"
def global_aligner(
dust3r_output: Dust3rResult,
device: Literal["cpu", "cuda", "mps"],
mode: GlobalAlignerMode = GlobalAlignerMode.PointCloudOptimizer,
**optim_kw,
):
# extract all inputs
view1, view2, pred1, pred2 = [
dust3r_output[k] for k in "view1 view2 pred1 pred2".split()
]
# build the optimizer
if mode == GlobalAlignerMode.PointCloudOptimizer:
net = PointCloudOptimizer(view1, view2, pred1, pred2, **optim_kw).to(device)
elif mode == GlobalAlignerMode.ModularPointCloudOptimizer:
net = ModularPointCloudOptimizer(view1, view2, pred1, pred2, **optim_kw).to(
device
)
elif mode == GlobalAlignerMode.PairViewer:
net = PairViewer(view1, view2, pred1, pred2, **optim_kw).to(device)
else:
raise NotImplementedError(f"Unknown mode {mode}")
return net |