Skip to content

Instantly share code, notes, and snippets.

@ChristianSchott

ChristianSchott/RenderPC.cs

Created November 12, 2024 19:06
Show Gist options
  • Save ChristianSchott/afc8eb36e39e39b982aa6c70f3422582 to your computer and use it in GitHub Desktop.
Save ChristianSchott/afc8eb36e39e39b982aa6c70f3422582 to your computer and use it in GitHub Desktop.
Download ZIP
Unity Point Cloud - Compute Rasterizer
Raw
RenderPC.cs
// based on https://github.com/m-schuetz/compute_rasterizer
using UnityEngine;
public struct Particle
{
public Vector4 position;
public Vector4 color;
}
// PUT THIS ON THE CAMERA! Otherwise OnRenderImage will not be called..
public class RenderPC : MonoBehaviour
{
public ComputeShader splatShader;
public Vector2Int size = new Vector2Int(2560, 1440);
public float epsilonColor = .01f;
public float epsilonExpand = .05f;
public int expandIters = 1;
public Transform target;
public Mesh mesh;
private GraphicsBuffer points;
private GraphicsBuffer depth;
private GraphicsBuffer color;
private GraphicsBuffer depth2;
private GraphicsBuffer color2;
void Awake()
{
var particleCount = mesh.vertexCount;
var pointData = new Particle[particleCount];
var verts = mesh.vertices;
var colors = mesh.colors;
for (int i = 0; i < particleCount; i++) {
pointData[i] = new Particle() {
position = verts[i],
color = colors[i],
};
}
points = new GraphicsBuffer(GraphicsBuffer.Target.Structured, particleCount, sizeof(float) * 8);
points.SetData(pointData);
depth = new GraphicsBuffer(GraphicsBuffer.Target.Structured, size.x * size.y, sizeof(uint));
color = new GraphicsBuffer(GraphicsBuffer.Target.Structured, size.x * size.y, sizeof(uint) * 4);
depth2 = new GraphicsBuffer(GraphicsBuffer.Target.Structured, size.x * size.y, sizeof(uint));
color2 = new GraphicsBuffer(GraphicsBuffer.Target.Structured, size.x * size.y, sizeof(uint) * 4);
}
void OnRenderImage(RenderTexture source, RenderTexture destination)
{
splatShader.SetInts("size", size.x, size.y);
splatShader.SetFloat("epsilonColor", epsilonColor);
splatShader.SetFloat("epsilonExpand", epsilonExpand);
var mainCam = Camera.main;
splatShader.SetMatrix("CameraVP", mainCam.projectionMatrix * mainCam.worldToCameraMatrix * target.localToWorldMatrix);
{
var kernel = splatShader.FindKernel("clear");
splatShader.SetTextureFromGlobal(kernel, "CameraDepth", "_CameraDepthTexture");
splatShader.SetBuffer(kernel, "DepthBuffer", depth);
splatShader.SetBuffer(kernel, "ColorBuffer", color);
splatShader.Dispatch(kernel, Mathf.CeilToInt(size.x / (float)32), Mathf.CeilToInt(size.y / (float)32), 1);
}
{
var kernel = splatShader.FindKernel("splat_depth");
splatShader.SetBuffer(kernel, "points", points);
splatShader.SetBuffer(kernel, "DepthBuffer", depth);
splatShader.Dispatch(kernel, Mathf.CeilToInt(points.count / (float)1024), 1, 1);
}
{
var kernel = splatShader.FindKernel("splat_color");
splatShader.SetBuffer(kernel, "points", points);
splatShader.SetBuffer(kernel, "DepthBuffer", depth);
splatShader.SetBuffer(kernel, "ColorBuffer", color);
splatShader.Dispatch(kernel, Mathf.CeilToInt(points.count / (float)1024), 1, 1);
}
for (int i = 0; i < expandIters; i++) {
var kernel = splatShader.FindKernel("expand");
splatShader.SetBuffer(kernel, "DepthBuffer", depth);
splatShader.SetBuffer(kernel, "ColorBuffer", color);
splatShader.SetBuffer(kernel, "DepthBuffer2", depth2);
splatShader.SetBuffer(kernel, "ColorBuffer2", color2);
splatShader.Dispatch(kernel, Mathf.CeilToInt(size.x / (float)32), Mathf.CeilToInt(size.y / (float)32), 1);
(depth, depth2) = (depth2, depth);
(color, color2) = (color2, color);
}
{
var desc = source.descriptor;
desc.enableRandomWrite = true;
desc.msaaSamples = 1;
var tmp = RenderTexture.GetTemporary(desc);
var kernel = splatShader.FindKernel("blit");
splatShader.SetBuffer(kernel, "ColorBuffer", color);
splatShader.SetTexture(kernel, "source", source);
splatShader.SetTexture(kernel, "res", tmp);
splatShader.SetTextureFromGlobal(kernel, "CameraDepth", "_CameraDepthTexture");
splatShader.Dispatch(kernel, Mathf.CeilToInt(size.x / (float)32), Mathf.CeilToInt(size.y / (float)32), 1);
Graphics.Blit(tmp, destination);
RenderTexture.ReleaseTemporary(tmp);
}
}
void OnDestroy()
{
points.Release();
depth.Release();
color.Release();
depth2.Release();
color2.Release();
}
}
Raw
splatPC.compute
// based on https://github.com/m-schuetz/compute_rasterizer
#pragma kernel clear
#pragma kernel splat_depth
#pragma kernel splat_color
#pragma kernel expand
#pragma kernel blit
// #pragma use_dxc
// #pragma require Int64BufferAtomics
#include "UnityCG.cginc"
struct Particle {
float4 position;
float4 color;
};
RWStructuredBuffer<Particle> points;
uint2 size;
float epsilonColor;
float epsilonExpand;
Texture2D<float> CameraDepth;
RWStructuredBuffer<uint> DepthBuffer;
RWStructuredBuffer<uint4> ColorBuffer;
float4x4 CameraVP;
uint index(uint2 coords) {
return coords.x + coords.y * size.x;
}
[numthreads(32, 32, 1)]
void clear (uint3 id : SV_DispatchThreadID)
{
uint pixelIdx = index(id.xy);
float depth = CameraDepth[id.xy];
DepthBuffer[pixelIdx] = asuint(LinearEyeDepth(depth));
ColorBuffer[pixelIdx] = 0;
}
[numthreads(1024, 1, 1)]
void splat_depth (uint3 id : SV_DispatchThreadID)
{
Particle particle = points[id.x];
float4 position = mul(CameraVP, float4(particle.position.xyz, 1));
float3 pos = position.xyz / position.w;
if (pos.z <= -1 || pos.z >= 1
|| pos.x <= -1 || pos.x >= 1
|| pos.y <= -1 || pos.y >= 1) {
return;
}
uint pixelIdx = index((pos.xy * 0.5 + 0.5) * size);
uint oldDepth = DepthBuffer[pixelIdx];
uint depth = asuint(position.w);
if (depth < oldDepth) {
InterlockedMin(DepthBuffer[pixelIdx], depth);
}
}
[numthreads(1024, 1, 1)]
void splat_color (uint3 id : SV_DispatchThreadID)
{
Particle particle = points[id.x];
float4 position = mul(CameraVP, float4(particle.position.xyz, 1));
float3 pos = position.xyz / position.w;
if (pos.z <= -1 || pos.z >= 1
|| pos.x <= -1 || pos.x >= 1
|| pos.y <= -1 || pos.y >= 1) {
return;
}
uint pixelIdx = index((pos.xy * 0.5 + 0.5) * size);
float depth = position.w;//asuint(position.w);
float oldDepth = asfloat(DepthBuffer[pixelIdx]);
if (depth <= oldDepth + epsilonColor) {
uint3 color = particle.color.rgb * 256;
InterlockedAdd(ColorBuffer[pixelIdx].r, color.r);
InterlockedAdd(ColorBuffer[pixelIdx].g, color.g);
InterlockedAdd(ColorBuffer[pixelIdx].b, color.b);
InterlockedAdd(ColorBuffer[pixelIdx].a, 1);
}
}
RWStructuredBuffer<uint> DepthBuffer2;
RWStructuredBuffer<uint4> ColorBuffer2;
[numthreads(32, 32, 1)]
void expand (uint3 id : SV_DispatchThreadID)
{
uint idx = index(id.xy);
uint idl = index(id.xy + uint2(-1, 0));
uint idr = index(id.xy + uint2(+1, 0));
uint idu = index(id.xy + uint2(0, -1));
uint idd = index(id.xy + uint2(0, +1));
float depth = asfloat(DepthBuffer[index(id.xy)]);
float4 n = float4(asfloat(DepthBuffer[idl]), asfloat(DepthBuffer[idr]), asfloat(DepthBuffer[idu]), asfloat(DepthBuffer[idd]));
float minDepth = min(depth, min(min(n.x, n.y), min(n.z, n.w)));
float expandDepth = minDepth + epsilonExpand;
bool dIn = depth <= expandDepth;
bool4 nIn = bool4((n.x <= expandDepth) && !dIn,
(n.y <= expandDepth) && !dIn,
(n.z <= expandDepth) && !dIn,
(n.w <= expandDepth) && !dIn);
uint sum = dIn + nIn.x + nIn.y + nIn.z + nIn.w;
DepthBuffer2[idx] = asuint(minDepth);
// DepthBuffer2[idx] = sum == 0 ? depth :
// (((dIn ? depth : 0)
// + (nIn.x ? n.x : 0)
// + (nIn.y ? n.y : 0)
// + (nIn.z ? n.z : 0)
// + (nIn.w ? n.w : 0)) / sum);
uint4 color = ColorBuffer[idx];
ColorBuffer2[idx] = sum == 0 ? color :
((dIn ? color : 0) +
(nIn.x ? ColorBuffer[idl] : 0) +
(nIn.y ? ColorBuffer[idr] : 0) +
(nIn.z ? ColorBuffer[idu] : 0) +
(nIn.w ? ColorBuffer[idd] : 0));
}
Texture2D<float4> source;
RWTexture2D<float4> res;
[numthreads(32, 32, 1)]
void blit (uint3 id : SV_DispatchThreadID)
{
float4 raw = ColorBuffer[index(id.xy)];
// smoothing
// raw += (ColorBuffer[index(id.xy + uint2(1, 0))]
// + ColorBuffer[index(id.xy + uint2(-1, 0))]
// + ColorBuffer[index(id.xy + uint2(0, 1))]
// + ColorBuffer[index(id.xy + uint2(0, -1))]) * 0.5;
// raw += (ColorBuffer[index(id.xy + uint2(1, 1))]
// + ColorBuffer[index(id.xy + uint2(-1, 1))]
// + ColorBuffer[index(id.xy + uint2(1, -1))]
// + ColorBuffer[index(id.xy + uint2(-1, -1))]) * 0.25;
float3 color = raw.w == 0 ? float3(0,0,0) : (raw.rgb * (1.0 / 256.0)) / (float)raw.w;
res[id.xy] = lerp(source[id.xy], float4(color, 1), min(1, raw.w));
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment