diff --git a/Assets/vfx-graph-gs/Scripts/SplatImporter.cs b/Assets/vfx-graph-gs/Scripts/SplatImporter.cs index 813e067..193a8b6 100644 --- a/Assets/vfx-graph-gs/Scripts/SplatImporter.cs +++ b/Assets/vfx-graph-gs/Scripts/SplatImporter.cs @@ -18,19 +18,10 @@ using Unity.VisualScripting; using System.Collections; using Unity.Collections; using UnityEditor.Experimental.GraphView; - -/******************************************************************************* - TODO: - 1. Move header parsing to ParsePLYFile - 2. Improve exception reporting and handling - 3. Use Burst compiler for parsing binary data - 4. If we can reasonably assume that items are contigious (e.g. position) - use MemoryMarshal - 5. Use a struct for header_info - 6. Write handler for data other than 32 bit float - 7. Pass around Span instead of byte[] for memory safety - 8. Implement opacity clamping in the VFX Graph -*******************************************************************************/ +using UnityEngine.Playables; +using Unity.Android.Gradle.Manifest; +using UnityEngine.Timeline; +using Unity.Jobs; namespace SplatVfx.Editor { @@ -39,16 +30,14 @@ namespace SplatVfx.Editor public sealed class SplatImporter : ScriptedImporter { const float SH_C0 = 0.28209479177387814f; + const string DefaultVfxPath = "Assets/vfx-graph-gs/SplatVFX.vfx"; - #region ScriptedImporter implementation - - [BurstCompile] public override void OnImportAsset(AssetImportContext context) { - var ply = ParsePLYFile(context.assetPath); - var header_info = ParseHeader(ply.header); + var (headerString, endHeaderPosition) = GetHeader(context.assetPath); + var headerData = ParseMetadata(headerString); - var data = ImportAsSplatData(header_info, ply.binaryData); + var data = ImportData(headerData, context.assetPath, endHeaderPosition); data.name = Path.GetFileNameWithoutExtension(context.assetPath); context.AddObjectToAsset("data", data); @@ -57,14 +46,7 @@ namespace SplatVfx.Editor context.SetMainObject(prefab); } - #endregion - - #region Prefab construction - - const string DefaultVfxPath = "Assets/vfx-graph-gs/SplatVFX.vfx"; - - [BurstCompile] - GameObject CreatePrefab(SplatData data) + private GameObject CreatePrefab(SplatData data) { var go = new GameObject(); var vfx = go.AddComponent(); @@ -78,12 +60,62 @@ namespace SplatVfx.Editor return go; } - #endregion + const int MAX_CHUNK_SIZE = 1048576 * 256; // 256MiB - #region Reader implementation + private enum TypeEnum + { + None, Float32, Float64, Float16, Int8, UInt8, Char, Short, UInt16 + } - [BurstCompile] - (string header, byte[] binaryData) ParsePLYFile(string filePath) + private enum DataEnum + { + x, y, z, f_dc_0, f_dc_1, f_dc_2, f_rest_0, f_rest_1, f_rest_2, + f_rest_3, f_rest_4, f_rest_5, f_rest_6, f_rest_7, f_rest_8, + f_rest_9, f_rest_10, f_rest_11, f_rest_12, f_rest_13, f_rest_14, + f_rest_15, f_rest_16, f_rest_17, f_rest_18, f_rest_19, f_rest_20, + f_rest_21, f_rest_22, f_rest_23, f_rest_24, f_rest_25, f_rest_26, + f_rest_27, f_rest_28, f_rest_29, f_rest_30, f_rest_31, f_rest_32, + f_rest_33, f_rest_34, f_rest_35, f_rest_36, f_rest_37, f_rest_38, + f_rest_39, f_rest_40, f_rest_41, f_rest_42, f_rest_43, f_rest_44, + opacity, scale_0, scale_1, scale_2, rot_0, rot_1, rot_2, rot_3 + } + + /*********************************************************************** + - TODO: Add support for more data types + ***********************************************************************/ + private static readonly Dictionary TypeToByteLengths = + new() + { + { TypeEnum.Float32, 4 }, + { TypeEnum.Float64, 8 }, + { TypeEnum.Float16, 2 }, // keeping this here to support half precision floats in the future + { TypeEnum.UInt8, 1 }, + { TypeEnum.Int8, 1 } + }; + + private static readonly Dictionary TypeMappings = + new() + { + { "float", TypeEnum.Float32 }, + { "double", TypeEnum.Float64 }, + { "uint8", TypeEnum.UInt8 }, + { "int8", TypeEnum.Int8 } + }; + + /*********************************************************************** + - TODO: Add endianness + - TODO: Add float size + ***********************************************************************/ + private struct HeaderData + { + public uint NumElements; + public uint BinaryDataStartOffset; + public int[] ElementByteOffsets; + public TypeEnum[] ElementTypes; + public int BytesPerElement; + } + + private static (string headerString, int endHeaderPosition) GetHeader(string filePath) { using var fileStream = new FileStream(filePath, FileMode.Open, FileAccess.Read); using var reader = new StreamReader(fileStream); @@ -100,58 +132,183 @@ namespace SplatVfx.Editor } } - long endHeaderPosition = headerBuilder.Length - headerBuilder.ToString().Split('\n').Length + 1; - fileStream.Seek(endHeaderPosition, SeekOrigin.Begin); + // Old logic that works because AppendLine adds another \n + // TODO: Fix this logic. Start with using Append instead of AppendLine + int endHeaderPosition = headerBuilder.Length - headerBuilder.ToString().Split('\n').Length + 1; - var binaryReader = new BinaryReader(fileStream); - var binaryData = binaryReader.ReadBytes((int)(fileStream.Length - endHeaderPosition)); - - return (headerBuilder.ToString(), binaryData); + return (headerBuilder.ToString(), endHeaderPosition); } - [BurstCompile] - (string format, int num_element, List<(string, string)> properties) - ParseHeader(string header) + private static HeaderData ParseMetadata(string headerString) { - using var reader = new StringReader(header); + HeaderData result = new(); string line; - string format = ""; - int num_element = 0; - var properties = new List<(string, string)>(); + result.ElementByteOffsets = new int[59]; + result.ElementTypes = new TypeEnum[59]; + + using var reader = new StringReader(headerString); while ((line = reader.ReadLine()) != null) { - if (line.StartsWith("end_header")) - { + if (line.Trim() == "end_header") break; - } var strings = line.Split(" "); switch (strings.FirstOrDefault()) { - case "format": - format = strings[1]; - - break; - case "element": - num_element = int.Parse(strings[2]); - + result.NumElements = uint.Parse(strings[2]); break; case "property": - properties.Add((strings[2].Trim(), strings[1].Trim())); + TypeEnum dataType = TypeMappings[strings[1]]; + // If in DataEnum, store byte offset and type + if (DataEnum.TryParse(strings[2], out DataEnum dataName)) + { + int ix = (int)dataName; + result.ElementByteOffsets[ix] = result.BytesPerElement; + result.ElementTypes[ix] = dataType; + } + + result.BytesPerElement += TypeToByteLengths[dataType]; break; } } - return (format, num_element, properties); + return result; } [BurstCompile] + private struct SplatDataNative + { + NativeArray position; + NativeArray color; + NativeArray sh; + NativeArray opacity; + NativeArray scale; + NativeArray rotation; + + public SplatDataNative(int count) + { + position = new NativeArray(count, Allocator.Temp); + color = new NativeArray(count, Allocator.Temp); + sh = new NativeArray(count, Allocator.Temp); + opacity = new NativeArray(count, Allocator.Temp); + scale = new NativeArray(count, Allocator.Temp); + rotation = new NativeArray(count, Allocator.Temp); + } + + public void Dispose() + { + position.Dispose(); + color.Dispose(); + sh.Dispose(); + opacity.Dispose(); + scale.Dispose(); + rotation.Dispose(); + } + } + + private static SplatData ImportData(HeaderData headerData, string filePath, int endHeaderPosition) + { + var splatData = ScriptableObject.CreateInstance(); + + int elementsPerChunk = MAX_CHUNK_SIZE / headerData.BytesPerElement; + int chunkSize = elementsPerChunk * headerData.BytesPerElement; + + using var fileStream = new FileStream(filePath, FileMode.Open, FileAccess.Read); + using var reader = new BinaryReader(fileStream); + fileStream.Seek(endHeaderPosition, SeekOrigin.Begin); + + var dataList = new List(); + var jobHandles = new List(); + + for (int pos = endHeaderPosition; pos < fileStream.Length; pos += chunkSize) + { + int currChunkSize = (int)(pos + chunkSize < fileStream.Length ? chunkSize : fileStream.Length - pos); + byte[] chunkBytes = reader.ReadBytes(currChunkSize); + int numElemInChunk = currChunkSize / headerData.BytesPerElement; + + var splatDataNative = new SplatDataNative(numElemInChunk); + dataList.Add(splatDataNative); + + var splatChunkJob = new SplatChunkJob + { + chunkBytes = chunkBytes, + splatDataNative = splatDataNative, + headerData = headerData, + numElem = numElemInChunk + }; + + var handle = splatChunkJob.Schedule(); + jobHandles.Add(handle); + } + + // Copy data here + + splatData.ReleaseGpuResources(); + return splatData; + } + + [BurstCompile] + private struct SplatChunkJob : IJob + { + public byte[] chunkBytes; + public SplatDataNative splatDataNative; + public HeaderData headerData; + public int numElem; + + public void Execute() + { + throw new NotImplementedException(); + } + } + + [BurstCompile] + private struct SplatElementJob : IJobParallelFor + { + public byte[] chunkBytes; + public SplatDataNative splatDataNative; + public HeaderData headerData; + + public void Execute(int ix) + { + throw new NotImplementedException(); + } + } + + // private struct ReadChunkJob + + /*********************************************************************** + ImportData function + - Takes a HeaderData struct and file path string as input + - Following the JOB_SZ and the length of a single element's data + -- Computes how many bytes a single job can process + -- Reads them in chunks + -- As each chunk is read + --- Creates NativeArrays representing all the data points + --- Creates a ReadChunk Jobs to parse the data + - Once each Job is complete, begins copying the NativeArrays into a + newly created SplatData scriptable object + ***********************************************************************/ + + /*********************************************************************** + ReadChunk Job + - This MUST be Burst compiled + - Takes in NativeArrays, the chunk's byte data, and the HeaderData + struct + - It may be more efficient to place the NativeArrays inside a Burst + compiled struct + - Iterates through the byte data using the total length of a data point + as stored in HeaderData + - For each data point, reads the data points and stores them into the + appropriate NativeArray, referencing the byte offset table in + HeaderData to find their locations in the byte data + ***********************************************************************/ + SplatData ImportAsSplatData((string format, int num_element, List<(string, string)> properties) header_info, byte[] binaryData) { var splatData = ScriptableObject.CreateInstance(); @@ -322,7 +479,6 @@ namespace SplatVfx.Editor // color = (Vector4)math.float4(src.r, src.g, src.b, src.a) / 255; // } - #endregion } -} // namespace SplatVfx.Editor +}