splatviz/splatviz.py at main · Florian-Barthel/splatviz · GitHub

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
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
import numpy as np
import torch
from imgui_bundle import imgui

torch.set_printoptions(precision=2, sci_mode=False)
np.set_printoptions(precision=2)
from renderer.renderer_wrapper import RendererWrapper
from renderer.gaussian_renderer import GaussianRenderer
from renderer.attach_renderer import AttachRenderer
from splatviz_utils.gui_utils import imgui_window
from splatviz_utils.gui_utils.window_helper import WindowHelper
from splatviz_utils.dict_utils import EasyDict
from widgets import (
    edit,
    eval,
    performance,
    load_ply,
    camera,
    save,
    render,
    training,
    video,
    features,
)


class Splatviz(WindowHelper, imgui_window.ImguiWindow):
    def __init__(self, mode, host, port):
        self.code_font_path = "resources/fonts/jetbrainsmono/JetBrainsMono-Regular.ttf"
        self.regular_font_path = "resources/fonts/source_sans_pro/SourceSansPro-Regular.otf"

        super().__init__(
            title="splatviz",
            window_width=1920,
            window_height=1080,
            font=self.regular_font_path,
            code_font=self.code_font_path,
        )

        self.code_font = imgui.get_io().fonts.add_font_from_file_ttf(self.code_font_path, 14)
        self.regular_font = imgui.get_io().fonts.add_font_from_file_ttf(self.code_font_path, 14)
        self._imgui_renderer.refresh_font_texture()

        self.widgets = []
        update_all_the_time = False
        if mode == "default":
            self.widgets = [
                load_ply.LoadWidget(self),
                camera.CamWidget(self),
                video.VideoWidget(self),
                performance.PerformanceWidget(self),
                save.CaptureWidget(self),
                render.RenderWidget(self),
                features.FeatureWidget(self),
                edit.EditWidget(self),
                eval.EvalWidget(self),
            ]
            renderer = GaussianRenderer()
        elif mode == "attach":
            self.widgets = [
                camera.CamWidget(self),
                video.VideoWidget(self),
                performance.PerformanceWidget(self),
                render.RenderWidget(self),
                features.FeatureWidget(self),
                edit.EditWidget(self),
                training.TrainingWidget(self),
            ]
            renderer = AttachRenderer(host=host, port=port)
            update_all_the_time = True
        else:
            raise NotImplementedError(f"Mode '{mode}' not recognized.")

        self.renderer = RendererWrapper(renderer, update_all_the_time)
        self.init_window_helper()

        # Widget interface.
        self.args = EasyDict()
        self.result = EasyDict()
        self.eval_result = ""

        # Initialize window.
        self.set_position(0, 0)
        self.adjust_font_size()
        self.skip_frame()
        self.preprocessed_images = []

    def close(self):
        for widget in self.widgets:
            widget.close()
        super().close()

    def draw_frame(self):
        # main loop
        self.begin_frame()
        self.args = EasyDict()
        self.set_sizes()

        self.draw_widgets_pane()
        self.draw_splitter()
        self.update_render_args(self.args)

        # Render
        if self.is_skipping_frames():
            pass
        else:
            self.renderer.set_args(**self.args)
            result = self.renderer.result
            if result is not None:
                self.result = result

        self.draw_rendering_pane()

        if "eval" in self.result:
            self.eval_result = self.result.eval
        else:
            self.eval_result = None

        if "preprocessed_images" in self.result:
            self.preprocessed_images = self.result.preprocessed_images

        # End frame.
        self.adjust_font_size()
        self.end_frame()