Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 7 additions & 2 deletions Project.toml
Original file line number Diff line number Diff line change
@@ -1,9 +1,10 @@
name = "GaussianSplatting"
uuid = "991e6b22-92f0-46ac-8b70-8a93e7beee5d"
version = "1.1.0"
version = "1.2.0"
authors = ["Anton Smirnov <[email protected]>"]

[deps]
AcceleratedKernels = "6a4ca0a5-0e36-4168-a932-d9be78d558f1"
Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
BSON = "fbb218c0-5317-5bc6-957e-2ee96dd4b1f0"
CImGui = "5d785b6c-b76f-510e-a07c-3070796c7e87"
Expand Down Expand Up @@ -38,13 +39,16 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
[weakdeps]
AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
Metal = "dde4c033-4e86-420c-a63e-0dd931031962"

[extensions]
GaussianSplattingAMDGPUExt = "AMDGPU"
GaussianSplattingCUDAExt = "CUDA"
GaussianSplattingMetalExt = "Metal"

[compat]
AMDGPU = "2"
AcceleratedKernels = "0.4.3"
Adapt = "4"
BSON = "0.3"
CImGui = "6"
Expand All @@ -61,10 +65,11 @@ ImageIO = "0.6"
ImageMagick = "1.3"
ImageTransformations = "0.10"
KernelAbstractions = "0.9.34"
Metal = "1.9"
ModernGL = "1.1"
NearestNeighbors = "0.4"
NerfUtils = "0.2"
NeuralGraphicsGL = "0.5"
NeuralGraphicsGL = "0.5.1"
PlyIO = "1.2"
Quaternions = "0.7"
Rotations = "1.7"
Expand Down
15 changes: 9 additions & 6 deletions docs/src/index.md
Original file line number Diff line number Diff line change
Expand Up @@ -7,8 +7,9 @@ Gaussian Splatting algorithm in pure Julia.
## Requirements

