A Simple GPU-based Drawing App in Unity

This week I’m escaping the complicated world of voxels for a bit to do something lighter. We’re building a simple GPU-based drawing app. What does that mean? A drawing app that runs entirely on the GPU using compute shaders. It’ll be just like Photoshop, except without all those complicated features like layers, undo, or a colour picker. Let’s go over the plan.

What’s the plan?

Our canvas will be a render texture that’s stretched across the screen. We’ll dispatch a compute shader kernel every frame, with a reference to our canvas and the current mouse position, among other things. We’ll dispatch one thread per pixel. So, each thread will compare its pixel position to the mouse position, taking into account the brush size. If it falls within the brush’s position and size, we’ll colour that pixel. So we’ll have a wonderfully parallel drawing app.

Once we complete the first part, we’ll tackle simple brush position interpolation. The issue is the mouse can move a lot from one frame to the next. If we don’t interpolate the position, we’ll end up with several disconnected round brush strokes.

Finally, I’ll demonstrate how we can add some GUI to control variables like the brush size. Let’s start.

Set up

Create a new compute shader. I called mine Draw.compute because I didn’t think too hard about a name. We’ll need two kernels, one to initialize the background and one to update every frame. We also need an RWTexture<float4> to act as our canvas. In case you forgot, that’s a writeable texture where each pixel is a float4.

#pragma kernel Update
#pragma kernel InitBackground

RWTexture2D<float4> _Canvas;

[numthreads(8,8,1)]
void InitBackground(uint3 id : SV_DispatchThreadID)
{
    _Canvas[id.xy] = float4(0, 0, 0, 0);
}

[numthreads(8,8,1)]
void Update(uint3 id : SV_DispatchThreadID)
{
    //TODO
}

With the compute shader bones out of the way, let’s set up the C# script. Create a MonoBehaviour called DrawManager.cs to manage our compute shader. We’ll create a RenderTexture and assign it to canvas. We’ll also dispatch InitBackground on Start and Update on Update. We’ll also use OnRenderImage to display the canvas. The component has to be attached to a camera in the scene. The callback receives the camera’s frame buffer, and we’ll draw our canvas on top of it.

public class DrawManager : MonoBehaviour
{
    [SerializeField] ComputeShader _drawComputeShader;
    RenderTexture _canvasRenderTexture;

    void Start()
    {
        _canvasRenderTexture = new RenderTexture(Screen.width, Screen.height, 24);
        _canvasRenderTexture.filterMode = FilterMode.Point;
        _canvasRenderTexture.enableRandomWrite = true;
        _canvasRenderTexture.Create();

        int initBackgroundKernel = _drawComputeShader.FindKernel("InitBackground");
        _drawComputeShader.SetTexture(initBackgroundKernel, "_Canvas", _canvasRenderTexture);
        _drawComputeShader.Dispatch(initBackgroundKernel, _canvasRenderTexture.width / 8,
            _canvasRenderTexture.height / 8, 1);
    }

    void Update()
    {
            int updateKernel = _drawComputeShader.FindKernel("Update");
            _drawComputeShader.SetTexture(updateKernel, "_Canvas", _canvasRenderTexture);
            _drawComputeShader.Dispatch(updateKernel, _canvasRenderTexture.width / 8,
                _canvasRenderTexture.height / 8, 1);
    }

    void OnRenderImage(RenderTexture src, RenderTexture dest)
    {
        Graphics.Blit(_canvasRenderTexture, dest);
    }
}

With the boilerplate out of the way, we can start working.

Drawing

Let’s jump to the fun part and implement drawing right away. All we need to make this interactive is the current mouse position in the Update kernel. Add float4 _MousePosition and bool _MouseDown to the compute shader. We’ll use _MouseDown to prevent drawing when the mouse button is up. The Update kernel runs once per pixel, so the function must check if the current thread is close enough to the mouse position to be coloured. On that note, let’s add float _BrushSize to the compute shader as well.

float4 _MousePosition;
bool _MouseDown;
float _BrushSize;

[numthreads(8,8,1)]
void Update(uint3 id : SV_DispatchThreadID)
{
    if (!_MouseDown) return;

    float2 pixelPos = id.xy;
    float2 mousePos = _MousePosition.xy;
    if (length(pixelPos - mousePos) < _BrushSize)
        _Canvas[id.xy] = float4(1, 0, 0, 1);
}

This block of code is everything newly added to the compute shader. As promised, in Update, we check the distance from the thread’s pixel to the mouse position and draw if necessary. Now you can draw pretty pictures like this.

