Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Editing: merge, move splats, bake transform when exporting #51

Merged
merged 27 commits into from
Nov 9, 2023
Merged
Changes from 1 commit
Commits
Show all changes
27 commits
Select commit Hold shift + click to select a range
3d96c00
Validator against both SBIR and our own results
aras-p Oct 18, 2023
1a03ecd
Make Very High quality preset not use any chunking at all
aras-p Oct 18, 2023
7b8c256
Split off GaussianSplatsTool into separate file
aras-p Oct 19, 2023
2d174fb
Change gaussian edit tool to tool context
aras-p Oct 19, 2023
6e51203
Beginnings of move tool
aras-p Oct 19, 2023
99b1f21
Move tool handles handle orientation, renaming variables for clarity
aras-p Oct 19, 2023
bb811cb
Beginnings of scale tool
aras-p Oct 19, 2023
431a514
Scale tool handles global vs local orientation modes
aras-p Oct 20, 2023
ee38ced
Add initial SH rotation code
aras-p Oct 20, 2023
d036594
SH rotation matrix column vs row major fix
aras-p Oct 20, 2023
626e710
Fix regular (debug points, debug boxes) shaders not working
aras-p Oct 20, 2023
b0b1a82
Add option to "bake" GS transform when exporting PLY file
aras-p Oct 20, 2023
0fbc310
Scale splats by transfrom local scale when exporting
aras-p Oct 20, 2023
018b5ff
Beginnings of rotate tool (not working properly yet)
aras-p Oct 20, 2023
9c4389e
Fix cutouts not working properly with bake transform lol
aras-p Oct 21, 2023
68c405b
Towards having ability for splat count in renderer to not be the same…
aras-p Oct 21, 2023
ec79cdc
Towards rotate tool, but not working fine just yet
aras-p Oct 25, 2023
88ee1df
Basic multi-select inspector for GS assets
aras-p Oct 25, 2023
f632550
Towards ability to merge splats (SH is not correct yet)
aras-p Nov 9, 2023
41e2b87
Add option to render only the SH contribution upon gray color
aras-p Nov 9, 2023
6907cce
Fix SH rotation to finally be proper :fingers-crossed:
aras-p Nov 9, 2023
e98549a
Merge branch 'main' into more-edit-tools
aras-p Nov 9, 2023
6160df7
Always display splat count, merge marks as edited, Reset button
aras-p Nov 9, 2023
b35b60d
Carry over "which splats are deleted" info when merging splats
aras-p Nov 9, 2023
e9bb447
Clearer label for bake transform
aras-p Nov 9, 2023
2f87543
Disable scale tool for now
aras-p Nov 9, 2023
6047c10
Cleanup
aras-p Nov 9, 2023
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
Prev Previous commit
Next Next commit
Add option to "bake" GS transform when exporting PLY file
Allow exporting even when no edits are done (e.g. when we only want to bake in new transform)
aras-p committed Oct 20, 2023

Verified

