Add PLY export & fix format

Author: pxl-th

src/GaussianSplatting.jl

@@ -154,9 +154,9 @@ end
 function gui(path::String; scale::Maybe{Int} = nothing, fullscreen::Bool = false)
     ispath(path) || error("Path does not exist: `$path`.")
 
-    viewer_mode = endswith(path, ".bson")
+    viewer_mode = endswith(path, ".bson") || endswith(path, ".ply")
     !viewer_mode && !isdir(path) && error(
-        "`path` must be either a `.bson` model checkpoint or " *
+        "`path` must be either a [`.bson`, '.ply'] model checkpoint or " *
         "a directory with COLMAP dataset, instead: `$path`.")
     !viewer_mode && scale ≡ nothing && error(
         "`scale` keyword argument must be specified if `path` is a COLMAP dataset.")
@@ -166,10 +166,18 @@ function gui(path::String; scale::Maybe{Int} = nothing, fullscreen::Bool = false
         (1024, 1024, true)
 
     gui = if viewer_mode
-        θ = BSON.load(path)
-        gaussians = GaussianModel(gpu_backend())
-        set_from_bson!(gaussians, θ[:gaussians])
-        camera = θ[:camera]
+        kab = gpu_backend()
+        if endswith(path, ".bson")
+            θ = BSON.load(path)
+            gaussians = GaussianModel(kab)
+            set_from_bson!(gaussians, θ[:gaussians])
+            camera = θ[:camera]
+        else
+            gaussians = import_ply(path, kab)
+            width = 1024
+            fov = NU.fov2focal(1024, 45f0)
+            camera = Camera(; fx=fov, fy=fov, width, height=width)
+        end
 
         GSGUI(gaussians, camera; width, height, fullscreen, resizable)
     else

src/gaussians.jl

@@ -157,7 +157,7 @@ sh_2_rgb(x) = x * SH0 + 0.5f0
 
 inverse_sigmoid(x) = log(x / (1f0 - x))
 
-function export_ply!(g::GaussianModel, filename::String)
+function export_ply(g::GaussianModel, filename::String)
     ply = PlyIO.Ply()
 
     n = size(g.points, 2)
@@ -174,16 +174,16 @@ function export_ply!(g::GaussianModel, filename::String)
         PlyIO.ArrayProperty("y", xyz[2, :]),
         PlyIO.ArrayProperty("z", xyz[3, :]),
 
-        [PlyIO.ArrayProperty("f_dc_$i", features_dc[i, :]) for i in 1:size(features_dc, 1)]...,
-        [PlyIO.ArrayProperty("f_rest_$i", features_rest[i, :]) for i in 1:size(features_rest, 1)]...,
-        [PlyIO.ArrayProperty("scale_$i", scales[i, :]) for i in 1:size(scales, 1)]...,
-        [PlyIO.ArrayProperty("rot_$i", rotations[i, :]) for i in 1:size(rotations, 1)]...,
+        [PlyIO.ArrayProperty("f_dc_$(i - 1)", features_dc[i, :]) for i in 1:size(features_dc, 1)]...,
+        [PlyIO.ArrayProperty("f_rest_$(i - 1)", features_rest[i, :]) for i in 1:size(features_rest, 1)]...,
+        [PlyIO.ArrayProperty("scale_$(i - 1)", scales[i, :]) for i in 1:size(scales, 1)]...,
+        [PlyIO.ArrayProperty("rot_$(i - 1)", rotations[i, :]) for i in 1:size(rotations, 1)]...,
 
         PlyIO.ArrayProperty("opacity", opacities),
     )
     push!(ply, vertex)
 
-    PlyIO.save_ply(ply, filename)
+    PlyIO.save_ply(ply, filename; ascii=false)
     return
 end
 
@@ -196,14 +196,14 @@ function import_ply(filename::String, kab)
 
     n = length(vertex["x"])
     xyz = vcat([reshape(vertex[i], 1, n) for i in ("x", "y", "z")]...)
-    scales = vcat([reshape(vertex["scale_$i"], 1, n) for i in 1:3]...)
-    rotations = vcat([reshape(vertex["rot_$i"], 1, n) for i in 1:4]...)
+    scales = vcat([reshape(vertex["scale_$(i - 1)"], 1, n) for i in 1:3]...)
+    rotations = vcat([reshape(vertex["rot_$(i - 1)"], 1, n) for i in 1:4]...)
     opacities = reshape(Array(vertex["opacity"]), 1, n)
 
-    features_dc = vcat([reshape(vertex["f_dc_$i"], 1, 1, n) for i in 1:3]...)
+    features_dc = vcat([reshape(vertex["f_dc_$(i - 1)"], 1, 1, n) for i in 1:3]...)
     features_rest = if n_frest > 0
         reshape(
-            vcat([reshape(vertex["f_rest_$i"], 1, n) for i in 1:n_frest]...),
+            vcat([reshape(vertex["f_rest_$(i - 1)"], 1, n) for i in 1:n_frest]...),
             3, :, n)
     else
         Array{Float32}(undef, 3, 0, n)

src/gui/gui.jl

@@ -354,6 +354,17 @@ function handle_ui!(gui::GSGUI; frame_time)
                 end
                 CImGui.Separator()
 
+                if CImGui.Button("Export PLY", CImGui.ImVec2(-1, 0))
+                    save_dir = unsafe_string(pointer(gui.ui_state.save_directory_path))
+                    isdir(save_dir) || mkpath(save_dir)
+
+                    tstmp = now()
+                    fmt = "timestamp-$(month(tstmp))M-$(day(tstmp))D-$(hour(tstmp)):$(minute(tstmp))"
+                    save_file = joinpath(save_dir, "state-(step-$(gui.trainer.step))-($fmt).ply")
+                    export_ply(gui.gaussians, save_file)
+                end
+                CImGui.Separator()
+
                 CImGui.Text("Path to State File (.bson):")
                 CImGui.PushItemWidth(-1)
                 CImGui.InputText(