Final Script

extends Node

var red_score := 0
var blue_score := 0

var thread: Thread
var semaphore: Semaphore
var mutex: Mutex
var exit := false

var paint_texture: ViewportTexture

var rd: RenderingDevice
var shader
var pipeline
var buffer

var buffer_uniform
var v_tex
var samp
var tex_uniform

var uniform_set

func _ready():
    paint_texture = $'../DrawViewport'.get_texture()

    mutex = Mutex.new()
    semaphore = Semaphore.new()
    thread = Thread.new()
    thread.start(_thread_calculate_score)

    # We will be using our own RenderingDevice to handle the compute commands
    rd = RenderingServer.create_local_rendering_device()

    # Create shader and pipeline
    var shader_file = load("res://compute.glsl")
    var shader_spirv = shader_file.get_spirv()
    shader = rd.shader_create_from_spirv(shader_spirv)
    pipeline = rd.compute_pipeline_create(shader)

    var pb = PackedInt32Array([0,0])
    var pbb = pb.to_byte_array()

    buffer = rd.storage_buffer_create(pbb.size(), pbb)

    buffer_uniform = RDUniform.new()
    buffer_uniform.uniform_type = RenderingDevice.UNIFORM_TYPE_STORAGE_BUFFER
    buffer_uniform.binding = 0
    buffer_uniform.add_id(buffer)

    var img = paint_texture.get_image()
    var ta = img.get_data()

    var fmt = RDTextureFormat.new()
    fmt.width = paint_texture.get_width()
    fmt.height = paint_texture.get_height()
    fmt.usage_bits = RenderingDevice.TEXTURE_USAGE_CAN_UPDATE_BIT | RenderingDevice.TEXTURE_USAGE_SAMPLING_BIT
    fmt.format = RenderingDevice.DATA_FORMAT_R8G8B8A8_SRGB

    v_tex = rd.texture_create(fmt, RDTextureView.new(), [ta])
    var samp_state = RDSamplerState.new()
    samp_state.unnormalized_uvw = true
    samp = rd.sampler_create(samp_state)

    tex_uniform = RDUniform.new()
    tex_uniform.uniform_type = RenderingDevice.UNIFORM_TYPE_SAMPLER_WITH_TEXTURE
    tex_uniform.binding = 1
    tex_uniform.add_id(samp)
    tex_uniform.add_id(v_tex)

    uniform_set = rd.uniform_set_create([buffer_uniform, tex_uniform], shader, 0)

func recalculate_score():
    await RenderingServer.frame_post_draw
    semaphore.post()

func get_blue_score() -> int:
    var result := blue_score

    return result

func get_red_score() -> int:
    var result := red_score

    return result

func _thread_calculate_score():
    while true:
        semaphore.wait()

        mutex.lock()
        var should_exit := exit
        mutex.unlock()

        if should_exit:
            break

        mutex.lock()
        var image := paint_texture.get_image()
        mutex.unlock()

        var size = image.get_size()
        var blue_pixels := 0
        var red_pixels := 0
        var ta = image.get_data()

        rd.texture_update(v_tex, 0, ta)

        # Start compute list to start recording our compute commands
        var compute_list = rd.compute_list_begin()
        # Bind the pipeline, this tells the GPU what shader to use
        rd.compute_list_bind_compute_pipeline(compute_list, pipeline)
        # Binds the uniform set with the data we want to give our shader
        rd.compute_list_bind_uniform_set(compute_list, uniform_set, 0)
        # Dispatch 1x1x1 (XxYxZ) work groups
        rd.compute_list_dispatch(compute_list, 2048/8, 2048/8, 1)
        #rd.compute_list_add_barrier(compute_list)
        # Tell the GPU we are done with this compute task
        rd.compute_list_end()

        # Force the GPU to start our commands
        rd.submit()
        # Force the CPU to wait for the GPU to finish with the recorded commands
        rd.sync()

        mutex.lock()
        # Now we can grab our data from the storage buffer
        var byte_data = rd.buffer_get_data(buffer)
        var output := byte_data.to_int32_array()

        print(output)
        blue_score = floor(output[1]/1000)
        red_score = floor(output[0]/1000)
        var pb = PackedInt32Array([0,0])
        var pbb = pb.to_byte_array()
        rd.buffer_update(buffer, 0, pbb.size(), pbb)
        mutex.unlock()

func _exit_tree():
    mutex.lock()
    exit = true
    mutex.unlock()

    semaphore.post()
    thread.wait_to_finish()