This commit was signed with the committer’s verified signature.
jeffkala Jeff Kala
commit b0b1a82421204cb924108c1675f37a536fc1a688
57 changes: 36 additions & 21 deletions package/Editor/GaussianSplatRendererEditor.cs
Original file line number Diff line number Diff line change
@@ -17,6 +17,8 @@ namespace GaussianSplatting.Editor
[CustomEditor(typeof(GaussianSplatRenderer))]
public class GaussianSplatRendererEditor : UnityEditor.Editor
{
const string kPrefExportBake = "nesnausk.GaussianSplatting.ExportBakeTransform";

SerializedProperty m_PropAsset;
SerializedProperty m_PropSplatScale;
SerializedProperty m_PropOpacityScale;
@@ -34,6 +36,8 @@ public class GaussianSplatRendererEditor : UnityEditor.Editor
bool m_ResourcesExpanded = false;
int m_CameraIndex = 0;

bool m_ExportBakeTransform;

static int s_EditStatsUpdateCounter = 0;

static HashSet<GaussianSplatRendererEditor> s_AllEditors = new();
@@ -51,6 +55,8 @@ public static void RepaintAll()

public void OnEnable()
{
m_ExportBakeTransform = EditorPrefs.GetBool(kPrefExportBake, false);

m_PropAsset = serializedObject.FindProperty("m_Asset");
m_PropSplatScale = serializedObject.FindProperty("m_SplatScale");
m_PropOpacityScale = serializedObject.FindProperty("m_OpacityScale");
@@ -201,34 +207,39 @@ void EditGUI(GaussianSplatRenderer gs)
}
GUILayout.EndHorizontal();
EditorGUILayout.PropertyField(m_PropCutouts);
EditorGUILayout.Space();

bool hasCutouts = gs.m_Cutouts != null && gs.m_Cutouts.Length != 0;
bool modifiedOrHasCutouts = gs.editModified || hasCutouts;
bool displayEditTools = isToolActive || modifiedOrHasCutouts;

if (displayEditTools)
var asset = gs.asset;
EditorGUILayout.Space();
EditorGUI.BeginChangeCheck();
m_ExportBakeTransform = EditorGUILayout.Toggle("Bake Transform", m_ExportBakeTransform);
if (EditorGUI.EndChangeCheck())
{
var asset = gs.asset;
EditorPrefs.SetBool(kPrefExportBake, m_ExportBakeTransform);
}

if (GUILayout.Button("Export PLY"))
ExportPlyFile(gs, m_ExportBakeTransform);
if (asset.m_PosFormat > GaussianSplatAsset.VectorFormat.Norm16 ||
asset.m_ScaleFormat > GaussianSplatAsset.VectorFormat.Norm16 ||
!GraphicsFormatUtility.IsFloatFormat(asset.m_ColorFormat) ||
asset.m_SHFormat > GaussianSplatAsset.SHFormat.Float16)
{
EditorGUILayout.HelpBox(
"It is recommended to use High or VeryHigh quality preset for editing splats, lower levels are lossy",
MessageType.Warning);
}

bool displayEditStats = isToolActive || modifiedOrHasCutouts;
if (displayEditStats)
{
EditorGUILayout.Space();
EditorGUILayout.LabelField("Splats", $"{asset.m_SplatCount:N0}");
EditorGUILayout.LabelField("Cut", $"{gs.editCutSplats:N0}");
EditorGUILayout.LabelField("Deleted", $"{gs.editDeletedSplats:N0}");
EditorGUILayout.LabelField("Selected", $"{gs.editSelectedSplats:N0}");
if (modifiedOrHasCutouts)
{
if (GUILayout.Button("Export modified PLY"))
ExportPlyFile(gs);
if (asset.m_PosFormat > GaussianSplatAsset.VectorFormat.Norm16 ||
asset.m_ScaleFormat > GaussianSplatAsset.VectorFormat.Norm16 ||
!GraphicsFormatUtility.IsFloatFormat(asset.m_ColorFormat) ||
asset.m_SHFormat > GaussianSplatAsset.SHFormat.Float16)
{
EditorGUILayout.HelpBox(
"It is recommended to use High or VeryHigh quality preset for editing splats, lower levels are lossy",
MessageType.Warning);
}
}

if (hasCutouts)
{
if (s_EditStatsUpdateCounter > 10)
@@ -277,15 +288,17 @@ public static Bounds TransformBounds(Transform tr, Bounds bounds )
return new Bounds { center = center, extents = ext };
}

static unsafe void ExportPlyFile(GaussianSplatRenderer gs)
static unsafe void ExportPlyFile(GaussianSplatRenderer gs, bool bakeTransform)
{
var path = EditorUtility.SaveFilePanel(
"Export Gaussian Splat PLY file", "", $"{gs.asset.name}-edit.ply", "ply");
if (string.IsNullOrWhiteSpace(path))
return;

int kSplatSize = UnsafeUtility.SizeOf<GaussianSplatAssetCreator.InputSplatData>();
using var gpuData = new GraphicsBuffer(GraphicsBuffer.Target.Structured, gs.asset.m_SplatCount, kSplatSize);
if (!gs.EditExportData(gpuData))

if (!gs.EditExportData(gpuData, bakeTransform))
return;

GaussianSplatAssetCreator.InputSplatData[] data = new GaussianSplatAssetCreator.InputSplatData[gpuData.count];
@@ -324,6 +337,8 @@ static unsafe void ExportPlyFile(GaussianSplatRenderer gs)
fs.Write(new ReadOnlySpan<byte>(ptr, kSplatSize));
}
}

Debug.Log($"Exported PLY {path} with {aliveCount:N0} splats");
}
}
}
18 changes: 15 additions & 3 deletions package/Runtime/GaussianSplatRenderer.cs
Original file line number Diff line number Diff line change
@@ -130,15 +130,13 @@ public Material SortAndRenderSplats(Camera cam, CommandBuffer cmb)

