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.

Sign up for GitHub

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

Support Scaniverse SPZ file format input for creating gaussian assets #161

Merged
merged 2 commits into from
Jan 30, 2025
Merged

Support Scaniverse SPZ file format input for creating gaussian assets #161

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
5a39f92
Support reading Scaniverse SPZ files
aras-p Jan 30, 2025
ad43176
Readme
aras-p Jan 30, 2025
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
109 changes: 12 additions & 97 deletions package/Editor/GaussianSplatAssetCreator.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,7 +47,7 @@ enum DataQuality
[SerializeField] GaussianSplatAsset.SHFormat m_FormatSH;

string m_ErrorMessage;
string m_PrevPlyPath;
string m_PrevFilePath;
int m_PrevVertexCount;
long m_PrevFileSize;

Expand Down Expand Up @@ -82,14 +82,14 @@ void OnGUI()
EditorGUILayout.Space();
GUILayout.Label("Input data", EditorStyles.boldLabel);
var rect = EditorGUILayout.GetControlRect(true);
m_InputFile = m_FilePicker.PathFieldGUI(rect, new GUIContent("Input PLY File"), m_InputFile, "ply", "PointCloudFile");
m_InputFile = m_FilePicker.PathFieldGUI(rect, new GUIContent("Input PLY/SPZ File"), m_InputFile, "ply,spz", "PointCloudFile");
m_ImportCameras = EditorGUILayout.Toggle("Import Cameras", m_ImportCameras);

if (m_InputFile != m_PrevPlyPath && !string.IsNullOrWhiteSpace(m_InputFile))
if (m_InputFile != m_PrevFilePath && !string.IsNullOrWhiteSpace(m_InputFile))
{
PLYFileReader.ReadFileHeader(m_InputFile, out m_PrevVertexCount, out var _, out var _);
m_PrevVertexCount = GaussianFileReader.ReadFileHeader(m_InputFile);
m_PrevFileSize = File.Exists(m_InputFile) ? new FileInfo(m_InputFile).Length : 0;
m_PrevPlyPath = m_InputFile;
m_PrevFilePath = m_InputFile;
}

if (m_PrevVertexCount > 0)
Expand Down Expand Up @@ -217,17 +217,6 @@ void ApplyQualityLevel()
}
}

// input file splat data is expected to be in this format
public struct InputSplatData
{
public Vector3 pos;
public Vector3 nor;
public Vector3 dc0;
public Vector3 sh1, sh2, sh3, sh4, sh5, sh6, sh7, sh8, sh9, shA, shB, shC, shD, shE, shF;
public float opacity;
public Vector3 scale;
public Quaternion rot;
}

