ScreenQuadFragmentShader.metal

struct FragmentOut {
    float4 color [[color(0)]];
    float depth [[depth(any)]];
};

fragment FragmentOut screenQuadFragment(ScreenQuadVertex inVertex [[stage_in]],
                                        texture2d<uint, access::sample> segmentationTexture [[texture(0)]],
                                        constant Uniforms& uniforms [[buffer(0)]])
{
    FragmentOut out;
    out.color = float4(0);
    
    float2 uv = inVertex.texcoord;
    // Scales the texture coordinates to the region of interest.
    // regionOfInterest is a CGRect (x, y, width, height) stored in a float4 (x, y, z, w).
    uv.x -= uniforms.regionOfInterest.x;
    uv.y -= 1.0 - (uniforms.regionOfInterest.y + uniforms.regionOfInterest.w);
    uv /= uniforms.regionOfInterest.zw;
    
    uint4 texColor = segmentationTexture.sample(s, uv);
    if(texColor.r == uniforms.classificationLabelIndex) {
        float time = (sin(uniforms.time * 5.0) + 1) * 0.5;
        out.color = float4(1, 1, 0, 0.25 * time);
        out.depth = uniforms.depthBufferZ;
    } else {
        // Discard color and depth information.
        discard_fragment();
    }
    
    return out;
}