blanchon commited on
Commit
08d80be
1 Parent(s): f9f1617

🔥 First commit

Browse files
.devcontainer/devcontainer.json ADDED
@@ -0,0 +1,31 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ // For format details, see https://aka.ms/devcontainer.json. For config options, see the
2
+ // README at: https://github.com/devcontainers/templates/tree/main/src/docker-existing-dockerfile
3
+ {
4
+ "name": "Cuda",
5
+
6
+ "image": "test",
7
+
8
+ // Features to add to the dev container. More info: https://containers.dev/features.
9
+ // "features": {},
10
+
11
+ // Use 'forwardPorts' to make a list of ports inside the container available locally.
12
+ "forwardPorts": [7860],
13
+
14
+ // Uncomment the next line to run commands after the container is created.
15
+ // "postCreateCommand": "cat /etc/os-release",
16
+
17
+ // Configure tool-specific properties.
18
+ // "customizations": {},
19
+
20
+ "containerEnv": {
21
+ "NVIDIA_VISIBLE_DEVICES": "0"
22
+ },
23
+
24
+ "runArgs": [
25
+ "--gpus","all",
26
+ "--runtime=nvidia"
27
+ ],
28
+
29
+ // Uncomment to connect as an existing user other than the container default. More info: https://aka.ms/dev-containers-non-root.
30
+ // "remoteUser": "devcontainer"
31
+ }
.gitignore CHANGED
@@ -1,3 +1,10 @@
 
 
 
 
 
 
 
1
  # Byte-compiled / optimized / DLL files
2
  __pycache__/
3
  *.py[cod]
 
1
+ ### Others
2
+ deps
3
+ debug
4
+ build
5
+ parameter
6
+
7
+ ### Python
8
  # Byte-compiled / optimized / DLL files
9
  __pycache__/
10
  *.py[cod]
