Spaces:
				
			
			
	
			
			
		Build error
		
	
	
	
			
			
	
	
	
	
		
		
		Build error
		
	🔥 First commit
Browse files- .devcontainer/devcontainer.json +31 -0
- .gitignore +7 -0
- .gitmodules +16 -0
- Dockerfile +145 -0
- deps/colmap +1 -0
- deps/gaussian-splatting-cuda +1 -0
- deps/rerun +1 -0
- deps/splat +1 -0
- requirements.txt +8 -0
- server.py +422 -0
- services/colmap.py +244 -0
- services/ffmpeg.py +100 -0
- services/gaussian_splatting_cuda.py +108 -0
- services/http.py +26 -0
- services/rerun.py +85 -0
- services/utils/read_write_model.py +514 -0
    	
        .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 = [email protected]:MrNeRF/gaussian-splatting-cuda.git
         | 
| 4 | 
            +
            	branch = master
         | 
| 5 | 
            +
            [submodule "deps/colmap"]
         | 
| 6 | 
            +
            	path = deps/colmap
         | 
| 7 | 
            +
            	url = [email protected]:colmap/colmap.git
         | 
| 8 | 
            +
            	branch = main
         | 
| 9 | 
            +
            [submodule "deps/rerun"]
         | 
| 10 | 
            +
            	path = deps/rerun
         | 
| 11 | 
            +
            	url = [email protected]:rerun-io/rerun.git
         | 
| 12 | 
            +
            	branch = release-0.8.2
         | 
| 13 | 
            +
            [submodule "deps/splat"]
         | 
| 14 | 
            +
            	path = deps/splat
         | 
| 15 | 
            +
            	url = [email protected]: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()
         |