539 lines
21 KiB
C#
539 lines
21 KiB
C#
using Unity.Burst;
|
|
using Unity.Mathematics;
|
|
using UnityEngine;
|
|
using UnityEngine.VFX;
|
|
using UnityEngine.VFX.Utility;
|
|
using UnityEditor;
|
|
using UnityEditor.AssetImporters;
|
|
using UnityEditor.Experimental;
|
|
using System;
|
|
using System.IO;
|
|
using System.Runtime.InteropServices;
|
|
using System.Text;
|
|
using UnityEditor.Rendering.LookDev;
|
|
using System.Linq;
|
|
using System.Collections.Generic;
|
|
using Unity.VectorGraphics;
|
|
using Unity.VisualScripting;
|
|
using System.Collections;
|
|
using Unity.Collections;
|
|
using UnityEditor.Experimental.GraphView;
|
|
using UnityEngine.Playables;
|
|
using Unity.Android.Gradle.Manifest;
|
|
using UnityEngine.Timeline;
|
|
using Unity.Jobs;
|
|
using System.Xml.Schema;
|
|
using Unity.Android.Gradle;
|
|
|
|
namespace SplatVfx.Editor
|
|
{
|
|
|
|
[ScriptedImporter(1, "ply"), BurstCompile]
|
|
public sealed class SplatImporter : ScriptedImporter
|
|
{
|
|
const float SH_C0 = 0.28209479177387814f;
|
|
const string DefaultVfxPath = "Assets/vfx-graph-gs/SplatVFX.vfx";
|
|
|
|
public override void OnImportAsset(AssetImportContext context)
|
|
{
|
|
var (headerString, endHeaderPosition) = GetHeader(context.assetPath);
|
|
var headerData = ParseMetadata(headerString);
|
|
|
|
var data = ImportData(headerData, context.assetPath, endHeaderPosition);
|
|
data.name = Path.GetFileNameWithoutExtension(context.assetPath);
|
|
context.AddObjectToAsset("data", data);
|
|
|
|
var prefab = CreatePrefab(data);
|
|
context.AddObjectToAsset("prefab", prefab);
|
|
context.SetMainObject(prefab);
|
|
}
|
|
|
|
private GameObject CreatePrefab(SplatData data)
|
|
{
|
|
var go = new GameObject();
|
|
var vfx = go.AddComponent<VisualEffect>();
|
|
vfx.visualEffectAsset = EditorResources.Load<VisualEffectAsset>(DefaultVfxPath);
|
|
|
|
var binderBase = go.AddComponent<VFXPropertyBinder>();
|
|
var binder = binderBase.AddPropertyBinder<VFXSplatDataBinder>();
|
|
|
|
binder.SplatData = data;
|
|
|
|
return go;
|
|
}
|
|
|
|
const int MAX_CHUNK_SIZE = 1048576 * 256; // 256MiB
|
|
|
|
private enum TypeEnum
|
|
{
|
|
None, Float32, Float64, Float16, Int8, UInt8, Char, Short, UInt16
|
|
}
|
|
|
|
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<TypeEnum, int> 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<string, TypeEnum> 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 NativeArray<int> ElementByteOffsets;
|
|
public NativeArray<int> ElementTypes;
|
|
public int BytesPerElement;
|
|
|
|
public HeaderData(int count)
|
|
{
|
|
NumElements = 0;
|
|
BinaryDataStartOffset = 0;
|
|
ElementByteOffsets = new NativeArray<int>(count, Allocator.Persistent);
|
|
ElementTypes = new NativeArray<int>(count, Allocator.Persistent);
|
|
BytesPerElement = 0;
|
|
}
|
|
|
|
public void Dispose()
|
|
{
|
|
ElementByteOffsets.Dispose();
|
|
ElementTypes.Dispose();
|
|
}
|
|
}
|
|
|
|
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);
|
|
|
|
var headerBuilder = new StringBuilder();
|
|
string line;
|
|
|
|
while ((line = reader.ReadLine()) != null)
|
|
{
|
|
headerBuilder.AppendLine(line);
|
|
if (line.Trim() == "end_header")
|
|
{
|
|
break;
|
|
}
|
|
}
|
|
|
|
// Old logic that works because AppendLine adds another \n
|
|
// TODO: Fix this logic
|
|
int endHeaderPosition = headerBuilder.Length - headerBuilder.ToString().Split('\n').Length + 1;
|
|
Debug.Log($"Total filesize: {fileStream.Length}");
|
|
Debug.Log($"Byte stream start position: {endHeaderPosition}");
|
|
|
|
return (headerBuilder.ToString(), endHeaderPosition);
|
|
}
|
|
|
|
private static HeaderData ParseMetadata(string headerString)
|
|
{
|
|
HeaderData result = new(59);
|
|
string line;
|
|
|
|
using var reader = new StringReader(headerString);
|
|
|
|
while ((line = reader.ReadLine()) != null)
|
|
{
|
|
if (line.Trim() == "end_header")
|
|
break;
|
|
|
|
var strings = line.Split(" ");
|
|
|
|
switch (strings.FirstOrDefault())
|
|
{
|
|
case "element":
|
|
result.NumElements = uint.Parse(strings[2]);
|
|
break;
|
|
|
|
case "property":
|
|
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] = (int)dataType;
|
|
}
|
|
|
|
result.BytesPerElement += TypeToByteLengths[dataType];
|
|
break;
|
|
}
|
|
}
|
|
Debug.Log($"Bytes per Element: {result.BytesPerElement}");
|
|
|
|
return result;
|
|
}
|
|
|
|
[BurstCompile]
|
|
private struct SplatDataNative
|
|
{
|
|
public NativeArray<float3> position;
|
|
public NativeArray<float3> color;
|
|
public NativeArray<float> sh;
|
|
public NativeArray<float> opacity;
|
|
public NativeArray<float3> scale;
|
|
public NativeArray<float4> rotation;
|
|
|
|
public SplatDataNative(int count)
|
|
{
|
|
position = new NativeArray<float3>(count, Allocator.Persistent);
|
|
color = new NativeArray<float3>(count, Allocator.Persistent);
|
|
sh = new NativeArray<float>(count * 45, Allocator.Persistent);
|
|
opacity = new NativeArray<float>(count, Allocator.Persistent);
|
|
scale = new NativeArray<float3>(count, Allocator.Persistent);
|
|
rotation = new NativeArray<float4>(count, Allocator.Persistent);
|
|
}
|
|
|
|
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<SplatData>();
|
|
|
|
splatData.PositionArray = new Vector3[headerData.NumElements];
|
|
splatData.ScaleArray = new Vector3[headerData.NumElements];
|
|
splatData.RotationArray = new Quaternion[headerData.NumElements];
|
|
splatData.ColorArray = new Color[headerData.NumElements];
|
|
splatData.OpacityArray = new float[headerData.NumElements];
|
|
splatData.SHArray = new float[headerData.NumElements * 45];
|
|
|
|
int elementsPerChunk = Mathf.Min(
|
|
MAX_CHUNK_SIZE / headerData.BytesPerElement,
|
|
(int)headerData.NumElements
|
|
);
|
|
int chunkSize = elementsPerChunk * headerData.BytesPerElement;
|
|
|
|
Debug.Log($"Elements per chunk: {elementsPerChunk}");
|
|
Debug.Log($"Chunk size: {chunkSize}");
|
|
|
|
using var fileStream = new FileStream(filePath, FileMode.Open, FileAccess.Read);
|
|
fileStream.Seek(endHeaderPosition, SeekOrigin.Begin);
|
|
|
|
int totalBytes = (int)fileStream.Length - endHeaderPosition;
|
|
Debug.Log($"Total bytes: {totalBytes}");
|
|
|
|
var splatDataNative = new SplatDataNative(elementsPerChunk);
|
|
var chunkBuffer = new NativeArray<byte>(chunkSize, Allocator.Persistent);
|
|
var chunkBufferSpan = MemoryMarshal.AsBytes<byte>(chunkBuffer);
|
|
int numBytesRead;
|
|
int numElemsRead;
|
|
int totalElemsRead = 0;
|
|
|
|
while ((numBytesRead = fileStream.Read(chunkBufferSpan)) > 0)
|
|
{
|
|
numElemsRead = numBytesRead / headerData.BytesPerElement;
|
|
var splatElementJob = new SplatElementJob
|
|
{
|
|
chunkBytes = chunkBuffer,
|
|
splatDataNative = splatDataNative,
|
|
headerData = headerData,
|
|
totalElemsRead = totalElemsRead,
|
|
numElemsRead = numElemsRead
|
|
};
|
|
|
|
splatElementJob.Schedule().Complete();
|
|
|
|
for (int i = 0; i < numElemsRead; i++)
|
|
{
|
|
splatData.PositionArray[totalElemsRead + i] = splatDataNative.position[i];
|
|
splatData.ColorArray[totalElemsRead + i].r = splatDataNative.color[i].x;
|
|
splatData.ColorArray[totalElemsRead + i].g = splatDataNative.color[i].y;
|
|
splatData.ColorArray[totalElemsRead + i].b = splatDataNative.color[i].z;
|
|
splatData.OpacityArray[totalElemsRead + i] = splatDataNative.opacity[i];
|
|
splatData.ScaleArray[totalElemsRead + i] = splatDataNative.scale[i];
|
|
splatData.RotationArray[totalElemsRead + i] = new quaternion(splatDataNative.rotation[i]);
|
|
}
|
|
|
|
totalElemsRead += numElemsRead;
|
|
Debug.Log($"Current Elements Read: {totalElemsRead}");
|
|
}
|
|
|
|
Debug.Log($"First position Vector3: {splatData.PositionArray[0]}");
|
|
|
|
headerData.Dispose();
|
|
splatDataNative.Dispose();
|
|
chunkBuffer.Dispose();
|
|
|
|
splatData.ReleaseGpuResources();
|
|
|
|
return splatData;
|
|
}
|
|
|
|
[BurstCompile]
|
|
private struct SplatElementJob : IJob
|
|
{
|
|
public NativeArray<byte> chunkBytes;
|
|
public SplatDataNative splatDataNative;
|
|
[ReadOnly] public HeaderData headerData;
|
|
[ReadOnly] public int totalElemsRead;
|
|
[ReadOnly] public int numElemsRead;
|
|
|
|
public void Execute()
|
|
{
|
|
for (int i = 0; i < numElemsRead; i++)
|
|
{
|
|
int elemIx = i;
|
|
int start = elemIx * headerData.BytesPerElement;
|
|
|
|
splatDataNative.position[elemIx] = new()
|
|
{
|
|
x = chunkBytes.ReinterpretLoad<float>(start + headerData.ElementByteOffsets[(int)DataEnum.x]),
|
|
y = chunkBytes.ReinterpretLoad<float>(start + headerData.ElementByteOffsets[(int)DataEnum.y]),
|
|
z = chunkBytes.ReinterpretLoad<float>(start + headerData.ElementByteOffsets[(int)DataEnum.z])
|
|
};
|
|
|
|
float3 color = new()
|
|
{
|
|
x = chunkBytes.ReinterpretLoad<float>(start + headerData.ElementByteOffsets[(int)DataEnum.f_dc_0]),
|
|
y = chunkBytes.ReinterpretLoad<float>(start + headerData.ElementByteOffsets[(int)DataEnum.f_dc_1]),
|
|
z = chunkBytes.ReinterpretLoad<float>(start + headerData.ElementByteOffsets[(int)DataEnum.f_dc_2])
|
|
};
|
|
color = 0.5f + color * SH_C0;
|
|
splatDataNative.color[elemIx] = color;
|
|
|
|
for (int j = 0; j < 45; j++)
|
|
{
|
|
splatDataNative.sh[elemIx + j] = chunkBytes[start + headerData.ElementByteOffsets[(int)DataEnum.f_rest_0 + j]];
|
|
}
|
|
|
|
float opacity = chunkBytes.ReinterpretLoad<float>(start + headerData.ElementByteOffsets[(int)DataEnum.opacity]);
|
|
opacity = 1 / (1 + math.exp(-opacity));
|
|
splatDataNative.opacity[elemIx] = opacity;
|
|
|
|
float3 scale = new()
|
|
{
|
|
x = chunkBytes.ReinterpretLoad<float>(start + headerData.ElementByteOffsets[(int)DataEnum.scale_0]),
|
|
y = chunkBytes.ReinterpretLoad<float>(start + headerData.ElementByteOffsets[(int)DataEnum.scale_1]),
|
|
z = chunkBytes.ReinterpretLoad<float>(start + headerData.ElementByteOffsets[(int)DataEnum.scale_2])
|
|
};
|
|
scale = math.exp(scale);
|
|
splatDataNative.scale[elemIx] = scale;
|
|
|
|
splatDataNative.rotation[elemIx] = new()
|
|
{
|
|
w = chunkBytes.ReinterpretLoad<float>(start + headerData.ElementByteOffsets[(int)DataEnum.rot_0]),
|
|
x = chunkBytes.ReinterpretLoad<float>(start + headerData.ElementByteOffsets[(int)DataEnum.rot_1]),
|
|
y = chunkBytes.ReinterpretLoad<float>(start + headerData.ElementByteOffsets[(int)DataEnum.rot_2]),
|
|
z = chunkBytes.ReinterpretLoad<float>(start + headerData.ElementByteOffsets[(int)DataEnum.rot_3])
|
|
};
|
|
}
|
|
}
|
|
}
|
|
|
|
SplatData ImportAsSplatData((string format, int num_element, List<(string, string)> properties) header_info, byte[] binaryData)
|
|
{
|
|
var splatData = ScriptableObject.CreateInstance<SplatData>();
|
|
|
|
var stream = new MemoryStream(binaryData);
|
|
var reader = new BinaryReader(stream);
|
|
|
|
splatData.PositionArray = new Vector3[header_info.num_element];
|
|
splatData.ScaleArray = new Vector3[header_info.num_element];
|
|
splatData.RotationArray = new Quaternion[header_info.num_element];
|
|
splatData.ColorArray = new Color[header_info.num_element];
|
|
splatData.OpacityArray = new float[header_info.num_element];
|
|
splatData.SHArray = new float[header_info.num_element * 45];
|
|
|
|
for (int i = 0; i < header_info.num_element; i++)
|
|
{
|
|
foreach ((string name, string dtype) property in header_info.properties)
|
|
{
|
|
float val = reader.ReadSingle();
|
|
|
|
switch (property.name)
|
|
{
|
|
case "x":
|
|
splatData.PositionArray[i].x = val;
|
|
|
|
break;
|
|
|
|
case "y":
|
|
splatData.PositionArray[i].y = val;
|
|
|
|
break;
|
|
|
|
case "z":
|
|
splatData.PositionArray[i].z = val;
|
|
|
|
break;
|
|
|
|
case "opacity":
|
|
splatData.OpacityArray[i] = 1 / (1 + math.exp(-val));
|
|
|
|
break;
|
|
|
|
default:
|
|
if (property.name.StartsWith("f_dc"))
|
|
{
|
|
switch (property.name[5])
|
|
{
|
|
case '0':
|
|
splatData.ColorArray[i].r = 0.5f + val * SH_C0;
|
|
|
|
break;
|
|
|
|
case '1':
|
|
splatData.ColorArray[i].g = 0.5f + val * SH_C0;
|
|
|
|
break;
|
|
|
|
case '2':
|
|
splatData.ColorArray[i].b = 0.5f + val * SH_C0;
|
|
|
|
break;
|
|
}
|
|
|
|
}
|
|
else if (property.name.StartsWith("scale"))
|
|
{
|
|
switch (property.name[6])
|
|
{
|
|
case '0':
|
|
splatData.ScaleArray[i].x = math.exp(val);
|
|
|
|
break;
|
|
|
|
case '1':
|
|
splatData.ScaleArray[i].y = math.exp(val);
|
|
|
|
break;
|
|
|
|
case '2':
|
|
splatData.ScaleArray[i].z = math.exp(val);
|
|
|
|
break;
|
|
}
|
|
|
|
}
|
|
else if (property.name.StartsWith("rot"))
|
|
{
|
|
switch (property.name[4])
|
|
{
|
|
case '0':
|
|
splatData.RotationArray[i].w = val;
|
|
|
|
break;
|
|
|
|
case '1':
|
|
splatData.RotationArray[i].x = val;
|
|
|
|
break;
|
|
|
|
case '2':
|
|
splatData.RotationArray[i].y = val;
|
|
|
|
break;
|
|
|
|
case '3':
|
|
splatData.RotationArray[i].z = val;
|
|
|
|
break;
|
|
}
|
|
|
|
}
|
|
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
|
|
splatData.ReleaseGpuResources();
|
|
|
|
return splatData;
|
|
}
|
|
|
|
// struct ReadData
|
|
// {
|
|
// public float px, py, pz;
|
|
// public float sx, sy, sz;
|
|
// public byte r, g, b, a;
|
|
// public byte rw, rx, ry, rz;
|
|
// }
|
|
|
|
// (Vector3[] position, Vector3[] axis, Color[] color)
|
|
// LoadDataArrays(string path)
|
|
// {
|
|
// var bytes = (Span<byte>)File.ReadAllBytes(path);
|
|
// var count = bytes.Length / 32;
|
|
|
|
// var source = MemoryMarshal.Cast<byte, ReadData>(bytes);
|
|
|
|
// var position = new Vector3[count];
|
|
// var axis = new Vector3[count * 3];
|
|
// var color = new Color[count];
|
|
|
|
// for (var i = 0; i < count; i++)
|
|
// ParseReadData(source[i],
|
|
// out position[i],
|
|
// out axis[i * 3],
|
|
// out axis[i * 3 + 1],
|
|
// out axis[i * 3 + 2],
|
|
// out color[i]);
|
|
|
|
// return (position, axis, color);
|
|
// }
|
|
|
|
// [BurstCompile]
|
|
// void ParseReadData(in ReadData src,
|
|
// out Vector3 position,
|
|
// out Vector3 axis1,
|
|
// out Vector3 axis2,
|
|
// out Vector3 axis3,
|
|
// out Color color)
|
|
// {
|
|
// var rv = (math.float4(src.rx, src.ry, src.rz, src.rw) - 128) / 128;
|
|
// var q = math.quaternion(-rv.x, -rv.y, rv.z, rv.w);
|
|
// position = math.float3(src.px, src.py, -src.pz);
|
|
// axis1 = math.mul(q, math.float3(src.sx, 0, 0));
|
|
// axis2 = math.mul(q, math.float3(0, src.sy, 0));
|
|
// axis3 = math.mul(q, math.float3(0, 0, src.sz));
|
|
// color = (Vector4)math.float4(src.r, src.g, src.b, src.a) / 255;
|
|
// }
|
|
|
|
}
|
|
|
|
}
|