gs.SetAssetDataOnMaterial(mpb);
mpb.SetBuffer(GaussianSplatRenderer.Props.SplatChunks, gs.m_GpuChunks);
mpb.SetInteger(GaussianSplatRenderer.Props.SplatChunkCount, gs.m_GpuChunksValid ? gs.m_GpuChunks.count : 0);

mpb.SetBuffer(GaussianSplatRenderer.Props.SplatViewData, gs.m_GpuView);

mpb.SetBuffer(GaussianSplatRenderer.Props.OrderBuffer, gs.m_GpuSortKeys);
mpb.SetFloat(GaussianSplatRenderer.Props.SplatScale, gs.m_SplatScale);
mpb.SetFloat(GaussianSplatRenderer.Props.SplatOpacityScale, gs.m_OpacityScale);
mpb.SetFloat(GaussianSplatRenderer.Props.SplatSize, gs.m_PointDisplaySize);
mpb.SetInteger(GaussianSplatRenderer.Props.SplatCount, gs.asset.m_SplatCount);
mpb.SetInteger(GaussianSplatRenderer.Props.SHOrder, gs.m_SHOrder);
mpb.SetInteger(GaussianSplatRenderer.Props.DisplayIndex, gs.m_RenderMode == GaussianSplatRenderer.RenderMode.DebugPointIndices ? 1 : 0);
mpb.SetInteger(GaussianSplatRenderer.Props.DisplayChunks, gs.m_RenderMode == GaussianSplatRenderer.RenderMode.DebugChunkBounds ? 1 : 0);
@@ -461,6 +459,8 @@ internal void SetAssetDataOnMaterial(MaterialPropertyBlock mat)
mat.SetInt(Props.SplatBitsValid, m_GpuEditSelected != null && m_GpuEditDeleted != null ? 1 : 0);
uint format = (uint)m_Asset.m_PosFormat | ((uint)m_Asset.m_ScaleFormat << 8) | ((uint)m_Asset.m_SHFormat << 16);
mat.SetInteger(Props.SplatFormat, (int)format);
mat.SetInteger(Props.SplatCount, m_Asset.m_SplatCount);
mat.SetInteger(Props.SplatChunkCount, m_GpuChunksValid ? m_GpuChunks.count : 0);
}

static void DisposeBuffer(ref GraphicsBuffer buf)
@@ -836,12 +836,24 @@ public void EditInvertSelection()
UpdateEditCountsAndBounds();
}

public bool EditExportData(GraphicsBuffer dstData)
public bool EditExportData(GraphicsBuffer dstData, bool bakeTransform)
{
if (!EnsureEditingBuffers()) return false;

int flags = 0;
var tr = transform;
Quaternion bakeRot = tr.localRotation;
Vector3 bakeScale = tr.localScale;

if (bakeTransform)
flags = 1;

using var cmb = new CommandBuffer { name = "SplatExportData" };
SetAssetDataOnCS(cmb, KernelIndices.ExportData);
cmb.SetComputeIntParam(m_CSSplatUtilities, "_ExportTransformFlags", flags);
cmb.SetComputeVectorParam(m_CSSplatUtilities, "_ExportTransformRotation", new Vector4(bakeRot.x, bakeRot.y, bakeRot.z, bakeRot.w));
cmb.SetComputeVectorParam(m_CSSplatUtilities, "_ExportTransformScale", bakeScale);
cmb.SetComputeMatrixParam(m_CSSplatUtilities, Props.MatrixObjectToWorld, tr.localToWorldMatrix);
cmb.SetComputeBufferParam(m_CSSplatUtilities, (int)KernelIndices.ExportData, "_ExportBuffer", dstData);

DispatchUtilsAndExecute(cmb, KernelIndices.ExportData, m_Asset.m_SplatCount);
5 changes: 5 additions & 0 deletions package/Shaders/GaussianSplatting.hlsl
Original file line number Diff line number Diff line change
@@ -16,6 +16,11 @@ float3 QuatRotateVector(float3 v, float4 r)
return v + r.w * t + cross(r.xyz, t);
}