- Julia 1.10 or higher.
- AMD ([AMDGPU.jl](https://github.com/JuliaGPU/AMDGPU.jl)) or
Nvidia ([CUDA.jl](https://github.com/JuliaGPU/CUDA.jl)) capable machine.
- [AMDGPU.jl](https://github.com/JuliaGPU/AMDGPU.jl) or
[CUDA.jl](https://github.com/JuliaGPU/CUDA.jl) or
[Metal.jl](https://github.com/JuliaGPU/Metal.jl) capable machine.

## Install

Expand All @@ -24,14 +25,16 @@ GaussianSplatting.jl comes with a GUI application to train & view the gaussians.

1. Add necessary packages:
```julia
] add AMDGPU # for AMD GPU
] add CUDA # for Nvidia GPU
] add AMDGPU # for AMD GPU
] add CUDA # for Nvidia GPU
] add Metal # for Apple GPU
```

2. Run:
```julia
julia> using AMDGPU; kab = ROCBackend() # for AMD GPU
julia> using CUDA; kab = CUDABackend() # for Nvidia GPU
julia> using AMDGPU; kab = ROCBackend() # for AMD GPU
julia> using CUDA; kab = CUDABackend() # for Nvidia GPU
julia> using Metal; kab = MetalBackend() # for Apple GPU
julia> GaussianSplatting.gui(kab, "path-to-colmap-dataset-directory"; scale=1)
```

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ GaussianSplatting.base_array_type(::ROCBackend) = ROCArray

GaussianSplatting.use_ak(::ROCBackend) = true

function GaussianSplatting.allocate_pinned(kab, ::Type{T}, shape) where T
function GaussianSplatting.allocate_pinned(::ROCBackend, ::Type{T}, shape) where T
x = Array{T}(undef, shape)
xd = unsafe_wrap(ROCArray, pointer(x), size(x))
return x, xd
Expand Down
19 changes: 19 additions & 0 deletions ext/GaussianSplattingMetalExt.jl
Original file line number Diff line number Diff line change
@@ -0,0 +1,19 @@
module GaussianSplattingMetalExt

using Metal
using GaussianSplatting

GaussianSplatting.base_array_type(::MetalBackend) = MtlArray

GaussianSplatting.use_ak(::MetalBackend) = true

function GaussianSplatting.allocate_pinned(::MetalBackend, ::Type{T}, shape) where T
xd = MtlArray{T, length(shape), Metal.SharedStorage}(undef, shape)
x = reshape(unsafe_wrap(Vector{T}, reshape(xd, :)), shape)
return x, xd
end

# Unregistered automatically in the array dtor.
GaussianSplatting.unpin_memory(::MtlArray) = return

end
1 change: 1 addition & 0 deletions src/GaussianSplatting.jl
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@ using GLFW

import CImGui.lib as iglib

import AcceleratedKernels as AK
import BSON
import ChainRulesCore as CRC
import ImageFiltering
Expand Down
4 changes: 2 additions & 2 deletions src/gui/gui.jl
Original file line number Diff line number Diff line change
Expand Up @@ -111,7 +111,7 @@ end

# Viewer-only mode.
function GSGUI(kab, gaussians::GaussianModel, camera::Camera; gl_kwargs...)
NGL.init(3, 0)
NGL.init(3, 2)
context = NGL.Context("GaussianSplatting.jl"; gl_kwargs...)
NGL.set_resize_callback!(context, resize_callback)

Expand Down Expand Up @@ -144,7 +144,7 @@ end

# Training mode.
function GSGUI(kab, dataset_path::String, scale::Int; gl_kwargs...)
NGL.init(3, 0)
NGL.init(3, 2)
context = NGL.Context("GaussianSplatting.jl"; gl_kwargs...)
NGL.set_resize_callback!(context, resize_callback)

Expand Down
16 changes: 8 additions & 8 deletions src/rasterization/projection.jl
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@ function project(
rotations::AbstractMatrix{Float32};
rast::GaussianRasterizer, camera::Camera,
near_plane::Float32, far_plane::Float32,
radius_clip::Float32, blur_ϵ::Float32,
radius_clip::Int32, blur_ϵ::Float32,
)
(; width, height) = resolution(camera)
@assert width % 16 == 0 && height % 16 == 0
Expand Down Expand Up @@ -57,7 +57,7 @@ function ∇project(
conics::AbstractMatrix{Float32};
rast::GaussianRasterizer, camera::Camera,
near_plane::Float32, far_plane::Float32,
radius_clip::Float32, blur_ϵ::Float32,
radius_clip::Int32, blur_ϵ::Float32,
)
K = camera.intrinsics
R_w2c = SMatrix{3, 3, Float32}(camera.w2c[1:3, 1:3])
Expand Down Expand Up @@ -110,7 +110,7 @@ function ChainRulesCore.rrule(::typeof(project),

rast::GaussianRasterizer, camera::Camera,
near_plane::Float32, far_plane::Float32,
radius_clip::Float32, blur_ϵ::Float32,
radius_clip::Int32, blur_ϵ::Float32,
)
means_2d, conics, compensations, depths = project(
means_3d, scales, rotations;
Expand Down Expand Up @@ -151,7 +151,7 @@ end
# Config.
near_plane::Float32,
far_plane::Float32,
radius_clip::Float32,
radius_clip::Int32,
blur_ϵ::Float32,
) where {C <: Maybe{AbstractMatrix{Float32}}, RM}
i = @index(Global)
Expand Down Expand Up @@ -194,10 +194,10 @@ end

# Discard Gaussians outside of image plane.
if (
(mean_2D[1] + radius) ≤ 0 ||
(mean_2D[1] - radius) ≥ resolution[1] ||
(mean_2D[2] + radius) ≤ 0 ||
(mean_2D[2] - radius) ≥ resolution[2]
(mean_2D[1] + radius) ≤ 0f0 ||
(mean_2D[1] - radius) ≥ Float32(resolution[1]) ||
(mean_2D[2] + radius) ≤ 0f0 ||
(mean_2D[2] - radius) ≥ Float32(resolution[2])
)
radii[i] = 0i32
return
Expand Down
6 changes: 3 additions & 3 deletions src/rasterization/rasterizer.jl
Original file line number Diff line number Diff line change
Expand Up @@ -192,7 +192,7 @@ function (rast::GaussianRasterizer)(
means_2d, conics, compensations, depths = project(
means_3d, scales_act, rotations;
rast, camera, near_plane=0.2f0, far_plane=1000f0,
radius_clip=3f0, blur_ϵ=0.3f0)
radius_clip=Int32(3), blur_ϵ=0.3f0)

colors = spherical_harmonics(means_3d, shs; rast, camera, sh_degree)

Expand Down Expand Up @@ -252,7 +252,7 @@ function rasterize(

# TODO make configurable.
near_plane, far_plane = 0.2f0, 1000f0
radius_clip = 3f0 # In pixels.
radius_clip = Int32(3) # In pixels.
blur_ϵ = 0.3f0

project!(kab)(
Expand Down Expand Up @@ -315,7 +315,7 @@ function rasterize(
rast.gstate.radii, rast.grid, BLOCK; ndrange=n)

if use_ak(kab)
sortperm!(
AK.sortperm!(
@view(rast.bstate.permutation[1:n_rendered]),
@view(rast.bstate.gaussian_keys_unsorted[1:n_rendered]);
temp=@view(rast.bstate.permutation_tmp[1:n_rendered]))
Expand Down
12 changes: 1 addition & 11 deletions src/rasterization/utils.jl
Original file line number Diff line number Diff line change
Expand Up @@ -13,8 +13,7 @@ sdiagm(x, y, z) = SMatrix{3, 3, Float32, 9}(

gpu_floor(T, x) = unsafe_trunc(T, floor(x))
gpu_ceil(T, x) = unsafe_trunc(T, ceil(x))

gpu_cld(x, y::T) where T = (x + y - one(T)) ÷ y
gpu_cld(x::X, y::T) where {X, T} = unsafe_trunc(T, floor(Float32(x + y - one(X)) / Float32(y)))

Base.@propagate_inbounds function get_rect(
pixel::SVector{2, Float32}, max_radius::Int32,
Expand All @@ -26,15 +25,6 @@ Base.@propagate_inbounds function get_rect(
@inbounds rmax = SVector{2, Int32}(
clamp(gpu_floor(Int32, gpu_cld(pixel[1] + max_radius, block[1])), 0i32, grid[1]),
clamp(gpu_floor(Int32, gpu_cld(pixel[2] + max_radius, block[2])), 0i32, grid[2]))

# rblock = inv.(block)

# rmin = gpu_floor.(Int32, (pixel .- max_radius) .* rblock)
# rmin = clamp.(rmin, 0i32, grid)

# rmax = gpu_ceil.(Int32, (pixel .+ max_radius) .* rblock)
# # rmax = gpu_cld.(pixel .+ max_radius, block)
# rmax = clamp.(rmax, 0i32, grid)
return rmin, rmax
end

Expand Down