static T CreateOrReplaceAsset<T>(T asset, string path) where T : UnityEngine.Object
{
Expand All @@ -250,7 +239,7 @@ unsafe void CreateAsset()
m_ErrorMessage = null;
if (string.IsNullOrWhiteSpace(m_InputFile))
{
m_ErrorMessage = $"Select input PLY file";
m_ErrorMessage = $"Select input PLY/SPZ file";
return;
}

Expand All @@ -263,7 +252,7 @@ unsafe void CreateAsset()

EditorUtility.DisplayProgressBar(kProgressTitle, "Reading data files", 0.0f);
GaussianSplatAsset.CameraInfo[] cameras = LoadJsonCamerasFile(m_InputFile, m_ImportCameras);
using NativeArray<InputSplatData> inputSplats = LoadPLYSplatFile(m_InputFile);
using NativeArray<InputSplatData> inputSplats = LoadInputSplatFile(m_InputFile);
if (inputSplats.Length == 0)
{
EditorUtility.ClearProgressBar();
Expand Down Expand Up @@ -304,7 +293,6 @@ unsafe void CreateAsset()
string pathOther = $"{m_OutputFolder}/{baseName}_oth.bytes";
string pathCol = $"{m_OutputFolder}/{baseName}_col.bytes";
string pathSh = $"{m_OutputFolder}/{baseName}_shs.bytes";
LinearizeData(inputSplats);

// if we are using full lossless (FP32) data, then do not use any chunking, and keep data as-is
bool useChunks = isUsingChunks;
Expand Down Expand Up @@ -341,64 +329,23 @@ unsafe void CreateAsset()
Selection.activeObject = savedAsset;
}

unsafe NativeArray<InputSplatData> LoadPLYSplatFile(string plyPath)
NativeArray<InputSplatData> LoadInputSplatFile(string filePath)
{
NativeArray<InputSplatData> data = default;
if (!File.Exists(plyPath))
if (!File.Exists(filePath))
{
m_ErrorMessage = $"Did not find {plyPath} file";
m_ErrorMessage = $"Did not find {filePath} file";
return data;
}

int splatCount;
int vertexStride;
NativeArray<byte> verticesRawData;
try
{
PLYFileReader.ReadFile(plyPath, out splatCount, out vertexStride, out _, out verticesRawData);
GaussianFileReader.ReadFile(filePath, out data);
}
catch (Exception ex)
{
m_ErrorMessage = ex.Message;
return data;
}

if (UnsafeUtility.SizeOf<InputSplatData>() != vertexStride)
{
m_ErrorMessage = $"PLY vertex size mismatch, expected {UnsafeUtility.SizeOf<InputSplatData>()} but file has {vertexStride}";
return data;
}

// reorder SHs
NativeArray<float> floatData = verticesRawData.Reinterpret<float>(1);
ReorderSHs(splatCount, (float*)floatData.GetUnsafePtr());

return verticesRawData.Reinterpret<InputSplatData>(1);
}

[BurstCompile]
static unsafe void ReorderSHs(int splatCount, float* data)
{
int splatStride = UnsafeUtility.SizeOf<InputSplatData>() / 4;
int shStartOffset = 9, shCount = 15;
float* tmp = stackalloc float[shCount * 3];
int idx = shStartOffset;
for (int i = 0; i < splatCount; ++i)
{
for (int j = 0; j < shCount; ++j)
{
tmp[j * 3 + 0] = data[idx + j];
tmp[j * 3 + 1] = data[idx + j + shCount];
tmp[j * 3 + 2] = data[idx + j + shCount * 2];
}

for (int j = 0; j < shCount * 3; ++j)
{
data[idx + j] = tmp[j];
}

idx += splatStride;
}
return data;
}

[BurstCompile]
Expand Down Expand Up @@ -560,38 +507,6 @@ static unsafe void ClusterSHs(NativeArray<InputSplatData> splatData, GaussianSpl
Debug.Log($"GS: clustered {splatData.Length/1000000.0:F2}M SHs into {shCount/1024}K ({passesOverData:F1}pass/{kBatchSize}batch) in {t1-t0:F0}s");
}

[BurstCompile]
struct LinearizeDataJob : IJobParallelFor
{
public NativeArray<InputSplatData> splatData;
public void Execute(int index)
{
var splat = splatData[index];

// rot
var q = splat.rot;
var qq = GaussianUtils.NormalizeSwizzleRotation(new float4(q.x, q.y, q.z, q.w));
qq = GaussianUtils.PackSmallest3Rotation(qq);
splat.rot = new Quaternion(qq.x, qq.y, qq.z, qq.w);

// scale
splat.scale = GaussianUtils.LinearScale(splat.scale);

// color
splat.dc0 = GaussianUtils.SH0ToColor(splat.dc0);
splat.opacity = GaussianUtils.Sigmoid(splat.opacity);

splatData[index] = splat;
}
}

static void LinearizeData(NativeArray<InputSplatData> splatData)
{
LinearizeDataJob job = new LinearizeDataJob();
job.splatData = splatData;
job.Schedule(splatData.Length, 4096).Complete();
}

[BurstCompile]
struct CalcChunkDataJob : IJobParallelFor
{
Expand Down
4 changes: 2 additions & 2 deletions package/Editor/GaussianSplatRendererEditor.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -398,13 +398,13 @@ static unsafe void ExportPlyFile(GaussianSplatRenderer gs, bool bakeTransform)
if (string.IsNullOrWhiteSpace(path))
return;

int kSplatSize = UnsafeUtility.SizeOf<GaussianSplatAssetCreator.InputSplatData>();
int kSplatSize = UnsafeUtility.SizeOf<Utils.InputSplatData>();
using var gpuData = new GraphicsBuffer(GraphicsBuffer.Target.Structured, gs.splatCount, kSplatSize);

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

GaussianSplatAssetCreator.InputSplatData[] data = new GaussianSplatAssetCreator.InputSplatData[gpuData.count];
Utils.InputSplatData[] data = new Utils.InputSplatData[gpuData.count];
gpuData.GetData(data);

var gpuDeleted = gs.GpuEditDeleted;
Expand Down
128 changes: 128 additions & 0 deletions package/Editor/Utils/GaussianFileReader.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,128 @@
using System.Collections.Generic;
using System.Globalization;
using System.IO;
using GaussianSplatting.Runtime;
using Unity.Burst;
using Unity.Collections;
using Unity.Collections.LowLevel.Unsafe;
using Unity.Jobs;
using Unity.Mathematics;
using UnityEngine;

namespace GaussianSplatting.Editor.Utils
{
// input file splat data is expected to be in this format
public struct InputSplatData
{
public Vector3 pos;
public Vector3 nor;
public Vector3 dc0;
public Vector3 sh1, sh2, sh3, sh4, sh5, sh6, sh7, sh8, sh9, shA, shB, shC, shD, shE, shF;
public float opacity;
public Vector3 scale;
public Quaternion rot;
}

[BurstCompile]
public class GaussianFileReader
{
// Returns splat count
public static int ReadFileHeader(string filePath)
{
int vertexCount = 0;
if (File.Exists(filePath))
{
if (isPLY(filePath))
PLYFileReader.ReadFileHeader(filePath, out vertexCount, out _, out _);
else if (isSPZ(filePath))
SPZFileReader.ReadFileHeader(filePath, out vertexCount);
}
return vertexCount;
}

public static unsafe void ReadFile(string filePath, out NativeArray<InputSplatData> splats)
{
if (isPLY(filePath))
{
NativeArray<byte> verticesRawData;
PLYFileReader.ReadFile(filePath, out var splatCount, out var splatStride, out List<string> _, out verticesRawData);
if (UnsafeUtility.SizeOf<InputSplatData>() != splatStride)
throw new IOException($"PLY vertex size mismatch, expected {UnsafeUtility.SizeOf<InputSplatData>()} but file has {splatStride}");

// reorder SHs
NativeArray<float> floatData = verticesRawData.Reinterpret<float>(1);
ReorderSHs(splatCount, (float*)floatData.GetUnsafePtr());

splats = verticesRawData.Reinterpret<InputSplatData>(1);
LinearizeData(splats);
return;
}
if (isSPZ(filePath))
{
SPZFileReader.ReadFile(filePath, out splats);
return;
}
throw new IOException($"File {filePath} is not a supported format");
}

static bool isPLY(string filePath) => filePath.EndsWith(".ply", true, CultureInfo.InvariantCulture);
static bool isSPZ(string filePath) => filePath.EndsWith(".spz", true, CultureInfo.InvariantCulture);

[BurstCompile]
static unsafe void ReorderSHs(int splatCount, float* data)
{
int splatStride = UnsafeUtility.SizeOf<InputSplatData>() / 4;
int shStartOffset = 9, shCount = 15;
float* tmp = stackalloc float[shCount * 3];
int idx = shStartOffset;
for (int i = 0; i < splatCount; ++i)
{
for (int j = 0; j < shCount; ++j)
{
tmp[j * 3 + 0] = data[idx + j];
tmp[j * 3 + 1] = data[idx + j + shCount];
tmp[j * 3 + 2] = data[idx + j + shCount * 2];
}

for (int j = 0; j < shCount * 3; ++j)
{
data[idx + j] = tmp[j];
}

idx += splatStride;
}
}

[BurstCompile]
struct LinearizeDataJob : IJobParallelFor
{
public NativeArray<InputSplatData> splatData;
public void Execute(int index)
{
var splat = splatData[index];

// rot
var q = splat.rot;
var qq = GaussianUtils.NormalizeSwizzleRotation(new float4(q.x, q.y, q.z, q.w));
qq = GaussianUtils.PackSmallest3Rotation(qq);
splat.rot = new Quaternion(qq.x, qq.y, qq.z, qq.w);

// scale
splat.scale = GaussianUtils.LinearScale(splat.scale);

// color
splat.dc0 = GaussianUtils.SH0ToColor(splat.dc0);
splat.opacity = GaussianUtils.Sigmoid(splat.opacity);

splatData[index] = splat;
}
}

static void LinearizeData(NativeArray<InputSplatData> splatData)
{
LinearizeDataJob job = new LinearizeDataJob();
job.splatData = splatData;
job.Schedule(splatData.Length, 4096).Complete();
}
}
}
3 changes: 3 additions & 0 deletions package/Editor/Utils/GaussianFileReader.cs.meta
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

Loading