float4 QuatMul(float4 a, float4 b)
{
return float4(a.wwww * b + (a.xyzx * b.wwwx + a.yzxy * b.zxyy) * float4(1,1,1,-1) - a.zxyz * b.yzxz);
}

float4 QuatInverse(float4 q)
{
return rcp(dot(q, q)) * q * float4(-1,-1,-1,1);
97 changes: 97 additions & 0 deletions package/Shaders/SplatUtilities.compute
Original file line number Diff line number Diff line change
@@ -490,13 +490,110 @@ float InvSigmoid(float v)
return log(v / max(1 - v, 1.0e-6));
}

// SH rotation
#include "SphericalHarmonics.hlsl"

void RotateSH(inout SplatSHData sh, float3x3 rot)
{
float3 shin[16];
float3 shout[16];
shin[0] = sh.col;
shin[1] = sh.sh1;
shin[2] = sh.sh2;
shin[3] = sh.sh3;
shin[4] = sh.sh4;
shin[5] = sh.sh5;
shin[6] = sh.sh6;
shin[7] = sh.sh7;
shin[8] = sh.sh8;
shin[9] = sh.sh9;
shin[10] = sh.sh10;
shin[11] = sh.sh11;
shin[12] = sh.sh12;
shin[13] = sh.sh13;
shin[14] = sh.sh14;
shin[15] = sh.sh15;
RotateSH(rot, 4, shin, shout);
sh.col = shout[0];
sh.sh1 = shout[1];
sh.sh2 = shout[2];
sh.sh3 = shout[3];
sh.sh4 = shout[4];
sh.sh5 = shout[5];
sh.sh6 = shout[6];
sh.sh7 = shout[7];
sh.sh8 = shout[8];
sh.sh9 = shout[9];
sh.sh10 = shout[10];
sh.sh11 = shout[11];
sh.sh12 = shout[12];
sh.sh13 = shout[13];
sh.sh14 = shout[14];
sh.sh15 = shout[15];
}

float3x3 CalcSHRotMatrix(float4x4 objToWorld)
{
float3x3 m = (float3x3)objToWorld;
float sx = length(float3(m[0][0], m[0][1], m[0][2]));
float sy = length(float3(m[1][0], m[1][1], m[1][2]));
float sz = length(float3(m[2][0], m[2][1], m[2][2]));

float invSX = 1.0 / sx;
float invSY = 1.0 / sy;
float invSZ = 1.0 / sz;

m[0][0] *= invSX;
m[0][1] *= invSX;
m[0][2] *= invSX;
m[1][0] *= invSY;
m[1][1] *= invSY;
m[1][2] *= invSY;
m[2][0] *= invSZ;
m[2][1] *= invSZ;
m[2][2] *= invSZ;

float det = determinant(m);
if (det < 0) {
m[0][2] = -m[0][2];
m[2][0] = -m[2][0];
m[1][2] = -m[1][2];
m[2][1] = -m[2][1];
}
return m;
}


float4 _ExportTransformRotation;
float3 _ExportTransformScale;
uint _ExportTransformFlags;

[numthreads(GROUP_SIZE,1,1)]
void CSExportData (uint3 id : SV_DispatchThreadID)
{
uint idx = id.x;
if (idx >= _SplatCount)
return;
SplatData src = LoadSplatData(idx);

// transform splat by matrix, if needed
if (_ExportTransformFlags != 0)
{
src.pos = mul(_MatrixObjectToWorld, float4(src.pos,1)).xyz;

// note: this only handles axis flips from scale, not any arbitrary scaling
if (_ExportTransformScale.x < 0)
src.rot.yz = -src.rot.yz;
if (_ExportTransformScale.y < 0)
src.rot.xz = -src.rot.xz;
if (_ExportTransformScale.z < 0)
src.rot.xy = -src.rot.xy;
src.rot = QuatMul(_ExportTransformRotation, src.rot);

float3x3 shRot = CalcSHRotMatrix(_MatrixObjectToWorld);
RotateSH(src.sh, shRot);
}

ExportSplatData dst;
dst.pos = src.pos;
dst.nor = 0;