Add uncertainty & covisibility calculation

benchmark/pipeline.jl

@@ -31,4 +31,4 @@ function main(dataset_path::String; scale::Int)
     end
     return
 end
-main("/home/pxl-th/Downloads/360_v2/bicycle"; scale=4)
+main("/home/pxlth/Downloads/360_v2/bicycle"; scale=4)

src/GaussianSplatting.jl

@@ -86,7 +86,7 @@ function main(dataset_path::String; scale::Int, save_path::Maybe{String} = nothi
 
     gaussians = GaussianModel(dataset.points, dataset.colors, dataset.scales; max_sh_degree=3)
     rasterizer = GaussianRasterizer(kab, camera;
-        antialias=false, fused=true, mode=:rgbd)
+        antialias=false, fused=false, mode=:rgbd)
 
     opt_params = OptimizationParams()
     trainer = Trainer(rasterizer, gaussians, dataset, opt_params)
@@ -95,6 +95,9 @@ function main(dataset_path::String; scale::Int, save_path::Maybe{String} = nothi
     @info "N train images: $(length(dataset.train_cameras))"
     @info "N test images: $(length(dataset.test_cameras))"
 
+    width, height = camera.intrinsics.resolution
+    uncertainties = adapt(kab, zeros(Float32, width, height))
+
     # res = resolution(camera)
     # writer = open_video_out(
     #     "./out.mp4", zeros(RGB{N0f8}, res.height, res.width);
@@ -111,7 +114,12 @@ function main(dataset_path::String; scale::Int, save_path::Maybe{String} = nothi
             image_features = rasterizer(
                 gaussians.points, gaussians.opacities, gaussians.scales,
                 gaussians.rotations, gaussians.features_dc, gaussians.features_rest;
-                camera, sh_degree=gaussians.sh_degree)
+                camera, sh_degree=gaussians.sh_degree, uncertainties)
+
+            uncertainties_h = Array(uncertainties)
+            uncertainties_image = colorview(Gray, clamp01!(transpose(uncertainties_h)))
+            save("uncertainties-$(trainer.step).png", uncertainties_image)
+            fill!(uncertainties, 0f0)
 
             host_image_features = Array(image_features)
             save("image-$(trainer.step).png",

src/rasterization/projection.jl

@@ -77,11 +77,15 @@ function ∇project(
         _as_T(SVector{3, Float32}, vscales),
         _as_T(SVector{4, Float32}, vrot),
 
+        # TODO support pose optimization
+        nothing,
+        nothing,
+
         # Input grad outputs.
         _as_T(SVector{2, Float32}, vmeans_2d),
         _as_T(SVector{3, Float32}, vconics),
-        vcompensations,
-        vdepths,
+        vcompensations isa ZeroTangent ? nothing : vcompensations,
+        vdepths isa ZeroTangent ? nothing : vdepths,
 
         _as_T(SVector{3, Float32}, conics),
         rast.gstate.radii,
@@ -213,12 +217,13 @@ end
     vmeans::AbstractVector{SVector{3, Float32}},
     vcov_scales::AbstractVector{SVector{3, Float32}},
     vcov_rotations::AbstractVector{SVector{4, Float32}},
-    vR_out::RG, vt_out,
+    vR_out::RG,
+    vt_out,
 
     # Input grad outputs.
     vmeans_2d::AbstractVector{SVector{2, Float32}},
     vconics::AbstractArray{SVector{3, Float32}},
-    vcompensations::VC, #::AbstractVector{Float32},
+    vcompensations::VC,
     vdepths::VD,
 
     conics::AbstractVector{SVector{3, Float32}},
@@ -236,7 +241,13 @@ end
     resolution::SVector{2, Int32},
     principal::SVector{2, Float32},
     ϵ::Float32,
-) where {C <: Maybe{AbstractMatrix{Float32}}, VC, VD, RM, RG}
+) where {
+    C <: Maybe{AbstractMatrix{Float32}},
+    VC <: Maybe{AbstractVector{Float32}},
+    VD <: Maybe{AbstractVector{Float32}},
+    RM,
+    RG <: Maybe{AbstractMatrix{Float32}},
+}
     i = @index(Global)
 
     conic = conics[i]

src/rasterization/rasterizer.jl

@@ -144,6 +144,9 @@ function (rast::GaussianRasterizer)(
     R_w2c = nothing, t_w2c = nothing;
     camera::Camera, sh_degree::Int,
     background::SVector{3, Float32} = zeros(SVector{3, Float32}),
+
+    covisibilities::Maybe{AbstractVector{Bool}} = nothing,
+    uncertainties::Maybe{AbstractMatrix{Float32}} = nothing,
 )
     # If rendering outside AD, use non-allocating path.
     inplace = NNlib.within_gradient(means_3d)
@@ -176,7 +179,7 @@ function (rast::GaussianRasterizer)(
     if rast.fused
         return rasterize(
             means_3d, shs, opacities_act, scales_act, rotations, R_w2c, t_w2c;
-            rast, camera, sh_degree, background)
+            rast, camera, sh_degree, background, covisibilities, uncertainties)
     else
         means_2d, conics, compensations, depths = project(
             means_3d, scales_act, rotations;
@@ -185,7 +188,6 @@ function (rast::GaussianRasterizer)(
 
         colors = spherical_harmonics(means_3d, shs; rast, camera, sh_degree)
 
-        # TODO handle :d mode
         color_features = if rast.mode == :rgbd
             vcat(colors, reshape(depths, 1, :))
         else
@@ -198,10 +200,9 @@ function (rast::GaussianRasterizer)(
             opacities_act
         end
 
-        image = render(
+        return render(
             means_2d, conics, opacities_scaled, color_features;
-            rast, camera, background, depths)
-        return image
+            rast, camera, background, depths, covisibilities, uncertainties)
     end
 end
 
@@ -214,6 +215,9 @@ function rasterize(
     R_w2c = nothing, t_w2c = nothing;
     rast::GaussianRasterizer, camera::Camera, sh_degree::Int,
     background::SVector{3, Float32},
+
+    covisibilities::Maybe{AbstractVector{Bool}} = nothing,
+    uncertainties::Maybe{AbstractMatrix{Float32}} = nothing,
 )
     render_depth = rast.mode == :rgbd
     render_depth && @assert rast.gstate.color_features ≢ nothing
@@ -336,6 +340,7 @@ function rasterize(
     render!(kab, (Int.(BLOCK)...,), (width, height))(
         # Outputs.
         rast.image, rast.istate.n_contrib, rast.istate.accum_α,
+        covisibilities, uncertainties,
         # Inputs.
         rast.bstate.gaussian_values_sorted,
         rast.gstate.means_2d,
@@ -476,10 +481,13 @@ function ChainRulesCore.rrule(::typeof(rasterize),
     R_w2c = nothing, t_w2c = nothing;
     rast::GaussianRasterizer, camera::Camera, sh_degree::Int,
     background::SVector{3, Float32},
+
+    covisibilities::Maybe{AbstractVector{Bool}} = nothing,
+    uncertainties::Maybe{AbstractMatrix{Float32}} = nothing,
 )
     image = rasterize(
         means_3d, shs, opacities, scales, rotations, R_w2c, t_w2c;
-        rast, camera, sh_degree, background)
+        rast, camera, sh_degree, background, covisibilities, uncertainties)
 
     function _pullback(vpixels)
         ∇ = ∇rasterize(

src/rasterization/render.jl

@@ -6,6 +6,9 @@ function render(
     rast::GaussianRasterizer, camera::Camera, background::SVector{3, Float32},
     # Used only for sorting.
     depths::AbstractVector{Float32},
+
+    covisibilities::Maybe{AbstractVector{Bool}} = nothing,
+    uncertainties::Maybe{AbstractMatrix{Float32}} = nothing,
 )
     kab = get_backend(rast)
     n = size(means_2d, 2)
@@ -78,6 +81,7 @@ function render(
     render!(kab, (Int.(BLOCK)...,), (width, height))(
         # Outputs.
         image, rast.istate.n_contrib, rast.istate.accum_α,
+        covisibilities, uncertainties,
         # Inputs.
         rast.bstate.gaussian_values_sorted,
         _as_T(SVector{2, Float32}, means_2d),
@@ -145,10 +149,13 @@ function ChainRulesCore.rrule(::typeof(render),
     colors::AbstractMatrix{Float32};
     rast::GaussianRasterizer, camera::Camera, background::SVector{3, Float32},
     depths::AbstractVector{Float32},
+
+    covisibilities::Maybe{AbstractVector{Bool}} = nothing,
+    uncertainties::Maybe{AbstractMatrix{Float32}} = nothing,
 )
     image = render(
         means_2d, conics, opacities, colors;
-        rast, camera, background, depths)
+        rast, camera, background, depths, covisibilities, uncertainties)
     function _render_pullback(vpixels)
         ∇ = ∇render(
             unthunk(vpixels), means_2d, conics, opacities, colors;
@@ -163,6 +170,9 @@ end
     out_color::AbstractArray{Float32, 3},
     n_contrib::AbstractMatrix{UInt32},
     accum_α::AbstractMatrix{Float32},
+
+    covisibilities::C,
+    uncertainties::U,
     # Input.
     gaussian_values_sorted::AbstractVector{UInt32},
     means_2d::AbstractVector{SVector{2, Float32}},
@@ -175,7 +185,11 @@ end
     background::SVector{channels, Float32},
     block::SVector{2, Int32},
     ::Val{block_size},
-) where {block_size, channels}
+) where {
+    block_size, channels,
+    C <: Maybe{AbstractVector{Bool}},
+    U <: Maybe{AbstractMatrix{Float32}},
+}
     gidx = @index(Group, NTuple) # ≡ group_index
     lidx = @index(Local, NTuple) # ≡ thread_index
     ridx = @index(Local)         # ≡ thread_rank
@@ -215,6 +229,7 @@ end
     last_contributor = 0u32
 
     color = zeros(MVector{channels, Float32})
+    uncertainty = 0f0 # Computed if `uncertainties ≢ nothing`.
     for round in 0i32:(rounds - 1i32)
         # Collectively fetch data from global to shared memory.
         progress = range[1] + block_size * round + ridx # 1-based.
@@ -258,6 +273,11 @@ end
                 color[c] += feature[c] * α * T
             end
 
+            U !== Nothing && (uncertainty += α * T)
+            # Since we are processing Gaussians front-to-back,
+            # mark visible only until cummulative opacity is > 0.5.
+            C !== Nothing && T > 0.5f0 && (covisibilities[gaussian_id] = true)
+
             # Keep track of last range entry to update this pixel.
             T = T_tmp
             last_contributor = contributor
@@ -272,6 +292,7 @@ end
         @unroll for c in 1i32:channels
             out_color[c, px, py] = color[c] + T * background[c]
         end
+        U !== Nothing && (uncertainties[px, py] = uncertainty)
     end
 end
 

src/training.jl

@@ -25,6 +25,10 @@ function Trainer(
     rast::GaussianRasterizer, gs::GaussianModel,
     dataset::ColmapDataset, opt_params::OptimizationParams;
 )
+    # If we are going to use trainer, invalidate its alloc cache.
+    AT = typeof(gs.points)
+    GPUArrays.AllocCache.invalidate!(AT, :train_step)
+
     ϵ = 1f-15
     kab = get_backend(gs)
     camera_extent = dataset.camera_extent