Right now, if you draw slightly fast in the app, you’ll notice an issue. We have a series of dots rather than a single continuous stroke.

To fix that, we need the previous mouse position as well. Then, we’ll interpolate between the two points, drawing several points along the way. Let’s extract our brush into a function at the same time. I call the function HardBrush. In the future, we could follow the same template to create a soft brush, textured brush, patterned brush, etc. So in the previous version, we would check our position against a single mouse position. Now, we’re going to slowly sweep from the last position mouse to the new mouse position, checking our position against each point in the sweep.

float4 HardBrush(float2 pixelPos, float4 currentColor, float4 brushColor, 
                 float brushSize, float2 previousMousePosition,
                 float2 mousePosition, float strokeSmoothingInterval)
{
    for (float i = 0; i < 1.0; i += strokeSmoothingInterval)
    {
        const float2 mousePos = lerp(previousMousePosition, mousePosition, i);
        if (length(pixelPos - mousePos) < brushSize)
            return brushColor;
    }

    return currentColor;
}

So the HardBrush function is the same as it was, but now with an added for loop. Let’s hook it up to the Update function.

float4 _PreviousMousePosition;
float _StrokeSmoothingInterval;
float _BrushSize;
float4 _BrushColour;

[numthreads(8,8,1)]
void Update(uint3 id : SV_DispatchThreadID)
{
    if (!_MouseDown) return;
    
    _Canvas[id.xy] = HardBrush(id.xy, _Canvas[id.xy], _BrushColour, 
                               _BrushSize, _PreviousMousePosition,
                               _MousePosition, _StrokeSmoothingInterval);
}

You’ll notice there are new variables as well; they’re from the C# side. Later, I’ll share the entire C# side, but all of these variables are directly from the inspector.

To finish, we’ll add a slider to control the brush size to demonstrate how to add GUI to our drawing app.

Tools GUI

First, add and position a slider in your scene. I’ll leave that part up to you. Next, create a new script to add to the slider called BrushSizeSlider.cs. We’ll use this script to keep track of whether the slider is currently in use. That way, we can stop drawing when manipulating the slider. Here’s the script.

using UnityEngine;
using UnityEngine.EventSystems;
using UnityEngine.UI;

public class BrushSizeSlider : MonoBehaviour, IPointerUpHandler, IPointerDownHandler
{
    public bool isInUse;
    public Slider slider;
    
    public void OnPointerUp(PointerEventData eventData)
    {
        isInUse = false;
    }

    public void OnPointerDown(PointerEventData eventData)
    {
        isInUse = true;
    }
}

So we set a flag when the slider is in use. Add a reference to the BrushSizeSlider in DrawManager.cs and check it on Update. If it’s in use, don’t dispatch the Update kernel. Also, set the min and max values on the slider to something reasonable, like 1 to 20. Finally, set the _BrushSize in the compute shader to the slider’s value. Now we have a simple GUI to control the brush size.

Before wrapping up, here’s the entire DrawManager.cs file for reference.

using UnityEngine;

public class DrawManager : MonoBehaviour
{
    [SerializeField] ComputeShader _drawComputeShader;
    [SerializeField] Color _backgroundColour;
    [SerializeField] Color _brushColour;
    [SerializeField] float _brushSize = 10f;

    [SerializeField] BrushSizeSlider _brushSizeSlider;
    [SerializeField, Range(0.01f, 1)] float _strokeSmoothingInterval = 0.1f;
    RenderTexture _canvasRenderTexture;

    Vector4 _previousMousePosition;

    void Start()
    {
        _brushSizeSlider.slider.SetValueWithoutNotify(_brushSize);

        _canvasRenderTexture = new RenderTexture(Screen.width, Screen.height, 24);
        _canvasRenderTexture.filterMode = FilterMode.Point;
        _canvasRenderTexture.enableRandomWrite = true;
        _canvasRenderTexture.Create();

        int initBackgroundKernel = _drawComputeShader.FindKernel("InitBackground");
        _drawComputeShader.SetVector("_BackgroundColour", _backgroundColour);
        _drawComputeShader.SetTexture(initBackgroundKernel, "_Canvas", _canvasRenderTexture);
        _drawComputeShader.SetFloat("_CanvasWidth", _canvasRenderTexture.width);
        _drawComputeShader.SetFloat("_CanvasHeight", _canvasRenderTexture.height);
        _drawComputeShader.GetKernelThreadGroupSizes(initBackgroundKernel,
            out uint xGroupSize, out uint yGroupSize, out _);
        _drawComputeShader.Dispatch(initBackgroundKernel,
            Mathf.CeilToInt(_canvasRenderTexture.width / (float) xGroupSize),
            Mathf.CeilToInt(_canvasRenderTexture.height / (float) yGroupSize),
            1);
        _drawComputeShader.Dispatch(initBackgroundKernel,
            Mathf.CeilToInt(_canvasRenderTexture.width / (float) xGroupSize),
            Mathf.CeilToInt(_canvasRenderTexture.height / (float) yGroupSize),
            1);


        _previousMousePosition = Input.mousePosition;
    }