.gitmodules ADDED
@@ -0,0 +1,16 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ [submodule "deps/gaussian-splatting-cuda"]
2
+ path = deps/gaussian-splatting-cuda
3
+ url = git@github.com:MrNeRF/gaussian-splatting-cuda.git
4
+ branch = master
5
+ [submodule "deps/colmap"]
6
+ path = deps/colmap
7
+ url = git@github.com:colmap/colmap.git
8
+ branch = main
9
+ [submodule "deps/rerun"]
10
+ path = deps/rerun
11
+ url = git@github.com:rerun-io/rerun.git
12
+ branch = release-0.8.2
13
+ [submodule "deps/splat"]
14
+ path = deps/splat
15
+ url = git@github.com:antimatter15/splat.git
16
+ branch = main
Dockerfile ADDED
@@ -0,0 +1,145 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # --- `colmap` Builder Stage ---
2
+ FROM nvidia/cuda:11.7.1-devel-ubuntu20.04 AS colmap_builder
3
+
4
+ ARG COLMAP_GIT_COMMIT=main
5
+ ARG CUDA_ARCHITECTURES=native
6
+ ENV QT_XCB_GL_INTEGRATION=xcb_egl
7
+
8
+ WORKDIR /workdir
9
+
10
+ # Prepare and empty machine for building.
11
+ RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y \
12
+ git \
13
+ cmake \
14
+ ninja-build \
15
+ build-essential \
16
+ libboost-program-options-dev \
17
+ libboost-filesystem-dev \
18
+ libboost-graph-dev \
19
+ libboost-system-dev \
20
+ libeigen3-dev \
21
+ libflann-dev \
22
+ libfreeimage-dev \
23
+ libmetis-dev \
24
+ libgoogle-glog-dev \
25
+ libgtest-dev \
26
+ libsqlite3-dev \
27
+ libglew-dev \
28
+ qtbase5-dev \
29
+ libqt5opengl5-dev \
30
+ libcgal-dev \
31
+ libceres-dev \
32
+ && rm -rf /var/lib/apt/lists/*
33
+
34
+ # Build and install COLMAP.
35
+ COPY deps/colmap /colmap
36
+ RUN cd /colmap && \
37
+ mkdir build && \
38
+ cd build && \
39
+ cmake .. -GNinja -DCMAKE_CUDA_ARCHITECTURES=${CUDA_ARCHITECTURES} && \
40
+ ninja && \
41
+ ninja install && \
42
+ cd .. && rm -rf colmap
43
+
44
+ # # --- `gaussian-splatting-cuda` Builder Stage ---
45
+ FROM nvidia/cuda:11.7.1-devel-ubuntu20.04 AS gs_builder
46
+
47
+ WORKDIR /workdir
48
+
49
+ # Install dependencies
50
+ # we could pin them to specific versions to be extra sure
51
+ RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y \
52
+ git \
53
+ python3-dev \
54
+ libtbb-dev \
55
+ libeigen3-dev \
56
+ unzip \
57
+ g++ \
58
+ libssl-dev \
59
+ build-essential \
60
+ checkinstall \
61
+ wget \
62
+ cmake \
63
+ protobuf-compiler \
64
+ && rm -rf /var/lib/apt/lists/*
65
+
66
+ # Install cmake 3.25
67
+ # RUN apt-get update && apt-get -y install
68
+ RUN wget https://github.com/Kitware/CMake/releases/download/v3.25.0/cmake-3.25.0.tar.gz \
69
+ && tar -zvxf cmake-3.25.0.tar.gz \
70
+ && cd cmake-3.25.0 \
71
+ && ./bootstrap \
72
+ && make -j8 \
73
+ && checkinstall --pkgname=cmake --pkgversion="3.25-custom" --default
74
+
75
+ # Copy necessary files
76
+ COPY deps/gaussian-splatting-cuda/cuda_rasterizer ./cuda_rasterizer
77
+ COPY deps/gaussian-splatting-cuda/external ./external
78
+ COPY deps/gaussian-splatting-cuda/includes ./includes
79
+ COPY deps/gaussian-splatting-cuda/parameter ./parameter
80
+ COPY deps/gaussian-splatting-cuda/src ./src
81
+ COPY deps/gaussian-splatting-cuda/CMakeLists.txt ./CMakeLists.txt
82
+
83
+ # Download and extract libtorch
84
+ RUN wget https://download.pytorch.org/libtorch/cu118/libtorch-cxx11-abi-shared-with-deps-2.0.1%2Bcu118.zip \
85
+ && unzip -o libtorch-cxx11-abi-shared-with-deps-2.0.1+cu118.zip -d external/ \
86
+ && rm libtorch-cxx11-abi-shared-with-deps-2.0.1+cu118.zip
87
+
88
+ # Build (on CPU, this will add compute_35 as build target, which we do not want)
89
+ ENV PATH /usr/local/cuda/bin:$PATH
90
+ ENV LD_LIBRARY_PATH /usr/local/cuda/lib64:$LD_LIBRARY_PATH
91
+ RUN cmake -B build -D CMAKE_BUILD_TYPE=Release -D CUDA_TOOLKIT_ROOT_DIR=/usr/local/cuda/ -D CUDA_VERSION=11.7 \
92
+ && cmake --build build -- -j8
93
+
94
+ # --- Runner Stage ---
95
+ FROM nvidia/cuda:11.7.1-devel-ubuntu20.04 AS runner
96
+
97
+ WORKDIR /app
98
+
99
+ RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y \
100
+ libboost-program-options-dev \
101
+ libboost-filesystem-dev \
102
+ libboost-graph-dev \
103
+ libboost-system-dev \
104
+ libeigen3-dev \
105
+ libflann-dev \
106
+ libfreeimage-dev \
107
+ libmetis-dev \
108
+ libgoogle-glog-dev \
109
+ libgtest-dev \
110
+ libsqlite3-dev \
111
+ libglew-dev \
112
+ qtbase5-dev \
113
+ libqt5opengl5-dev \
114
+ libcgal-dev \
115
+ libceres-dev \
116
+ imagemagick \
117
+ ffmpeg \
118
+ python3-pip \
119
+ && rm -rf /var/lib/apt/lists/*
120
+
121
+ # Copy built artifact from colmap_builder stage
122
+ COPY --from=colmap_builder /usr/local/bin/colmap /usr/local/bin/colmap
123
+
124
+ # Copy built artifact from builder stage
125
+ COPY --from=gs_builder /workdir/build/gaussian_splatting_cuda /usr/local/bin/gaussian_splatting_cuda
126
+ COPY --from=gs_builder /workdir/external/libtorch /usr/local/libtorch
127
+ COPY --from=gs_builder /workdir/parameter /usr/local/bin/parameter
128
+
129
+ # Setup environment
130
+ ENV PATH /usr/local/libtorch/bin:/usr/local/cuda/bin:$PATH
131
+ ENV LD_LIBRARY_PATH /usr/local/libtorch/lib:/usr/local/cuda/lib64:$LD_LIBRARY_PATH
132
+
133
+ # Install python dependencies
134
+ COPY requirements.txt /app/requirements.txt
135
+ RUN python3 -m pip install --upgrade pip
136
+ RUN python3 -m pip install -r /app/requirements.txt
137
+
138
+ COPY services /app/services
139
+ COPY server.py /app/server.py
140
+
141
+ # Fix bug
142
+ RUN mkdir /parameter && cp /usr/local/bin/parameter/optimization_params.json /parameter/optimization_params.json
143
+
144
+ EXPOSE 7860
145
+ CMD [ "python3", "-u", "/app/server.py" ]
deps/colmap ADDED
@@ -0,0 +1 @@
 
 
1
+ Subproject commit c04629017e7378b3046c6e8961277fbe98b56a32
deps/gaussian-splatting-cuda ADDED
@@ -0,0 +1 @@
 
 
1
+ Subproject commit b3aced9a3c80bed0e072f31540bcf9919cf1eb1d
deps/rerun ADDED
@@ -0,0 +1 @@
 
 
1
+ Subproject commit 9b76b7027b5cc34bf86b07c41968eb0988b383a3
deps/splat ADDED
@@ -0,0 +1 @@
 
 
1
+ Subproject commit db27473c34b21e9294bd80848380660808b21a4e
requirements.txt ADDED
@@ -0,0 +1,8 @@
 
 
 
 
 
 
 
 
 
1
+ requests
2
+ numpy
3
+ typing_extensions
4
+ rich
5
+ fastapi
6
+ uvicorn[standard]
7
+ gradio
8
+ # rerun-sdk==0.8.2 # if you want to use the rerun
server.py ADDED
@@ -0,0 +1,422 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from pathlib import Path
2
+ import shutil
3
+ import tempfile
4
+ import gradio as gr
5
+ import uuid
6
+ from typing_extensions import TypedDict, Tuple
7
+
8
+ from fastapi import FastAPI
9
+ from fastapi.staticfiles import StaticFiles
10
+ import uvicorn
11
+
12
+ app = FastAPI()
13
+
14
+ # create a static directory to store the static files
15
+ gs_dir = Path(str(tempfile.gettempdir())) / "gaussian_splatting_gradio"
16
+ gs_dir.mkdir(parents=True, exist_ok=True)
17
+
18
+ # mount FastAPI StaticFiles server
19
+ app.mount("/static", StaticFiles(directory=gs_dir), name="static")
20
+
21
+ StateDict = TypedDict("StateDict", {
22
+ "uuid": str,
23
+ })
24
+
25
+ def getHTML():
26
+ html_body = """
27
+ <body>
28
+ <div id="progress"></div>
29
+ <div id="message"></div>
30
+ <div class="scene" id="spinner">
31
+ <div class="cube-wrapper">
32
+ <div class="cube">
33
+ <div class="cube-faces">
34
+ <div class="cube-face bottom"></div>
35
+ <div class="cube-face top"></div>
36
+ <div class="cube-face left"></div>
37
+ <div class="cube-face right"></div>
38
+ <div class="cube-face back"></div>
39
+ <div class="cube-face front"></div>
40
+ </div>
41
+ </div>
42
+ </div>
43
+ </div>
44
+ <canvas id="canvas"></canvas>
45
+
46
+ <div id="quality">
47
+ <span id="fps"></span>
48
+ </div>
49
+
50
+ <style>
51
+ .cube-wrapper {
52
+ transform-style: preserve-3d;
53
+ }
54
+
55
+ .cube {
56
+ transform-style: preserve-3d;
57
+ transform: rotateX(45deg) rotateZ(45deg);
58
+ animation: rotation 2s infinite;
59
+ }
60
+
61
+ .cube-faces {
62
+ transform-style: preserve-3d;
63
+ height: 80px;
64
+ width: 80px;
65
+ position: relative;
66
+ transform-origin: 0 0;
67
+ transform: translateX(0) translateY(0) translateZ(-40px);
68
+ }
69
+
70
+ .cube-face {
71
+ position: absolute;
72
+ inset: 0;
73
+ background: #0017ff;
74
+ border: solid 1px #ffffff;
75
+ }
76
+ .cube-face.top {
77
+ transform: translateZ(80px);
78
+ }
79
+ .cube-face.front {
80
+ transform-origin: 0 50%;
81
+ transform: rotateY(-90deg);
82
+ }
83
+ .cube-face.back {
84
+ transform-origin: 0 50%;
85
+ transform: rotateY(-90deg) translateZ(-80px);
86
+ }
87
+ .cube-face.right {
88
+ transform-origin: 50% 0;
89
+ transform: rotateX(-90deg) translateY(-80px);
90
+ }
91
+ .cube-face.left {
92
+ transform-origin: 50% 0;
93
+ transform: rotateX(-90deg) translateY(-80px) translateZ(80px);
94
+ }
95
+
96
+ @keyframes rotation {
97
+ 0% {
98
+ transform: rotateX(45deg) rotateY(0) rotateZ(45deg);
99
+ animation-timing-function: cubic-bezier(
100
+ 0.17,
101
+ 0.84,
102
+ 0.44,
103
+ 1
104
+ );
105
+ }
106
+ 50% {
107
+ transform: rotateX(45deg) rotateY(0) rotateZ(225deg);
108
+ animation-timing-function: cubic-bezier(
109
+ 0.76,
110
+ 0.05,
111
+ 0.86,
112
+ 0.06
113
+ );
114
+ }
115
+ 100% {
116
+ transform: rotateX(45deg) rotateY(0) rotateZ(405deg);
117
+ animation-timing-function: cubic-bezier(
118
+ 0.17,
119
+ 0.84,
120
+ 0.44,
121
+ 1
122
+ );
123
+ }
124
+ }
125
+
126
+ .scene,
127
+ #message {
128
+ position: absolute;
129
+ display: flex;
130
+ top: 0;
131
+ right: 0;
132
+ left: 0;
133
+ bottom: 0;
134
+ z-index: 2;
135
+ height: 100%;
136
+ width: 100%;
137
+ align-items: center;
138
+ justify-content: center;
139
+ }
140
+ #message {
141
+ font-weight: bold;
142
+ font-size: large;
143
+ color: red;
144
+ pointer-events: none;
145
+ }
146
+
147
+ #progress {
148
+ position: absolute;
149
+ top: 0;
150
+ height: 5px;
151
+ background: blue;
152
+ z-index: 99;
153
+ transition: width 0.1s ease-in-out;
154
+ }
155
+
156
+ #quality {
157
+ position: absolute;
158
+ bottom: 10px;
159
+ z-index: 999;
160
+ right: 10px;
161
+ }
162
+
163
+ #canvas {
164
+ display: block;
165
+ position: absolute;
166
+ top: 0;
167
+ left: 0;
168
+ width: 100%;
169
+ height: 100%;
170
+ touch-action: none;
171
+ }
172
+
173
+ #instructions {
174
+ background: rgba(0,0,0,0.6);
175
+ white-space: pre-wrap;
176
+ padding: 10px;
177
+ border-radius: 10px;
178
+ font-size: x-small;
179
+ }
180
+ </style>
181
+ </body>
182
+ """
183
+
184
+ html = f"""
185
+ <head>
186
+ <title>3D Gaussian Splatting Viewer</title>
187
+ <script src="http://zeus.blanchon.cc/dropshare/main.js"></script>
188
+ </head>
189
+
190
+ {html_body}
191
+ """
192
+ return f"""<iframe style="width: 100%; height: 900px" srcdoc='{html}'></iframe>"""
193
+
194
+ def createStateSession() -> StateDict:
195
+ # Create new session
196
+ session_uuid = str(uuid.uuid4())
197
+ print("createStateSession")
198
+ print(session_uuid)
199
+ return StateDict(
200
+ uuid=session_uuid,
201
+ )
202
+
203
+ def removeStateSession(session_state_value: StateDict):
204
+ # Clean up previous session
205
+ return StateDict(
206
+ uuid=None,
207
+ )
208
+
209
+ def makeButtonVisible() -> Tuple[gr.Button, gr.Button]:
210
+ process_button = gr.Button(visible=True)
211
+ reset_button = gr.Button(visible=False) #TODO: I will bring this back when I figure out how to stop the process
212
+ return process_button, reset_button
213
+
214
+ def resetSession(state: StateDict) -> Tuple[StateDict, gr.Button, gr.Button]:
215
+ print("resetSession")
216
+ new_state = removeStateSession(state)
217
+ process_button = gr.Button(visible=False)
218
+ reset_button = gr.Button(visible=False)
219
+ return new_state, process_button, reset_button
220
+
221
+ def process(
222
+ # *args, **kwargs
223
+ session_state_value: StateDict,
224
+ filepath: str,
225
+ ffmpeg_fps: int,
226
+ ffmpeg_qscale: int,
227
+ colmap_camera: str,
228
+ ):
229
+ if session_state_value["uuid"] is None:
230
+ return
231
+ print("process")
232
+ # print(args)
233
+ # print(kwargs)
234
+ # return
235
+ print(session_state_value)
236
+ print(f"Processing {filepath}")
237
+
238
+ try:
239
+ session_tmpdirname = gs_dir / str(session_state_value['uuid'])
240
+ session_tmpdirname.mkdir(parents=True, exist_ok=True)
241
+ print('Created temporary directory', session_tmpdirname)
242
+
243
+ gs_dir_path = Path(session_tmpdirname)
244
+ logfile_path = Path(session_tmpdirname) / "log.txt"
245
+ logfile_path.touch()
246
+ with logfile_path.open("w") as log_file:
247
+ # Create log file
248
+ logfile_path.touch()
249
+
250
+ from services.ffmpeg import ffmpeg_run
251
+ ffmpeg_run(
252
+ video_path = Path(filepath),
253
+ output_path = gs_dir_path,
254
+ fps = int(ffmpeg_fps),
255
+ qscale = int(ffmpeg_qscale),
256
+ stream_file=log_file
257
+ )
258
+
259
+ from services.colmap import colmap
260
+ colmap(
261
+ source_path=gs_dir_path,
262
+ camera=str(colmap_camera),
263
+ stream_file=log_file
264
+ )
265
+
266
+ print("Done with colmap")
267
+
268
+ # Create a zip of the gs_dir_path folder
269
+ print(gs_dir, gs_dir_path)
270
+ print(gs_dir_path.name)
271
+ archive = shutil.make_archive("result", 'zip', gs_dir, gs_dir_path)
272
+ print('Created zip file', archive)
273
+
274
+ # Move the zip file to the gs_dir_path folder
275
+ shutil.move(archive, gs_dir_path)
276
+
277
+ from services.gaussian_splatting_cuda import gaussian_splatting_cuda
278
+ gaussian_splatting_cuda(
279
+ data_path = gs_dir_path,
280
+ output_path = gs_dir_path / "output",
281
+ gs_command = str(Path(__file__).parent.absolute() / "build" / 'gaussian_splatting_cuda'),
282
+ iterations = 100,
283
+ convergence_rate = 0.01,
284
+ resolution = 512,
285
+ enable_cr_monitoring = False,
286
+ force = False,
287
+ empty_gpu_cache = False,
288
+ stream_file = log_file
289
+ )
290
+
291
+ except Exception:
292
+ pass
293
+ # print('Error - Removing temporary directory', session_tmpdirname)
294
+ # shutil.rmtree(session_tmpdirname)
295
+
296
+ def updateLog(session_state_value: StateDict) -> str:
297
+ if session_state_value["uuid"] is None:
298
+ return ""
299
+
300
+ log_file = gs_dir / str(session_state_value['uuid']) / "log.txt"
301
+ if not log_file.exists():
302
+ return ""
303
+
304
+ with log_file.open("r") as log_file:
305
+ logs = log_file.read()
306
+
307
+ return logs
308
+
309
+ with gr.Blocks() as demo:
310
+ session_state = gr.State({
311
+ "uuid": None,
312
+ })
313
+
314
+ with gr.Row():
315
+
316
+ with gr.Column():
317
+ video_input = gr.PlayableVideo(
318
+ format="mp4",
319
+ source="upload",
320
+ label="Upload a video",
321
+ include_audio=False
322
+ )
323
+ with gr.Row(variant="panel"):
324
+ ffmpeg_fps = gr.Number(
325
+ label="FFMPEG FPE",
326
+ value=1,
327
+ minimum=1,
328
+ maximum=5,
329
+ step=0.10,
330
+ )
331
+ ffmpeg_qscale = gr.Number(
332
+ label="FFMPEG QSCALE",
333
+ value=1,
334
+ minimum=1,
335
+ maximum=5,
336
+ step=1,
337
+ )
338
+ colmap_camera = gr.Dropdown(
339
+ label="COLMAP Camera",
340
+ value="OPENCV",
341
+ choices=["OPENCV", "SIMPLE_PINHOLE", "PINHOLE", "SIMPLE_RADIAL", "RADIAL"],
342
+ )
343
+
344
+ text_log = gr.Textbox(
345
+ label="Logs",
346
+ info="Logs",
347
+ interactive=False,
348
+ show_copy_button=True
349
+ )
350
+ # text_log = gr.Code(
351
+ # label="Logs",
352
+ # language=None,
353
+ # interactive=False,
354
+ # )
355
+
356
+
357
+ process_button = gr.Button("Process", visible=False)
358
+ reset_button = gr.ClearButton(
359
+ components=[video_input, text_log, ffmpeg_fps, ffmpeg_qscale, colmap_camera],
360
+ label="Reset",
361
+ visible=False,
362
+ )
363
+
364
+ process_event = process_button.click(
365
+ fn=process,
366
+ inputs=[session_state, video_input, ffmpeg_fps, ffmpeg_qscale, colmap_camera],
367
+ outputs=[],
368
+ )
369
+
370
+ upload_event = video_input.upload(
371
+ fn=makeButtonVisible,
372
+ inputs=[],
373
+ outputs=[process_button, reset_button]
374
+ ).then(
375
+ fn=createStateSession,
376
+ inputs=[],
377
+ outputs=[session_state],
378
+ ).then(
379
+ fn=updateLog,
380
+ inputs=[session_state],
381
+ outputs=[text_log],
382
+ every=2,
383
+ )
384
+
385
+ reset_button.click(
386
+ fn=resetSession,
387
+ inputs=[session_state],
388
+ outputs=[session_state, process_button, reset_button],
389
+ cancels=[process_event]
390
+ )
391
+
392
+ video_input.clear(
393
+ fn=resetSession,
394
+ inputs=[session_state],
395
+ outputs=[session_state, process_button, reset_button],
396
+ cancels=[process_event]
397
+ )
398
+
399
+ demo.close
400
+
401
+
402
+ # gr.LoginButton, gr.LogoutButton
403
+ # gr.HuggingFaceDatasetSaver
404
+ # gr.OAuthProfile
405
+
406
+
407
+
408
+
409
+
410
+
411
+ # with gr.Tab("jsdn"):
412
+ # input_mic = gr.HTML(getHTML())
413
+
414
+ demo.queue()
415
+ # demo.launch()
416
+
417
+ # mount Gradio app to FastAPI app
418
+ app = gr.mount_gradio_app(app, demo, path="/")
419
+
420
+
421
+ if __name__ == "__main__":
422
+ uvicorn.run(app, host="0.0.0.0", port=7860)
services/colmap.py ADDED
@@ -0,0 +1,244 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Literal, Optional
2
+ from io import IOBase
3
+ import os
4
+ from pathlib import Path
5
+ import shutil
6
+ import subprocess
7
+ from rich.progress import Progress
8
+ from rich.console import Console
9
+
10
+ console = Console()
11
+
12
+ class FailedProcess(Exception):
13
+ pass
14
+
15
+ def colmap_feature_extraction(
16
+ database_path: Path,
17
+ image_path: Path,
18
+ camera: Literal["OPENCV"],
19
+ colmap_command: str = "colmap",
20
+ use_gpu: bool = True,
21
+ stream_file: Optional[IOBase] = None
22
+ ):
23
+ total = len(list(image_path.glob("*.jpg")))
24
+ with Progress(console=console) as progress:
25
+ task = progress.add_task("Feature Extraction", total=total)
26
+
27
+ database_path.parent.mkdir(parents=True, exist_ok=True)
28
+ cmd = [
29
+ colmap_command,
30
+ "feature_extractor",
31
+ "--database_path", database_path.as_posix(),
32
+ "--image_path", image_path.as_posix(),
33
+ "--ImageReader.single_camera", "1",
34
+ "--ImageReader.camera_model", camera,
35
+ "--SiftExtraction.use_gpu", "1" if use_gpu else "0",
36
+ # "--SiftExtraction.domain_size_pooling", "1",
37
+ # "--SiftExtraction.estimate_affine_shape", "1"
38
+ ]
39
+ console.log(f"💻 Executing command: {' '.join(cmd)}")
40
+
41
+ _stdout = stream_file if stream_file else subprocess.PIPE
42
+ with subprocess.Popen(cmd, stdout=_stdout, stderr=subprocess.STDOUT, text=True) as process:
43
+ if process.stdout:
44
+ for line in process.stdout:
45
+ if line.startswith("Processed file "):
46
+ line_process = line\
47
+ .replace("Processed file [", "")\
48
+ .replace("]", "")\
49
+ .replace("\n", "")
50
+ current, total = line_process.split("/")
51
+ progress.update(task, completed=int(current), total=int(total), refresh=True)
52
+
53
+ progress.update(task, completed=int(total), refresh=True)
54
+
55
+ return_code = process.returncode
56
+
57
+ if return_code == 0:
58
+
59
+ console.log(f'Feature stored in {database_path.as_posix()}.')
60
+ console.log('✅ Feature extraction completed.')
61
+ else:
62
+ raise FailedProcess("Feature extraction failed.")
63
+
64
+ def colmap_feature_matching(
65
+ database_path: Path,
66
+ image_path: Path,
67
+ colmap_command: str = "colmap",
68
+ use_gpu: bool = True,
69
+ stream_file: Optional[IOBase] = None
70
+ ):
71
+ total = len(list(image_path.glob("*.jpg")))
72
+ with Progress(console=console) as progress:
73
+ task = progress.add_task("Feature Matching", total=total)
74
+
75
+ database_path
76
+ cmd = [
77
+ colmap_command,
78
+ "exhaustive_matcher",
79
+ "--database_path", database_path.as_posix(),
80
+ "--SiftMatching.use_gpu", "1" if use_gpu else "0"
81
+ ]
82
+ console.log(f"💻 Executing command: {' '.join(cmd)}")
83
+
84
+ _stdout = stream_file if stream_file else subprocess.PIPE
85
+ with subprocess.Popen(cmd, stdout=_stdout, stderr=subprocess.STDOUT, text=True) as process:
86
+ if process.stdout:
87
+ for line in process.stdout:
88
+ pass
89
+
90
+ progress.update(task, completed=int(total), refresh=True)
91
+
92
+ return_code = process.returncode
93
+
94
+ if return_code == 0:
95
+
96
+ console.log('✅ Feature matching completed.')
97
+ else:
98
+ raise FailedProcess("Feature matching failed.")
99
+
100
+ def colmap_bundle_adjustment(
101
+ database_path: Path,
102
+ image_path: Path,
103
+ sparse_path: Path,
104
+ colmap_command: str = "colmap",
105
+ stream_file: Optional[IOBase] = None
106
+ ):
107
+ total = len(list(image_path.glob("*.jpg")))
108
+ with Progress(console=console) as progress:
109
+ task = progress.add_task("Bundle Adjustment", total=total)
110
+
111
+ cmd = [
112
+ colmap_command,
113
+ "mapper",
114
+ "--database_path", database_path.as_posix(),
115
+ "--image_path", image_path.as_posix(),
116
+ "--output_path", sparse_path.as_posix(),
117
+ "--Mapper.ba_global_function_tolerance=0.000001"
118
+ # "--Mapper.ba_local_max_num_iterations", "40",
119
+ # "--Mapper.ba_global_max_num_iterations", "100",
120
+ # "--Mapper.ba_local_max_refinements", "3",
121
+ # "--Mapper.ba_global_max_refinements", "5"
122
+ ]
123
+ console.log(f"💻 Executing command: {' '.join(cmd)}")
124
+
125
+ sparse_path.mkdir(parents=True, exist_ok=True)
126
+
127
+ _stdout = stream_file if stream_file else subprocess.PIPE
128
+ with subprocess.Popen(cmd, stdout=_stdout, stderr=subprocess.STDOUT, text=True) as process:
129
+ if process.stdout:
130
+ for line in process.stdout:
131
+ print(line)
132
+ if line.startswith("Registering image #"):
133
+ line_process = line\
134
+ .replace("Registering image #", "")\
135
+ .replace("\n", "")
136
+ *_, current = line_process.split("(")
137
+ current, *_ = current.split(")")
138
+ progress.update(task, completed=int(current), refresh=True)
139
+
140
+ progress.update(task, completed=int(total), refresh=True)
141
+
142
+ return_code = process.returncode
143
+
144
+ if return_code == 0:
145
+ console.log('✅ Bundle adjustment completed.')
146
+ else:
147
+ raise FailedProcess("Bundle adjustment failed.")
148
+
149
+ def colmap_image_undistortion(
150
+ image_path: Path,
151
+ sparse0_path: Path,
152
+ source_path: Path,
153
+ colmap_command: str = "colmap",
154
+ stream_file: Optional[IOBase] = None
155
+ ):
156
+ total = len(list(image_path.glob("*.jpg")))
157
+ with Progress(console=console) as progress:
158
+ task = progress.add_task("Image Undistortion", total=total)
159
+ cmd = [
160
+ colmap_command,
161
+ "image_undistorter",
162
+ "--image_path", image_path.as_posix(),
163
+ "--input_path", sparse0_path.as_posix(),
164
+ "--output_path", source_path.as_posix(),
165
+ "--output_type", "COLMAP"
166
+ ]
167
+ console.log(f"💻 Executing command: {' '.join(cmd)}")
168
+
169
+ _stdout = stream_file if stream_file else subprocess.PIPE
170
+ with subprocess.Popen(cmd, stdout=_stdout, stderr=subprocess.STDOUT, text=True) as process:
171
+ if process.stdout:
172
+ for line in process.stdout:
173
+ if line.startswith("Undistorting image ["):
174
+ line_process = line\
175
+ .replace("Undistorting image [", "")\
176
+ .replace("]", "")\
177
+ .replace("\n", "")
178
+ current, total = line_process.split("/")
179
+ progress.update(task, completed=int(current), total=int(total), refresh=True)
180
+
181
+ progress.update(task, completed=int(total), refresh=True)
182
+
183
+ return_code = process.returncode
184
+
185
+ if return_code == 0:
186
+ console.log('✅ Image undistortion completed.')
187
+ else:
188
+ raise FailedProcess("Image undistortion failed.")
189
+
190
+ def colmap(
191
+ source_path: Path,
192
+ camera: Literal["OPENCV"] = "OPENCV",
193
+ colmap_command: str = "colmap",
194
+ use_gpu: bool = True,
195
+ skip_matching: bool = False,
196
+ stream_file: Optional[IOBase] = None
197
+ ):
198
+ image_path = source_path / "input"
199
+ if not image_path.exists():
200
+ raise Exception(f"Image path {image_path} does not exist. Exiting.")
201
+
202
+ total = len(list(image_path.glob("*.jpg")))
203
+ if total == 0:
204
+ raise Exception(f"No images found in {image_path}. Exiting.")
205
+
206
+ database_path = source_path / "distorted" / "database.db"
207
+
208
+ sparse_path = source_path / "distorted" / "sparse"
209
+
210
+ if not skip_matching:
211
+ colmap_feature_extraction(database_path, image_path, camera, colmap_command, use_gpu, stream_file)
212
+ colmap_feature_matching(database_path, image_path, colmap_command, use_gpu, stream_file)
213
+ colmap_bundle_adjustment(database_path, image_path, sparse_path, colmap_command, stream_file)
214
+
215
+ colmap_image_undistortion(image_path, sparse_path / "0", source_path, colmap_command, stream_file)
216
+
217
+ origin_path = source_path / "sparse"
218
+ destination_path = source_path / "sparse" / "0"
219
+ destination_path.mkdir(exist_ok=True)
220
+ console.log(f"🌟 Moving files from {origin_path} to {destination_path}")
221
+ for file in os.listdir(origin_path):
222
+ if file == '0':
223
+ continue
224
+ source_file = os.path.join(origin_path, file)
225
+ destination_file = os.path.join(destination_path, file)
226
+ shutil.copy(source_file, destination_file)
227
+
228
+ if __name__ == "__main__":
229
+ import tempfile
230
+ with tempfile.NamedTemporaryFile(mode='w+t') as temp_file:
231
+ print(f"Using temp file: {temp_file.name}")
232
+ try:
233
+ colmap(
234
+ source_path = Path("/home/europe/Desktop/gaussian-splatting-kit/test/"),
235
+ camera = "OPENCV",
236
+ colmap_command = "colmap",
237
+ use_gpu = True,
238
+ skip_matching = False,
239
+ stream_file = open("/home/europe/Desktop/gaussian-splatting-kit/test.log", "w+t")
240
+ )
241
+ except FailedProcess:
242
+ console.log("🚨 Error executing colmap.")
243
+ temp_file.seek(0)
244
+ print(temp_file.read())
services/ffmpeg.py ADDED
@@ -0,0 +1,100 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from io import IOBase
2
+ import os
3
+ import subprocess
4
+ from typing import Optional
5
+ from pathlib import Path
6
+ from rich.console import Console
7
+
8
+ console = Console()
9
+
10
+ class FailedProcess(Exception):
11
+ pass
12
+
13
+ def ffmpeg_extract_frames(
14
+ video_path: Path,
15
+ frames_path: Path,
16
+ # TODO: Enable these options
17
+ # start_time: Optional[str] = None,
18
+ # duration: Optional[float] = None,
19
+ # end_time: Optional[str] = None,
20
+ fps: float = 1,
21
+ qscale: int = 1,
22
+ stream_file: Optional[IOBase] = None
23
+ ) -> str:
24
+ frame_destination = frames_path / "input"
25
+ console.log(f"🎞️ Extracting Images from {video_path} to {frame_destination} (fps: {fps}, qscale: {qscale}")
26
+ # Create the directory to store the frames
27
+ frames_path.mkdir(parents=True, exist_ok=True)
28
+ frame_destination.mkdir(parents=True, exist_ok=True)
29
+ # Store the current working directory
30
+ cwd = os.getcwd()
31
+ # Change the current working directory to frame_destination
32
+ os.chdir(frame_destination)
33
+
34
+ # Construct the ffmpeg command as a list of strings
35
+ cmd = [
36
+ 'ffmpeg',
37
+ '-i', str(video_path),
38
+ '-qscale:v', str(qscale),
39
+ '-qmin', '1',
40
+ '-vf', f"fps={fps}",
41
+ '%04d.jpg'
42
+ ]
43
+
44
+ console.log(f"💻 Executing command: {' '.join(cmd)}")
45
+
46
+ _stdout = stream_file if stream_file else subprocess.PIPE
47
+ with subprocess.Popen(cmd, stdout=_stdout, stderr=subprocess.STDOUT, text=True) as process:
48
+ if process.stdout:
49
+ for line in process.stdout:
50
+ print(line)
51
+
52
+ # Change the current working directory back to the original
53
+ os.chdir(cwd)
54
+
55
+ return_code = process.returncode
56
+
57
+ if return_code == 0:
58
+ console.log(f"✅ Images Successfully Extracted! Path: {frames_path}")
59
+ else:
60
+ raise FailedProcess("Error extracting frames.")
61
+
62
+ return frames_path
63
+
64
+ def ffmpeg_run(
65
+ video_path: Path,
66
+ output_path: Path,
67
+ ffmpeg_command: str = "ffmpeg",
68
+ # TODO: Enable these options
69
+ # start_time: Optional[str] = None,
70
+ # duration: Optional[float] = None,
71
+ # end_time: Optional[str] = None,
72
+ fps: float = 1,
73
+ qscale: int = 1,
74
+ stream_file: Optional[IOBase] = None
75
+ ) -> str:
76
+ console.log("🌟 Starting the Frames Extraction...")
77
+ frames_path = ffmpeg_extract_frames(
78
+ video_path,
79
+ output_path,
80
+ fps=fps, qscale=qscale,
81
+ stream_file=stream_file
82
+ )
83
+ console.log(f"🎉 Frames Extraction Complete! Path: {frames_path}")
84
+ return frames_path
85
+
86
+ if __name__ == "__main__":
87
+ import tempfile
88
+ with tempfile.NamedTemporaryFile(mode='w+t') as temp_file:
89
+ print(f"Using temp file: {temp_file.name}")
90
+ try:
91
+ ffmpeg_run(
92
+ Path("/home/europe/Desktop/gaussian-splatting-kit/test/test.mov"),
93
+ Path("/home/europe/Desktop/gaussian-splatting-kit/test"),
94
+ stream_file=temp_file
95
+ )
96
+ except FailedProcess:
97
+ console.log("🚨 Error extracting frames.")
98
+ temp_file.seek(0)
99
+ print(temp_file.read())
100
+
services/gaussian_splatting_cuda.py ADDED
@@ -0,0 +1,108 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from io import IOBase
2
+ from pathlib import Path
3
+ import subprocess
4
+ from typing import Optional
5
+ from rich.console import Console
6
+
7
+ console = Console()
8
+
9
+ def gaussian_splatting_cuda_training(
10
+ data_path: Path,
11
+ output_path: Path,
12
+ gs_command: str,
13
+ iterations: int = 10_000,
14
+ convergence_rate: float = 0.01,
15
+ resolution: int = 512,
16
+ enable_cr_monitoring: bool = False,
17
+ force: bool = False,
18
+ empty_gpu_cache: bool = False,
19
+ stream_file: Optional[IOBase] = None
20
+ ) -> str:
21
+ """
22
+ Core Options
23
+ -h, --help
24
+ Display this help menu.
25
+
26
+ -d, --data_path [PATH]
27
+ Specify the path to the training data.
28
+
29
+ -f, --force
30
+ Force overwriting of output folder. If not set, the program will exit if the output folder already exists.
31
+
32
+ -o, --output_path [PATH]
33
+ Specify the path to save the trained model. If this option is not specified, the trained model will be saved to the "output" folder located in the root directory of the project.
34
+
35
+ -i, --iter [NUM]
36
+ Specify the number of iterations to train the model. Although the paper sets the maximum number of iterations at 30k, you'll likely need far fewer. Starting with 6k or 7k iterations should yield preliminary results. Outputs are saved every 7k iterations and also at the end of the training. Therefore, even if you set it to 5k iterations, an output will be generated upon completion.
37
+
38
+ Advanced Options
39
+ --empty-gpu-cache Empty CUDA memory after ever 100 iterations. Attention! This has a considerable performance impact
40
+
41
+ --enable-cr-monitoring
42
+ Enable monitoring of the average convergence rate throughout training. If done, it will stop optimizing when the average convergence rate is below 0.008 per default after 15k iterations. This is useful for speeding up the training process when the gain starts to dimish. If not enabled, the training will stop after the specified number of iterations --iter. Otherwise its stops when max 30k iterations are reached.
43
+
44
+ -c, --convergence_rate [RATE]
45
+ Set custom average onvergence rate for the training process. Requires the flag --enable-cr-monitoring to be set.
46
+ """
47
+
48
+ cmd = [
49
+ gs_command,
50
+ f"--data-path={data_path.as_posix()}"
51
+ f"--output-path={output_path.as_posix()}"
52
+ f"--iter={iterations}",
53
+ # TODO: Enable these options and put the right defaults in the function signature
54
+ # f"--convergence-rate={convergence_rate}",
55
+ # f"--resolution={resolution}",
56
+ # "--enable-cr-monitoring" if enable_cr_monitoring else "",
57
+ # "--force" if force else "",
58
+ # "--empty-gpu-cache" if empty_gpu_cache else ""
59
+ ]
60
+
61
+ console.log(f"💻 Executing command: {' '.join(cmd)}")
62
+
63
+ _stdout = stream_file if stream_file else subprocess.PIPE
64
+ with subprocess.Popen(cmd, stdout=_stdout, stderr=subprocess.STDOUT, text=True) as process:
65
+ if process.stdout:
66
+ for line in process.stdout:
67
+ print(line)
68
+
69
+
70
+ # Check if the command was successful
71
+ return_code = process.returncode
72
+ if return_code == 0:
73
+ console.log('✅ Successfully splatted frames.')
74
+ else:
75
+ raise Exception('Error splatting frames.')
76
+
77
+ def gaussian_splatting_cuda(
78
+ data_path: Path,
79
+ output_path: Path,
80
+ gs_command: str,
81
+ iterations: int = 10_000,
82
+ convergence_rate: float = 0.01,
83
+ resolution: int = 512,
84
+ enable_cr_monitoring: bool = False,
85
+ force: bool = False,
86
+ empty_gpu_cache: bool = False,
87
+ stream_file: Optional[IOBase] = None
88
+ ) -> str:
89
+ # Check if the output path exists
90
+ if output_path.exists() and not force:
91
+ raise Exception(f"Output folder already exists. Path: {output_path}, use --force to overwrite.")
92
+
93
+ # Create the output path if it doesn't exist
94
+ output_path.mkdir(parents=True, exist_ok=True)
95
+
96
+ # Execute gaussian_splatting_cuda
97
+ gaussian_splatting_cuda_training(
98
+ data_path,
99
+ output_path,
100
+ gs_command,
101
+ iterations,
102
+ convergence_rate,
103
+ resolution,
104
+ enable_cr_monitoring,
105
+ force,
106
+ empty_gpu_cache,
107
+ stream_file
108
+ )
services/http.py ADDED
@@ -0,0 +1,26 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from pathlib import Path
2
+ import requests
3
+ from rich.console import Console
4
+
5
+ console = Console()
6
+
7
+ def download_file(url: str, file_path: Path) -> Path:
8
+ console.log(f"📥 Downloading File from URL: {url}")
9
+ response = requests.get(url, stream=True)
10
+ if response.status_code == 200:
11
+ with file_path.open('wb') as file:
12
+ for chunk in response.iter_content(chunk_size=1024):
13
+ if chunk:
14
+ file.write(chunk)
15
+ console.log(f"✅ File Successfully Downloaded! Path: {file_path}")
16
+ else:
17
+ console.log(f"🚨 Error downloading file from {url}.")
18
+ return file_path
19
+
20
+ def download_api(url: str, file_path: Path) -> Path:
21
+ # Download the video from internet
22
+ video_path = file_path + '/video.mp4'
23
+ console.log("🌟 Starting the Video Download...")
24
+ video_path = download_file(url, video_path)
25
+ console.log(f"🎉 Video Download Complete! Path: {video_path}")
26
+ return video_path
services/rerun.py ADDED
@@ -0,0 +1,85 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import re
3
+ from pathlib import Path
4
+
5
+ import numpy as np
6
+ import rerun as rr # pip install rerun-sdk
7
+ from utils.read_write_model import read_model
8
+
9
+ # From https://github.com/rerun-io/rerun/tree/main/examples/python/structure_from_motion
10
+ def read_and_log_sparse_reconstruction(
11
+ exp_name: str,
12
+ dataset_path: Path,
13
+ output_path: Path,
14
+ filter_output: bool = False,
15
+ filter_min_visible: int = 2_000
16
+ ) -> None:
17
+ rr.init(exp_name)
18
+
19
+ cameras, images, points3D = read_model(dataset_path / "sparse", ext=".bin")
20
+
21
+ if filter_output:
22
+ # Filter out noisy points
23
+ points3D = {id: point for id, point in points3D.items() if point.rgb.any() and len(point.image_ids) > 4}
24
+
25
+ rr.log_view_coordinates("/", up="-Y", timeless=True)
26
+
27
+ # Iterate through images (video frames) logging data related to each frame.
28
+ for image in sorted(images.values(), key=lambda im: im.name): # type: ignore[no-any-return]
29
+ image_file = dataset_path / "images" / image.name
30
+
31
+ if not os.path.exists(image_file):
32
+ continue
33
+
34
+ # COLMAP sets image ids that don't match the original video frame
35
+ idx_match = re.search(r"\d+", image.name)
36
+ assert idx_match is not None
37
+ frame_idx = int(idx_match.group(0))
38
+
39
+ quat_xyzw = image.qvec[[1, 2, 3, 0]] # COLMAP uses wxyz quaternions
40
+ camera = cameras[image.camera_id]
41
+ np.array([1.0, 1.0])
42
+
43
+ visible = [id != -1 and points3D.get(id) is not None for id in image.point3D_ids]
44
+ visible_ids = image.point3D_ids[visible]
45
+
46
+ if filter_output and len(visible_ids) < filter_min_visible:
47
+ continue
48
+
49
+ visible_xyzs = [points3D[id] for id in visible_ids]
50
+ visible_xys = image.xys[visible]
51
+
52
+ rr.set_time_sequence("frame", frame_idx)
53
+
54
+ points = [point.xyz for point in visible_xyzs]
55
+ point_colors = [point.rgb for point in visible_xyzs]
56
+ point_errors = [point.error for point in visible_xyzs]
57
+
58
+ rr.log_scalar("plot/avg_reproj_err", np.mean(point_errors), color=[240, 45, 58])
59
+
60
+ rr.log_points("points", points, colors=point_colors, ext={"error": point_errors})
61
+
62
+ # COLMAP's camera transform is "camera from world"
63
+ rr.log_transform3d(
64
+ "camera", rr.TranslationRotationScale3D(image.tvec, rr.Quaternion(xyzw=quat_xyzw)), from_parent=True
65
+ )
66
+ rr.log_view_coordinates("camera", xyz="RDF") # X=Right, Y=Down, Z=Forward
67
+
68
+ # Log camera intrinsics
69
+ assert camera.model == "PINHOLE"
70
+ rr.log_pinhole(
71
+ "camera/image",
72
+ width=camera.width,
73
+ height=camera.height,
74
+ focal_length_px=camera.params[:2],
75
+ principal_point_px=camera.params[2:],
76
+ )
77
+
78
+ rr.log_image_file("camera/image", img_path=dataset_path / "images" / image.name)
79
+ rr.log_points("camera/image/keypoints", visible_xys, colors=[34, 138, 167])
80
+
81
+ rerun_output_directory = output_path / "rerun"
82
+ rerun_output_directory.mkdir(parents=True, exist_ok=True)
83
+ rerun_output_file = rerun_output_directory / "recording.rrd"
84
+ rr.save(rerun_output_file.as_posix())
85
+
services/utils/read_write_model.py ADDED
@@ -0,0 +1,514 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # This file is adapted from
2
+ # https://github.com/colmap/colmap/blob/bf3e19140f491c3042bfd85b7192ef7d249808ec/scripts/python/read_write_model.py
3
+ # Copyright (c) 2023, ETH Zurich and UNC Chapel Hill.
4
+ # All rights reserved.
5
+ #
6
+ # Redistribution and use in source and binary forms, with or without
7
+ # modification, are permitted provided that the following conditions are met:
8
+ #
9
+ # * Redistributions of source code must retain the above copyright
10
+ # notice, this list of conditions and the following disclaimer.
11
+ #
12
+ # * Redistributions in binary form must reproduce the above copyright
13
+ # notice, this list of conditions and the following disclaimer in the
14
+ # documentation and/or other materials provided with the distribution.
15
+ #
16
+ # * Neither the name of ETH Zurich and UNC Chapel Hill nor the names of
17
+ # its contributors may be used to endorse or promote products derived
18
+ # from this software without specific prior written permission.
19
+ #
20
+ # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
21
+ # AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
22
+ # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
23
+ # ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR CONTRIBUTORS BE
24
+ # LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
25
+ # CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
26
+ # SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
27
+ # INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
28
+ # CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
29
+ # ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
30
+ # POSSIBILITY OF SUCH DAMAGE.
31
+ #
32
+ # Author: Johannes L. Schoenberger (jsch-at-demuc-dot-de)
33
+ # type: ignore
34
+ from __future__ import annotations
35
+
36
+ import argparse
37
+ import collections
38
+ import os
39
+ import struct
40
+ from pathlib import Path
41
+ from typing import Mapping
42
+
43
+ import numpy as np
44
+
45
+ CameraModel = collections.namedtuple("CameraModel", ["model_id", "model_name", "num_params"])
46
+ Camera = collections.namedtuple("Camera", ["id", "model", "width", "height", "params"])
47
+ BaseImage = collections.namedtuple("Image", ["id", "qvec", "tvec", "camera_id", "name", "xys", "point3D_ids"])
48
+ Point3D = collections.namedtuple("Point3D", ["id", "xyz", "rgb", "error", "image_ids", "point2D_idxs"])
49
+
50
+
51
+ class Image(BaseImage):
52
+ def qvec2rotmat(self):
53
+ return qvec2rotmat(self.qvec)
54
+
55
+
56
+ CAMERA_MODELS = {
57
+ CameraModel(model_id=0, model_name="SIMPLE_PINHOLE", num_params=3),
58
+ CameraModel(model_id=1, model_name="PINHOLE", num_params=4),
59
+ CameraModel(model_id=2, model_name="SIMPLE_RADIAL", num_params=4),
60
+ CameraModel(model_id=3, model_name="RADIAL", num_params=5),
61
+ CameraModel(model_id=4, model_name="OPENCV", num_params=8),
62
+ CameraModel(model_id=5, model_name="OPENCV_FISHEYE", num_params=8),
63
+ CameraModel(model_id=6, model_name="FULL_OPENCV", num_params=12),
64
+ CameraModel(model_id=7, model_name="FOV", num_params=5),
65
+ CameraModel(model_id=8, model_name="SIMPLE_RADIAL_FISHEYE", num_params=4),
66
+ CameraModel(model_id=9, model_name="RADIAL_FISHEYE", num_params=5),
67
+ CameraModel(model_id=10, model_name="THIN_PRISM_FISHEYE", num_params=12),
68
+ }
69
+ CAMERA_MODEL_IDS = {camera_model.model_id: camera_model for camera_model in CAMERA_MODELS}
70
+ CAMERA_MODEL_NAMES = {camera_model.model_name: camera_model for camera_model in CAMERA_MODELS}
71
+
72
+
73
+ def read_next_bytes(fid, num_bytes, format_char_sequence, endian_character="<"):
74
+ """
75
+ Read and unpack the next bytes from a binary file.
76
+
77
+ :param fid:
78
+ :param num_bytes: Sum of combination of {2, 4, 8}, e.g. 2, 6, 16, 30, etc.
79
+ :param format_char_sequence: List of {c, e, f, d, h, H, i, I, l, L, q, Q}.
80
+ :param endian_character: Any of {@, =, <, >, !}
81
+ :return: Tuple of read and unpacked values.
82
+ """
83
+ data = fid.read(num_bytes)
84
+ return struct.unpack(endian_character + format_char_sequence, data)
85
+
86
+
87
+ def write_next_bytes(fid, data, format_char_sequence, endian_character="<"):
88
+ """
89
+ Pack and write to a binary file.
90
+
91
+ :param fid:
92
+ :param data: data to send, if multiple elements are sent at the same time,
93
+ they should be encapsuled either in a list or a tuple
94
+ :param format_char_sequence: List of {c, e, f, d, h, H, i, I, l, L, q, Q}.
95
+ should be the same length as the data list or tuple
96
+ :param endian_character: Any of {@, =, <, >, !}
97
+ """
98
+ if isinstance(data, (list, tuple)):
99
+ bytes = struct.pack(endian_character + format_char_sequence, *data)
100
+ else:
101
+ bytes = struct.pack(endian_character + format_char_sequence, data)
102
+ fid.write(bytes)
103
+
104
+
105
+ def read_cameras_text(path: Path):
106
+ """
107
+ see: src/base/reconstruction.cc
108
+ void Reconstruction::WriteCamerasText(const std::string& path)
109
+ void Reconstruction::ReadCamerasText(const std::string& path)
110
+ """
111
+ cameras = {}
112
+ with open(path) as fid:
113
+ while True:
114
+ line = fid.readline()
115
+ if not line:
116
+ break
117
+ line = line.strip()
118
+ if len(line) > 0 and line[0] != "#":
119
+ elems = line.split()
120
+ camera_id = int(elems[0])
121
+ model = elems[1]
122
+ width = int(elems[2])
123
+ height = int(elems[3])
124
+ params = np.array(tuple(map(float, elems[4:])))
125
+ cameras[camera_id] = Camera(id=camera_id, model=model, width=width, height=height, params=params)
126
+ return cameras
127
+
128
+
129
+ def read_cameras_binary(path_to_model_file: Path) -> Mapping[int, Camera]:
130
+ """
131
+ see: src/base/reconstruction.cc
132
+ void Reconstruction::WriteCamerasBinary(const std::string& path)
133
+ void Reconstruction::ReadCamerasBinary(const std::string& path)
134
+ """
135
+ cameras = {}
136
+ with path_to_model_file.open("rb") as fid:
137
+ num_cameras = read_next_bytes(fid, 8, "Q")[0]
138
+ for _ in range(num_cameras):
139
+ camera_properties = read_next_bytes(fid, num_bytes=24, format_char_sequence="iiQQ")
140
+ camera_id = camera_properties[0]
141
+ model_id = camera_properties[1]
142
+ model_name = CAMERA_MODEL_IDS[camera_properties[1]].model_name
143
+ width = camera_properties[2]
144
+ height = camera_properties[3]
145
+ num_params = CAMERA_MODEL_IDS[model_id].num_params
146
+ params = read_next_bytes(fid, num_bytes=8 * num_params, format_char_sequence="d" * num_params)
147
+ cameras[camera_id] = Camera(
148
+ id=camera_id, model=model_name, width=width, height=height, params=np.array(params)
149
+ )
150
+ assert len(cameras) == num_cameras
151
+ return cameras
152
+
153
+
154
+ def write_cameras_text(cameras, path):
155
+ """
156
+ see: src/base/reconstruction.cc
157
+ void Reconstruction::WriteCamerasText(const std::string& path)
158
+ void Reconstruction::ReadCamerasText(const std::string& path)
159
+ """
160
+ HEADER = (
161
+ "# Camera list with one line of data per camera:\n"
162
+ + "# CAMERA_ID, MODEL, WIDTH, HEIGHT, PARAMS[]\n"
163
+ + f"# Number of cameras: {len(cameras)}\n"
164
+ )
165
+ with open(path, "w") as fid:
166
+ fid.write(HEADER)
167
+ for _, cam in cameras.items():
168
+ to_write = [cam.id, cam.model, cam.width, cam.height, *cam.params]
169
+ line = " ".join([str(elem) for elem in to_write])
170
+ fid.write(line + "\n")
171
+
172
+
173
+ def write_cameras_binary(cameras, path_to_model_file):
174
+ """
175
+ see: src/base/reconstruction.cc
176
+ void Reconstruction::WriteCamerasBinary(const std::string& path)
177
+ void Reconstruction::ReadCamerasBinary(const std::string& path)
178
+ """
179
+ with open(path_to_model_file, "wb") as fid:
180
+ write_next_bytes(fid, len(cameras), "Q")
181
+ for _, cam in cameras.items():
182
+ model_id = CAMERA_MODEL_NAMES[cam.model].model_id
183
+ camera_properties = [cam.id, model_id, cam.width, cam.height]
184
+ write_next_bytes(fid, camera_properties, "iiQQ")
185
+ for p in cam.params:
186
+ write_next_bytes(fid, float(p), "d")
187
+ return cameras
188
+
189
+
190
+ def read_images_text(path: Path):
191
+ """
192
+ see: src/base/reconstruction.cc
193
+ void Reconstruction::ReadImagesText(const std::string& path)
194
+ void Reconstruction::WriteImagesText(const std::string& path)
195
+ """
196
+ images = {}
197
+ with open(path) as fid:
198
+ while True:
199
+ line = fid.readline()
200
+ if not line:
201
+ break
202
+ line = line.strip()
203
+ if len(line) > 0 and line[0] != "#":
204
+ elems = line.split()
205
+ image_id = int(elems[0])
206
+ qvec = np.array(tuple(map(float, elems[1:5])))
207
+ tvec = np.array(tuple(map(float, elems[5:8])))
208
+ camera_id = int(elems[8])
209
+ image_name = elems[9]
210
+ elems = fid.readline().split()
211
+ xys = np.column_stack([tuple(map(float, elems[0::3])), tuple(map(float, elems[1::3]))])
212
+ point3D_ids = np.array(tuple(map(int, elems[2::3])))
213
+ images[image_id] = Image(
214
+ id=image_id,
215
+ qvec=qvec,
216
+ tvec=tvec,
217
+ camera_id=camera_id,
218
+ name=image_name,
219
+ xys=xys,
220
+ point3D_ids=point3D_ids,
221
+ )
222
+ return images
223
+
224
+
225
+ def read_images_binary(path_to_model_file: Path) -> Mapping[int, Image]:
226
+ """
227
+ see: src/base/reconstruction.cc
228
+ void Reconstruction::ReadImagesBinary(const std::string& path)
229
+ void Reconstruction::WriteImagesBinary(const std::string& path)
230
+ """
231
+ images = {}
232
+ with open(path_to_model_file, "rb") as fid:
233
+ num_reg_images = read_next_bytes(fid, 8, "Q")[0]
234
+ for _ in range(num_reg_images):
235
+ binary_image_properties = read_next_bytes(fid, num_bytes=64, format_char_sequence="idddddddi")
236
+ image_id = binary_image_properties[0]
237
+ qvec = np.array(binary_image_properties[1:5])
238
+ tvec = np.array(binary_image_properties[5:8])
239
+ camera_id = binary_image_properties[8]
240
+ image_name = ""
241
+ current_char = read_next_bytes(fid, 1, "c")[0]
242
+ while current_char != b"\x00": # look for the ASCII 0 entry
243
+ image_name += current_char.decode("utf-8")
244
+ current_char = read_next_bytes(fid, 1, "c")[0]
245
+ num_points2D = read_next_bytes(fid, num_bytes=8, format_char_sequence="Q")[0]
246
+ x_y_id_s = read_next_bytes(fid, num_bytes=24 * num_points2D, format_char_sequence="ddq" * num_points2D)
247
+ xys = np.column_stack([tuple(map(float, x_y_id_s[0::3])), tuple(map(float, x_y_id_s[1::3]))])
248
+ point3D_ids = np.array(tuple(map(int, x_y_id_s[2::3])))
249
+ images[image_id] = Image(
250
+ id=image_id,
251
+ qvec=qvec,
252
+ tvec=tvec,
253
+ camera_id=camera_id,
254
+ name=image_name,
255
+ xys=xys,
256
+ point3D_ids=point3D_ids,
257
+ )
258
+ return images
259
+
260
+
261
+ def write_images_text(images, path):
262
+ """
263
+ see: src/base/reconstruction.cc
264
+ void Reconstruction::ReadImagesText(const std::string& path)
265
+ void Reconstruction::WriteImagesText(const std::string& path)
266
+ """
267
+ if len(images) == 0:
268
+ mean_observations = 0
269
+ else:
270
+ mean_observations = sum((len(img.point3D_ids) for _, img in images.items())) / len(images)
271
+ HEADER = (
272
+ "# Image list with two lines of data per image:\n"
273
+ + "# IMAGE_ID, QW, QX, QY, QZ, TX, TY, TZ, CAMERA_ID, NAME\n"
274
+ + "# POINTS2D[] as (X, Y, POINT3D_ID)\n"
275
+ + f"# Number of images: {len(images)}, mean observations per image: {mean_observations}\n"
276
+ )
277
+
278
+ with open(path, "w") as fid:
279
+ fid.write(HEADER)
280
+ for _, img in images.items():
281
+ image_header = [img.id, *img.qvec, *img.tvec, img.camera_id, img.name]
282
+ first_line = " ".join(map(str, image_header))
283
+ fid.write(first_line + "\n")
284
+
285
+ points_strings = []
286
+ for xy, point3D_id in zip(img.xys, img.point3D_ids):
287
+ points_strings.append(" ".join(map(str, [*xy, point3D_id])))
288
+ fid.write(" ".join(points_strings) + "\n")
289
+
290
+
291
+ def write_images_binary(images, path_to_model_file):
292
+ """
293
+ see: src/base/reconstruction.cc
294
+ void Reconstruction::ReadImagesBinary(const std::string& path)
295
+ void Reconstruction::WriteImagesBinary(const std::string& path)
296
+ """
297
+ with open(path_to_model_file, "wb") as fid:
298
+ write_next_bytes(fid, len(images), "Q")
299
+ for _, img in images.items():
300
+ write_next_bytes(fid, img.id, "i")
301
+ write_next_bytes(fid, img.qvec.tolist(), "dddd")
302
+ write_next_bytes(fid, img.tvec.tolist(), "ddd")
303
+ write_next_bytes(fid, img.camera_id, "i")
304
+ for char in img.name:
305
+ write_next_bytes(fid, char.encode("utf-8"), "c")
306
+ write_next_bytes(fid, b"\x00", "c")
307
+ write_next_bytes(fid, len(img.point3D_ids), "Q")
308
+ for xy, p3d_id in zip(img.xys, img.point3D_ids):
309
+ write_next_bytes(fid, [*xy, p3d_id], "ddq")
310
+
311
+
312
+ def read_points3D_text(path):
313
+ """
314
+ see: src/base/reconstruction.cc
315
+ void Reconstruction::ReadPoints3DText(const std::string& path)
316
+ void Reconstruction::WritePoints3DText(const std::string& path)
317
+ """
318
+ points3D = {}
319
+ with open(path) as fid:
320
+ while True:
321
+ line = fid.readline()
322
+ if not line:
323
+ break
324
+ line = line.strip()
325
+ if len(line) > 0 and line[0] != "#":
326
+ elems = line.split()
327
+ point3D_id = int(elems[0])
328
+ xyz = np.array(tuple(map(float, elems[1:4])))
329
+ rgb = np.array(tuple(map(int, elems[4:7])))
330
+ error = float(elems[7])
331
+ image_ids = np.array(tuple(map(int, elems[8::2])))
332
+ point2D_idxs = np.array(tuple(map(int, elems[9::2])))
333
+ points3D[point3D_id] = Point3D(
334
+ id=point3D_id, xyz=xyz, rgb=rgb, error=error, image_ids=image_ids, point2D_idxs=point2D_idxs
335
+ )
336
+ return points3D
337
+
338
+
339
+ def read_points3D_binary(path_to_model_file: Path) -> Mapping[int, Point3D]:
340
+ """
341
+ see: src/base/reconstruction.cc
342
+ void Reconstruction::ReadPoints3DBinary(const std::string& path)
343
+ void Reconstruction::WritePoints3DBinary(const std::string& path)
344
+ """
345
+ points3D = {}
346
+ with open(path_to_model_file, "rb") as fid:
347
+ num_points = read_next_bytes(fid, 8, "Q")[0]
348
+ for _ in range(num_points):
349
+ binary_point_line_properties = read_next_bytes(fid, num_bytes=43, format_char_sequence="QdddBBBd")
350
+ point3D_id = binary_point_line_properties[0]
351
+ xyz = np.array(binary_point_line_properties[1:4])
352
+ rgb = np.array(binary_point_line_properties[4:7])
353
+ error = np.array(binary_point_line_properties[7])
354
+ track_length = read_next_bytes(fid, num_bytes=8, format_char_sequence="Q")[0]
355
+ track_elems = read_next_bytes(fid, num_bytes=8 * track_length, format_char_sequence="ii" * track_length)
356
+ image_ids = np.array(tuple(map(int, track_elems[0::2])))
357
+ point2D_idxs = np.array(tuple(map(int, track_elems[1::2])))
358
+ points3D[point3D_id] = Point3D(
359
+ id=point3D_id, xyz=xyz, rgb=rgb, error=error, image_ids=image_ids, point2D_idxs=point2D_idxs
360
+ )
361
+ return points3D
362
+
363
+
364
+ def write_points3D_text(points3D, path):
365
+ """
366
+ see: src/base/reconstruction.cc
367
+ void Reconstruction::ReadPoints3DText(const std::string& path)
368
+ void Reconstruction::WritePoints3DText(const std::string& path)
369
+ """
370
+ if len(points3D) == 0:
371
+ mean_track_length = 0
372
+ else:
373
+ mean_track_length = sum((len(pt.image_ids) for _, pt in points3D.items())) / len(points3D)
374
+ HEADER = (
375
+ "# 3D point list with one line of data per point:\n"
376
+ + "# POINT3D_ID, X, Y, Z, R, G, B, ERROR, TRACK[] as (IMAGE_ID, POINT2D_IDX)\n"
377
+ + f"# Number of points: {len(points3D)}, mean track length: {mean_track_length}\n"
378
+ )
379
+
380
+ with open(path, "w") as fid:
381
+ fid.write(HEADER)
382
+ for _, pt in points3D.items():
383
+ point_header = [pt.id, *pt.xyz, *pt.rgb, pt.error]
384
+ fid.write(" ".join(map(str, point_header)) + " ")
385
+ track_strings = []
386
+ for image_id, point2D in zip(pt.image_ids, pt.point2D_idxs):
387
+ track_strings.append(" ".join(map(str, [image_id, point2D])))
388
+ fid.write(" ".join(track_strings) + "\n")
389
+
390
+
391
+ def write_points3D_binary(points3D, path_to_model_file):
392
+ """
393
+ see: src/base/reconstruction.cc
394
+ void Reconstruction::ReadPoints3DBinary(const std::string& path)
395
+ void Reconstruction::WritePoints3DBinary(const std::string& path)
396
+ """
397
+ with open(path_to_model_file, "wb") as fid:
398
+ write_next_bytes(fid, len(points3D), "Q")
399
+ for _, pt in points3D.items():
400
+ write_next_bytes(fid, pt.id, "Q")
401
+ write_next_bytes(fid, pt.xyz.tolist(), "ddd")
402
+ write_next_bytes(fid, pt.rgb.tolist(), "BBB")
403
+ write_next_bytes(fid, pt.error, "d")
404
+ track_length = pt.image_ids.shape[0]
405
+ write_next_bytes(fid, track_length, "Q")
406
+ for image_id, point2D_id in zip(pt.image_ids, pt.point2D_idxs):
407
+ write_next_bytes(fid, [image_id, point2D_id], "ii")
408
+
409
+
410
+ def detect_model_format(path: Path, ext: str) -> bool:
411
+ parts = ["cameras", "images", "points3D"]
412
+ if all([(path / p).with_suffix(ext) for p in parts]):
413
+ print("Detected model format: '" + ext + "'")
414
+ return True
415
+
416
+ return False
417
+
418
+
419
+ def read_model(path: Path, ext: str = ""):
420
+ # try to detect the extension automatically
421
+ if ext == "":
422
+ if detect_model_format(path, ".bin"):
423
+ ext = ".bin"
424
+ elif detect_model_format(path, ".txt"):
425
+ ext = ".txt"
426
+ else:
427
+ print("Provide model format: '.bin' or '.txt'")
428
+ return
429
+
430
+ if ext == ".txt":
431
+ cameras = read_cameras_text((path / "cameras").with_suffix(ext))
432
+ images = read_images_text((path / "images").with_suffix(ext))
433
+ points3D = read_points3D_text((path / "points3D").with_suffix(ext))
434
+ else:
435
+ cameras = read_cameras_binary((path / "cameras").with_suffix(ext))
436
+ images = read_images_binary((path / "images").with_suffix(ext))
437
+ points3D = read_points3D_binary((path / "points3D").with_suffix(ext))
438
+ return cameras, images, points3D
439
+
440
+
441
+ def write_model(cameras, images, points3D, path, ext=".bin"):
442
+ if ext == ".txt":
443
+ write_cameras_text(cameras, os.path.join(path, "cameras" + ext))
444
+ write_images_text(images, os.path.join(path, "images" + ext))
445
+ write_points3D_text(points3D, os.path.join(path, "points3D") + ext)
446
+ else:
447
+ write_cameras_binary(cameras, os.path.join(path, "cameras" + ext))
448
+ write_images_binary(images, os.path.join(path, "images" + ext))
449
+ write_points3D_binary(points3D, os.path.join(path, "points3D") + ext)
450
+ return cameras, images, points3D
451
+
452
+
453
+ def qvec2rotmat(qvec):
454
+ return np.array(
455
+ [
456
+ [
457
+ 1 - 2 * qvec[2] ** 2 - 2 * qvec[3] ** 2,
458
+ 2 * qvec[1] * qvec[2] - 2 * qvec[0] * qvec[3],
459
+ 2 * qvec[3] * qvec[1] + 2 * qvec[0] * qvec[2],
460
+ ],
461
+ [
462
+ 2 * qvec[1] * qvec[2] + 2 * qvec[0] * qvec[3],
463
+ 1 - 2 * qvec[1] ** 2 - 2 * qvec[3] ** 2,
464
+ 2 * qvec[2] * qvec[3] - 2 * qvec[0] * qvec[1],
465
+ ],
466
+ [
467
+ 2 * qvec[3] * qvec[1] - 2 * qvec[0] * qvec[2],
468
+ 2 * qvec[2] * qvec[3] + 2 * qvec[0] * qvec[1],
469
+ 1 - 2 * qvec[1] ** 2 - 2 * qvec[2] ** 2,
470
+ ],
471
+ ]
472
+ )
473
+
474
+
475
+ def rotmat2qvec(R):
476
+ Rxx, Ryx, Rzx, Rxy, Ryy, Rzy, Rxz, Ryz, Rzz = R.flat
477
+ K = (
478
+ np.array(
479
+ [
480
+ [Rxx - Ryy - Rzz, 0, 0, 0],
481
+ [Ryx + Rxy, Ryy - Rxx - Rzz, 0, 0],
482
+ [Rzx + Rxz, Rzy + Ryz, Rzz - Rxx - Ryy, 0],
483
+ [Ryz - Rzy, Rzx - Rxz, Rxy - Ryx, Rxx + Ryy + Rzz],
484
+ ]
485
+ )
486
+ / 3.0
487
+ )
488
+ eigvals, eigvecs = np.linalg.eigh(K)
489
+ qvec = eigvecs[[3, 0, 1, 2], np.argmax(eigvals)]
490
+ if qvec[0] < 0:
491
+ qvec *= -1
492
+ return qvec
493
+
494
+
495
+ def main():
496
+ parser = argparse.ArgumentParser(description="Read and write COLMAP binary and text models")
497
+ parser.add_argument("--input_model", help="path to input model folder")
498
+ parser.add_argument("--input_format", choices=[".bin", ".txt"], help="input model format", default="")
499
+ parser.add_argument("--output_model", help="path to output model folder")
500
+ parser.add_argument("--output_format", choices=[".bin", ".txt"], help="output model format", default=".txt")
501
+ args = parser.parse_args()
502
+
503
+ cameras, images, points3D = read_model(path=args.input_model, ext=args.input_format)
504
+
505
+ print("num_cameras:", len(cameras))
506
+ print("num_images:", len(images))
507
+ print("num_points3D:", len(points3D))
508
+
509
+ if args.output_model is not None:
510
+ write_model(cameras, images, points3D, path=args.output_model, ext=args.output_format)
511
+
512
+
513
+ if __name__ == "__main__":
514
+ main()