    void Update()
    {
        if (!_brushSizeSlider.isInUse && Input.GetMouseButton(0))
        {
            int updateKernel = _drawComputeShader.FindKernel("Update");
            _drawComputeShader.SetVector("_PreviousMousePosition", _previousMousePosition);
            _drawComputeShader.SetVector("_MousePosition", Input.mousePosition);
            _drawComputeShader.SetBool("_MouseDown", Input.GetMouseButton(0));
            _drawComputeShader.SetFloat("_BrushSize", _brushSize);
            _drawComputeShader.SetVector("_BrushColour", _brushColour);
            _drawComputeShader.SetFloat("_StrokeSmoothingInterval", _strokeSmoothingInterval);
            _drawComputeShader.SetTexture(updateKernel, "_Canvas", _canvasRenderTexture);
            _drawComputeShader.SetFloat("_CanvasWidth", _canvasRenderTexture.width);
            _drawComputeShader.SetFloat("_CanvasHeight", _canvasRenderTexture.height);

            _drawComputeShader.GetKernelThreadGroupSizes(updateKernel,
                out uint xGroupSize, out uint yGroupSize, out _);
            _drawComputeShader.Dispatch(updateKernel,
                Mathf.CeilToInt(_canvasRenderTexture.width / (float) xGroupSize),
                Mathf.CeilToInt(_canvasRenderTexture.height / (float) yGroupSize),
                1);
        }

        _previousMousePosition = Input.mousePosition;
    }

    void OnRenderImage(RenderTexture src, RenderTexture dest)
    {
        Graphics.Blit(_canvasRenderTexture, dest);
    }

    public void OnBrushSizeChanged(float newValue)
    {
        _brushSize = newValue;
    }
}

So there’s our simple drawing app. I hooked up some variables behind the scenes. I also added some bounds checks to make sure we don’t try to draw outside the canvas. I also added an early exit on the compute shader side, so I’ll share that entire file for reference. By the way, the Github project is linked at the end, it’s easier to read the code there if you can.

#pragma kernel Update
#pragma kernel InitBackground

RWTexture2D<float4> _Canvas;
float _CanvasWidth;
float _CanvasHeight;
float4 _PreviousMousePosition;
float4 _MousePosition;
float _StrokeSmoothingInterval;
bool _MouseDown;
float _BrushSize;
float4 _BrushColour;
float4 _BackgroundColour;

float4 HardBrush(float2 pixelPos, float4 currentColor, float4 brushColor, float brushSize, float2 previousMousePosition,
                 float2 mousePosition, float strokeSmoothingInterval)
{
    for (float i = 0; i < 1.0; i += strokeSmoothingInterval)
    {
        const float2 mousePos = lerp(previousMousePosition, mousePosition, i);
        if (length(pixelPos - mousePos) < brushSize)
            return brushColor;
    }

    return currentColor;
}

[numthreads(8,8,1)]
void InitBackground(uint3 id : SV_DispatchThreadID)
{
    if (id.x >= _CanvasWidth || id.y >= _CanvasHeight)
        return;

    _Canvas[id.xy] = _BackgroundColour;
}

[numthreads(8,8,1)]
void Update(uint3 id : SV_DispatchThreadID)
{
    if (!_MouseDown)
        return;

    if (id.x >= _CanvasWidth || id.y >= _CanvasHeight)
        return;

    _Canvas[id.xy] = HardBrush(id.xy, _Canvas[id.xy], _BrushColour, 
	                       _BrushSize, _PreviousMousePosition, 
                               _MousePosition, _StrokeSmoothingInterval);
}

The entire compute shader is under 40 lines of code, which is pretty cool if you consider that it’s the bulk of the functionality. The whole project is around 100 lines of code, which is also impressively small, I think.

See the finished project here on GitHub. If you’re enjoying my work, join my mailing list to be notified whenever a new post is released.

Leave A Comment