code/tt_moge/reference/utils3d/numpy/rasterization.py

import os
import time
from typing import *
from contextlib import contextmanager
import threading
import weakref

import numpy as np
from numpy import ndarray
import moderngl

from .transforms import extrinsics_to_view, intrinsics_to_perspective

__all__ = [
    'RastContext',
    'rasterize_triangles',
    'rasterize_triangles_peeling',
    'rasterize_lines',
    'rasterize_point_cloud',
    'sample_texture',
    'test_rasterization',
]


def map_np_dtype(dtype) -> str:
    if dtype == int:
        return 'i4'
    elif dtype == np.uint8:
        return 'u1'
    elif dtype == np.uint32:
        return 'u2'
    elif dtype == np.float16:
        return 'f2'
    elif dtype == np.float32:
        return 'f4'
    

def one_value(dtype):
    if dtype == 'u1':
        return 255
    elif dtype == 'u2':
        return 65535
    else:
        return 1


RASTERIZE_TRIANGLES_VS = """
#version 430
uniform mat4 uViewMat;
uniform mat4 uProjectionMat;
in vec3 inVert;
in vec4 inAttr;
{return_interpolation} in vec2 inUV;
{flat} out vec4 vAttr;
out vec3 vViewPos;
{return_interpolation} out vec2 vUV;
void main() {
    vec4 viewPos = uViewMat * vec4(inVert, 1.0);
    vAttr = inAttr;
    vViewPos = viewPos.xyz;
    gl_Position = uProjectionMat * viewPos;
    {return_interpolation} vUV = inUV;
}
"""

RASTERIZE_TRIANGLES_FS = """
#version 430
{flat} in vec4 vAttr;
in vec2 vUV;
in vec3 vViewPos;
out vec4 outAttr;
out float outMask;
{return_interpolation} out int outID;
{return_interpolation} out vec2 outUV;
void main() {
    float currDepth = log2(-vViewPos.z) / 64.f + 0.5f;
    outAttr = vAttr;
    outMask = 1.0;
    {return_interpolation} outID = gl_PrimitiveID;
    {return_interpolation} outUV = vUV;
    gl_FragDepth = currDepth;
}
"""

RASTERIZE_TRIANGLES_PEELING_FS = """
#version 430
uniform sampler2D prevBufferDepthMap;
uniform vec2 uScreenSize;
{flat} in vec4 vAttr;
in vec2 vUV;
in vec3 vViewPos;
out vec4 outAttr;
out float outMask;
{return_interpolation} out int outID;
{return_interpolation} out vec2 outUV;
void main() {
    float currDepth = log2(-vViewPos.z) / 64.f + 0.5f;
    float prevDepth = texture(prevBufferDepthMap, gl_FragCoord.xy / uScreenSize).r;
    if (currDepth <= prevDepth) {
        discard; 
    }
    outAttr = vAttr;
    outMask = 1.0;
    {return_interpolation} outID = gl_PrimitiveID;
    {return_interpolation} outUV = vUV;
    gl_FragDepth = currDepth;
}
"""

RASTERIZE_POINT_CLOUD_VS = """
#version 430
uniform mat4 uViewMat;
uniform mat4 uProjectionMat;
in vec3 inVert;
in vec4 inAttr;
in float inPointSize;
{return_point_id} out int vPointID;
out vec3 vViewPos;
out vec4 vAttr;
out float vPointSize;
void main() {
    vec4 viewPos = uViewMat * vec4(inVert, 1.0);
    vAttr = inAttr;
    vViewPos = viewPos.xyz;
    vPointSize = inPointSize;
    {return_point_id} vPointID = gl_VertexID;
    gl_Position = uProjectionMat * viewPos;
}
"""

RASTERIZE_POINT_CLOUD_GS = """
#version 430 core
layout(points) in;
layout(triangle_strip, max_vertices = {num_point_shape_vertices}) out;
in vec3 vViewPos[];
in vec4 vAttr[];
in float vPointSize[];
{return_point_id} in int vPointID[];
flat out vec4 gAttr;
out vec3 gViewPos;
flat out vec3 gPointViewPos;
flat out vec2 gPointCenter2D;
flat out float gPointSize;
{return_point_id} flat out int gPointID;
uniform mat4 uViewMat;
uniform mat4 uProjectionMat;
uniform vec2 uScreenSize;
uniform bool uIsPointSize3D;
const vec3 triangle[3] = {vec3(0., 0.5, 0.0), vec3(-0.433013, -0.25, 0.0), vec3(0.433013,  -0.25, 0.0)};
const vec3 square[4] = {vec3(-0.5, -0.5, 0.0), vec3(0.5, -0.5, 0.0), vec3(-0.5,  0.5, 0.0), vec3( 0.5,  0.5, 0.0)};
const vec3 pentagon[5] = {vec3(0.0, 0.5, 0.0),  vec3(-0.475528, 0.154508, 0.0),  vec3(0.475528, 0.154508, 0.0),  vec3(-0.293893, -0.404508, 0.0),  vec3(0.293893, -0.404508, 0.0)};
const vec3 hexagon[6] = {vec3(0.0, 0.5, 0.0), vec3(-0.433013, 0.25, 0.0), vec3(0.433013, 0.25, 0.0), vec3(-0.433013, -0.25, 0.0), vec3(0.433013, -0.25, 0.0), vec3(0.0, -0.5, 0.0)};
const vec3 circle[4] = {vec3(-0.5, -0.5, 0.0), vec3(0.5, -0.5, 0.0), vec3(-0.5,  0.5, 0.0), vec3( 0.5,  0.5, 0.0)};
void main() {
    mat4 invProjection = inverse(uProjectionMat);
    gAttr = vAttr[0];
    {return_point_id} gPointID = vPointID[0];
    gPointViewPos = vViewPos[0];
    gPointCenter2D = (0.5 * gl_in[0].gl_Position.xy / gl_in[0].gl_Position.w + 0.5) * uScreenSize;
    gPointSize = vPointSize[0];
    for (int i = 0; i < {point_shape}.length(); ++i) {
        vec3 viewPos = uIsPointSize3D ? gPointViewPos + gPointSize * {point_shape}[i] : gPointViewPos;
        vec4 clipCoord = uProjectionMat * vec4(viewPos, 1.0);
        if (!uIsPointSize3D) clipCoord.xy += 2 * gPointSize * clipCoord.w / uScreenSize * {point_shape}[i].xy;
        gViewPos = (invProjection * clipCoord).xyz;
        gl_Position = clipCoord;
        EmitVertex();
    }
    EndPrimitive();
}
"""

RASTERIZE_POINT_CLOUD_FS = """
#version 430 core
flat  in vec4 gAttr;
in vec3 gViewPos;
flat in vec3 gPointViewPos;
flat in vec2 gPointCenter2D;
flat in float gPointSize;
{return_point_id} flat in int gPointID;
out vec4 outAttr;
out float outMask;
{return_point_id} out int outPointID;
uniform bool uIsPointSize3D;
void main() {
    outAttr = gAttr;
    outMask = 1.0;
    {return_point_id} outPointID = gPointID;
    gl_FragDepth = log2(-gViewPos.z) / 64.f + 0.5f;
    {circle_begin} 
    float distSquare;
    if (uIsPointSize3D) {
        vec3 x = gViewPos - gPointViewPos;
        distSquare = dot(x, x);
    }
    else {
        vec2 x = gl_FragCoord.xy - gPointCenter2D;
        distSquare = dot(x, x);   
    }
    if (distSquare > 0.25 * gPointSize * gPointSize)
        discard;
    {circle_end}
}
"""

FULL_SCREEN_VS = """
#version 430 core
out vec2 screenXY;
void main() {
    screenXY = vec2((gl_VertexID << 1) & 2, gl_VertexID & 2);
    gl_Position = vec4(screenXY * 2.0 - 1.0, 0.0, 1.0);
}
"""

CLEAR_TEXTURE_FS = """
#version 430
uniform {type} uCleanColor;
in vec2 screenXY;
out {type} outColor;
void main() {
    outColor = uCleanColor;
}
"""

SAMPLE_TEXTURE_FS = """
#version 430
uniform sampler2D texMap;
uniform sampler2D uvMap;
in vec2 screenXY;
out vec4 outColor;
void main() {
    outColor = vec4(texture(texMap, texture(uvMap, screenXY).xy));
}
"""

DEPTH_LINEAR_TO_BUFFER_FS = """
#version 430
uniform sampler2D linearDepthMap;
in vec2 screenXY;
void main() {
    float d = texture(linearDepthMap, screenXY).r;
    gl_FragDepth = min(1.f, log2(d) / 64.f + 0.5f);
}
"""

DEPTH_BUFFER_TO_LINEAR_FS = """
#version 430
uniform sampler2D bufferDepthMap;
in vec2 screenXY;
out float outLinearDepth;
void main() {
    float d = texture(bufferDepthMap, screenXY).r;
    outLinearDepth = d == 1.f ? 1.f / 0.f : exp2((d - 0.5f) * 64.f);
}
"""

class RastContext:
    """ Context for numpy-side rasterization. Based on moderngl.
    """
    
    _threading_local: threading.local = threading.local()
    "Thread-local static variable to store default context."

    programs: Dict[str, moderngl.Program]
    "Cache of compiled moderngl programs."

    def __init__(self, *args, **kwargs):
        """Create context for numpy-side rasterization. Based on moderngl.

        ## Parameters
            Leave empty to create context with default settings. Or refer to moderngl.create_context() for other parameters.
        """
        if len(args) == 0 and len(kwargs) == 0:
            try:
                # Default
                self.mgl_ctx = moderngl.create_context()    
            except:
                pass
            if not hasattr(self, 'mgl_ctx'):
                try:
                    # EGL
                    self.mgl_ctx = moderngl.create_context(standalone=True, backend='egl',)     
                except:
                    pass
            if not hasattr(self, 'mgl_ctx'):
                try:
                    # Default standalone
                    self.mgl_ctx = moderngl.create_context(standalone=True)     
                except:
                    pass
            if not hasattr(self, 'mgl_ctx'):
                try:
                    # EGL with more specifications
                    self.mgl_ctx = moderngl.create_context(standalone=True, backend='egl', libgl='libGL.so.1', libegl='libEGL.so.1')    
                except:
                    pass
            if not hasattr(self, 'mgl_ctx'):
                raise RuntimeError(
                    "Failed to create moderngl context with default settings. " 
                    "Please specify context creation settings following https://moderngl.readthedocs.io/en/latest/topics/context.html. "
                    "If running on headless ubuntu server, please refer to https://moderngl.readthedocs.io/en/latest/techniques/headless_ubuntu_18_server.html"
                )
            
        elif len(args) == 1 and isinstance(args[0], moderngl.Context):
            self.mgl_ctx = args[0]
        else:
            self.mgl_ctx = moderngl.create_context(*args, **kwargs)
        self.programs = {}

    def __enter__(self):
        self.mgl_ctx.__enter__()
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.mgl_ctx.__exit__(exc_type, exc_value, traceback)

    def __del__(self):
        try:
            self.mgl_ctx.release()
            for prog_name, prog in self.programs.items():
                prog.release()
        except:
            pass

    @staticmethod
    def get_default_context():
        """ Get the default RastContext in current thread. Create one if not exists.
        """
        if not hasattr(RastContext._threading_local, 'default_context'):
            RastContext._threading_local.default_context = RastContext()
        return RastContext._threading_local.default_context


def run_full_screen_program(
    ctx: RastContext, 
    prog: moderngl.Program, 
    in_tex: Union[moderngl.Texture, List[moderngl.Texture]], 
    out_tex_or_rbo: Union[Union[moderngl.Texture, moderngl.Renderbuffer], Sequence[Union[moderngl.Texture, moderngl.Renderbuffer]]] = None,
    out_depth: Optional[moderngl.Texture] = None,
    **uniforms
) -> moderngl.Texture:
    vao = ctx.mgl_ctx.vertex_array(prog, [])
    if isinstance(in_tex, list):
        for i, tex in enumerate(in_tex):
            tex.use(location=i)
    else:
        in_tex.use(location=0)
    for k, v in uniforms.items():   
        prog[k] = v
    if not isinstance(out_tex_or_rbo, list):
        out_tex_or_rbo = [out_tex_or_rbo] if out_tex_or_rbo is not None else []
    fbo = ctx.mgl_ctx.framebuffer(color_attachments=out_tex_or_rbo, depth_attachment=out_depth)
    fbo.use()
    ctx.mgl_ctx.disable(ctx.mgl_ctx.DEPTH_TEST)
    ctx.mgl_ctx.disable(ctx.mgl_ctx.CULL_FACE)
    ctx.mgl_ctx.disable(ctx.mgl_ctx.BLEND)
    vao.render(moderngl.TRIANGLES, vertices=3)
    vao.release()
    fbo.release()


def clear_texture(ctx: RastContext, tex: moderngl.Texture, value: Union[Tuple[float], Tuple[int], float, int]):
    if tex.depth:
        fbo = ctx.mgl_ctx.framebuffer(color_attachments=[], depth_attachment=tex)
        fbo.clear(depth=value)
        fbo.release()
    elif tex.dtype.startswith('f'):
        fbo = ctx.mgl_ctx.framebuffer(color_attachments=[tex])
        fbo.clear(color=value if isinstance(value, tuple) else (value,) * tex.components)
        fbo.release()
    elif tex.dtype.startswith('i') or tex.dtype.startswith('u'):
        fbo = ctx.mgl_ctx.framebuffer(color_attachments=[tex])
        vec_type = tex.dtype[0] + 'vec4'
        program_name = f'clear_texture_{tex.dtype}'
        if prog := ctx.programs.get(program_name, None) is None:
            prog = ctx.mgl_ctx.program(
                vertex_shader=FULL_SCREEN_VS,
                fragment_shader=CLEAR_TEXTURE_FS.replace('{type}', vec_type)
            )
            ctx.programs[program_name] = prog
        vao = ctx.mgl_ctx.vertex_array(prog, [])
        prog['uCleanColor'].value = value + (0,) * (4 - len(value))
        fbo.use()
        vao.render(moderngl.TRIANGLES, vertices=3)
        fbo.release()
        vao.release()


def depth_texture_buffer_to_linear(ctx: RastContext, buffer_depth_tex: moderngl.Texture, linear_depth_tex: moderngl.Texture):
    program_name = f'depth_buffer_to_linear'
    if (prog := ctx.programs.get(program_name, None)) is None:
        prog = ctx.mgl_ctx.program(
            vertex_shader=FULL_SCREEN_VS,
            fragment_shader=DEPTH_BUFFER_TO_LINEAR_FS
        )
        prog['bufferDepthMap'] = 0
        ctx.programs[program_name] = prog
    run_full_screen_program(ctx, prog, buffer_depth_tex, out_tex_or_rbo=linear_depth_tex)


def depth_texture_linear_to_buffer(ctx: RastContext, linear_depth_tex: moderngl.Texture, buffer_depth_tex: moderngl.Texture):
    program_name = f'depth_linear_to_buffer'
    if (prog := ctx.programs.get(program_name, None)) is None:
        prog = ctx.mgl_ctx.program(
            vertex_shader=FULL_SCREEN_VS,
            fragment_shader=DEPTH_LINEAR_TO_BUFFER_FS
        )
        prog['linearDepthMap'] = 0
        ctx.programs[program_name] = prog
    run_full_screen_program(ctx, prog, linear_depth_tex, out_depth=buffer_depth_tex)


def rasterize_triangles(
    size: Tuple[int, int],
    *,
    vertices: ndarray,
    attributes: Optional[ndarray] = None,
    attributes_domain: Optional[Literal['vertex', 'face']] = 'vertex',
    faces: Optional[ndarray] = None,
    view: ndarray = None,
    projection: ndarray = None,
    extrinsics: ndarray = None,
    intrinsics: ndarray = None,
    near: float = 0.01,
    far: float = float('inf'),
    cull_backface: bool = False,
    return_depth: bool = False,
    return_interpolation: bool = False,
    background: Optional[Dict[str, ndarray]] = None,
    ctx: Optional[RastContext] = None,
) -> Dict[str, ndarray]:
    """
    Rasterize triangles.

    Parameters
    ----
    - `size` (Tuple[int, int]): (height, width) of the output image
    - `vertices` (ndarray): (N, 3) or (T, 3, 3)
    - `faces` (Optional[ndarray]): (T, 3) or None. If `None`, the vertices must be an array with shape (T, 3, 3)
    - `attributes` (ndarray): (N, C), (T, 3, C) for vertex domain or (T, C) for face domain
    - `attributes_domain` (Literal['vertex', 'face']): domain of the attributes
    - `view` | `extrinsics` (ndarray): (4, 4) View matrix or extrinsics matrix. Provide either one of them.
    - `projection` | `intrinsics` (ndarray): (4, 4) Projection matrix or (3, 3) Intrinsics matrix. Provide either one of them.
    - `cull_backface` (bool): whether to cull backface
    - `background` (Optional[Dict[str, ndarray]]): background to composite with. Contains the same keys as the return dictionary, each with shape matching the output.
    - `ctx` (RastContext): rasterization context. Created by `RastContext()`. Default to the thread-local default context.

    Returns
    ----
    A dictionary containing
    
    - `mask` (ndarray): (H, W) bool mask of valid pixels

    if attributes is not None
    - `image` (ndarray): (H, W, C) float32 rendered image corresponding to the input attributes

    if return_depth is True
    - `depth` (ndarray): (H, W) float32 camera space linear depth. Empty pixels are filled with inf.

    if return_interpolation is True
    - `interpolation_id` (ndarray): (H, W) int32 triangle ID map
    - `interpolation_uv` (ndarray): (H, W, 2) float32 triangle UV (first two channels of barycentric coordinates)
    """
    if ctx is None:
        ctx = RastContext.get_default_context()

    height, width = size
    if faces is None:
        assert vertices.ndim == 3 and vertices.shape[1] == vertices.shape[2] == 3, "If faces is None, vertices must be an array with shape (T, 3, 3)"
    else:
        assert faces.ndim == 2 and faces.shape[1] == 3, f"Faces should be a 2D array with shape (T, 3), but got {faces.shape}"
        assert vertices.ndim == 2 and vertices.shape[1] == 3

    assert vertices.dtype == np.float32

    if attributes is not None:
        assert attributes.dtype == np.float32
        if attributes_domain == 'vertex':
            assert (attributes.shape[:-1] == vertices.shape[:-1]), f"Attribute shape {attributes.shape} does not match vertex shape {vertices.shape}"
        elif attributes_domain == 'face':
            if faces is None:
                assert attributes.shape[0] == vertices.shape[0], f"Attribute shape {attributes.shape} does not match vertex shape {vertices.shape}"
            else:
                assert attributes.shape[0] == faces.shape[0], f"Attribute shape {attributes.shape} does not match face shape {faces.shape}"
        else:
            raise ValueError(f"Unknown attributes domain: {attributes_domain}")

        assert attributes.shape[-1] in [1, 2, 3, 4], f'Vertex attribute only supports channels 1, 2, 3, 4, but got {attributes.shape[-1]}'

    assert intrinsics is None or (intrinsics.shape == (3, 3) and intrinsics.dtype == np.float32), f"Intrinsics should be a 3x3 float32 matrix, but got {intrinsics.shape} {intrinsics.dtype}"
    assert extrinsics is None or (extrinsics.shape == (4, 4) and extrinsics.dtype == np.float32), f"Extrinsics should be a 4x4 float32 matrix, but got {extrinsics.shape} {extrinsics.dtype}"
    assert view is None or (view.shape == (4, 4) and view.dtype == np.float32), f"View should be a 4x4 float32 matrix, but got {view.shape} {view.dtype}"
    assert projection is None or (projection.shape == (4, 4) and projection.dtype == np.float32), f"Projection should be a 4x4 float32 matrix, but got {projection.shape} {projection.dtype}"

    background_image, background_mask, background_depth, background_interpolation_id, background_interpolation_uv = [background.get(k, None) if background is not None else None for k in ['image', 'mask', 'depth', 'interpolation_id', 'interpolation_uv']]
    if background_image is not None:
        assert background_image.ndim == 3 and background_image.shape == (height, width, attributes.shape[-1]), f"Image should be a float32 array with shape (H, W, {attributes.shape[1]}), but got {background_image.shape} {background_image.dtype}"
    if background_depth is not None:
        assert background_depth.dtype == np.float32 and background_depth.ndim == 2 and background_depth.shape == (height, width), f"Depth should be a float32 array with shape (H, W), but got {background_depth.shape} {background_depth.dtype}"
    if background_interpolation_id is not None:
        assert background_interpolation_id.dtype == np.int32 and background_interpolation_id.ndim == 2 and background_interpolation_id.shape == (height, width), f"Interpolation ID should be a int32 array with shape (H, W), but got {background_interpolation_id.shape} {background_interpolation_id.dtype}"
    if background_interpolation_uv is not None:
        assert background_interpolation_uv.dtype == np.float32 and background_interpolation_uv.ndim == 3 and background_interpolation_uv.shape == (height, width, 2), f"Interpolation UV should be a float32 array with shape (H, W, 2), but got {background_interpolation_uv.shape} {background_interpolation_uv.dtype}"

    if faces is not None:
        vertices = vertices[faces]
    if attributes is not None:
        num_channels = attributes.shape[-1]
        attributes = np.concatenate([attributes, np.zeros((*attributes.shape[:-1], 4 - num_channels,), dtype=attributes.dtype)], axis=-1) if num_channels < 4 else attributes
        if attributes_domain == 'vertex':
            if faces is not None:
                attributes = attributes[faces]
        elif attributes_domain == 'face':
            attributes = attributes[:, None, :].repeat(3, axis=1)

    if intrinsics is not None:
        assert projection is None, "Cannot specify both intrinsics and projection."
        projection = intrinsics_to_perspective(intrinsics, near, far)
    if extrinsics is not None:
        assert view is None, "Cannot specify both extrinsics and view."
        view = extrinsics_to_view(extrinsics)
    if view is None:
        view = np.eye(4, dtype=np.float32)
    if projection is None:
        projection = np.eye(4, dtype=np.float32) 

    if background_image is not None:
        background_image = np.concatenate([background_image, np.ones((height, width, 4 - background_image.shape[-1]), dtype=background_image.dtype)], axis=-1) if background_image.shape[-1] < 4 else background_image

    # Get program
    flat = attributes_domain == 'face'
    program_name = f'rasterize_triangles(flat={flat}, return_interpolation={return_interpolation})'
    if  (prog := ctx.programs.get(program_name, None)) is None:
        prog = ctx.mgl_ctx.program(
            vertex_shader=RASTERIZE_TRIANGLES_VS.replace('{flat}', 'flat' if flat else '').replace('{return_interpolation}', '' if return_interpolation else '//'),
            fragment_shader=RASTERIZE_TRIANGLES_FS.replace('{flat}', 'flat' if flat else '').replace('{return_interpolation}', '' if return_interpolation else '//')
        )
        ctx.programs[program_name] = prog

    # Create buffers
    vbo_vertices = ctx.mgl_ctx.buffer(np.ascontiguousarray(vertices, dtype='f4'))
    if attributes is not None:
        vbo_attributes = ctx.mgl_ctx.buffer(np.ascontiguousarray(attributes, dtype='f4'))
    if return_interpolation:
        vbo_uv = ctx.mgl_ctx.buffer(np.ascontiguousarray((np.array([[1., 0.], [0., 1.], [0., 0.]], dtype=np.float32).reshape(1, 3, 2).repeat(vertices.shape[0], axis=0))))
    
    vao = ctx.mgl_ctx.vertex_array(
        prog,
        list(filter(lambda x: x is not None, [
            (vbo_vertices, '3f', 'inVert'),
            (vbo_attributes, f'4f', 'inAttr') if attributes is not None else None,
            (vbo_uv, '2f', 'inUV') if return_interpolation else None,
        ])),
        mode=moderngl.TRIANGLES,
    )

    # Create textures
    image_tex = ctx.mgl_ctx.texture((width, height), 4, dtype='f4', data=np.ascontiguousarray(background_image[::-1, :, :]) if background_image is not None else None)
    buffer_depth_tex = ctx.mgl_ctx.depth_texture((width, height))
    mask_tex = ctx.mgl_ctx.texture((width, height), 1, dtype='f1', data=np.ascontiguousarray(background_mask[::-1, :]) if background_mask is not None else None)
    if return_depth:
        linear_depth_tex = ctx.mgl_ctx.texture((width, height), 1, dtype='f4', data=np.ascontiguousarray(background_depth[::-1, :]) if background_depth is not None else None)
    if return_interpolation:
        interpolation_id_tex = ctx.mgl_ctx.texture((width, height), 1, dtype='i4', data=np.ascontiguousarray(background_interpolation_id[::-1, :]) if background_interpolation_id is not None else None)
        interpolation_uv_tex = ctx.mgl_ctx.texture((width, height), 2, dtype='f4', data=np.ascontiguousarray(background_interpolation_uv[::-1, :, :]) if background_interpolation_uv is not None else None)
    
    # Clear textures
    if background_image is None:
        clear_texture(ctx, image_tex, value=(0.0, 0.0, 0.0, 1.0))
    if background_mask is None:
        clear_texture(ctx, mask_tex, value=0.0)
    if background_depth is not None:
        depth_texture_linear_to_buffer(ctx, linear_depth_tex, buffer_depth_tex)
    else:
        clear_texture(ctx, buffer_depth_tex, value=1.0)
    if return_interpolation:
        if background_interpolation_id is None:
            clear_texture(ctx, interpolation_id_tex, value=(-1,))
        if background_interpolation_uv is None:
            clear_texture(ctx, interpolation_uv_tex, value=(0.0, 0.0))

    # Create framebuffer
    if return_interpolation:
        fbo = ctx.mgl_ctx.framebuffer(
            color_attachments=[image_tex, mask_tex, interpolation_id_tex, interpolation_uv_tex],
            depth_attachment=buffer_depth_tex,
        )
    else:
        fbo = ctx.mgl_ctx.framebuffer(
            color_attachments=[image_tex, mask_tex],
            depth_attachment=buffer_depth_tex,
        )
    fbo.viewport = (0, 0, width, height)

    # Set uniforms
    prog['uViewMat'].write(np.ascontiguousarray(view.transpose().astype('f4')))
    prog['uProjectionMat'].write(np.ascontiguousarray(projection.transpose().astype('f4')))

    # Set render states
    ctx.mgl_ctx.depth_func = '<'
    ctx.mgl_ctx.enable(ctx.mgl_ctx.DEPTH_TEST)
    if cull_backface:
        ctx.mgl_ctx.enable(ctx.mgl_ctx.CULL_FACE)
    else:
        ctx.mgl_ctx.disable(ctx.mgl_ctx.CULL_FACE)
    ctx.mgl_ctx.disable(ctx.mgl_ctx.BLEND)
    
    # Render
    fbo.use()
    vao.render()

    # Read
    mask = np.frombuffer(mask_tex.read(), dtype=bool).reshape((height, width))
    mask = np.flip(mask, axis=0)
    if attributes is not None:
        image = np.frombuffer(image_tex.read(), dtype='f4').reshape((height, width, 4))
        image = np.flip(image, axis=0)
        image = image[:, :, :num_channels]
    else:
        image = None
    if return_depth:
        depth_texture_buffer_to_linear(ctx, buffer_depth_tex, linear_depth_tex)
        depth = np.frombuffer(linear_depth_tex.read(), dtype='f4').reshape((height, width))
        depth = np.flip(depth, axis=0)
    else:
        depth = None
    if return_interpolation:
        interpolation_id = np.frombuffer(interpolation_id_tex.read(), dtype='i4').reshape((height, width))
        interpolation_id = np.flip(interpolation_id, axis=0)
        interpolation_uv = np.frombuffer(interpolation_uv_tex.read(), dtype='f4').reshape((height, width, 2))
        interpolation_uv = np.flip(interpolation_uv, axis=0)
    else:
        interpolation_id = None
        interpolation_uv = None

    # Release
    vao.release()
    vbo_vertices.release()
    if attributes is not None:
        vbo_attributes.release()
    if return_interpolation:
        vbo_uv.release()
    fbo.release()
    image_tex.release()
    mask_tex.release()
    buffer_depth_tex.release()
    if background_depth is not None or return_depth:
        linear_depth_tex.release()
    if return_interpolation:
        interpolation_id_tex.release()
        interpolation_uv_tex.release()

    output = {
        "image": image,
        "mask": mask,
        "depth": depth,
        "interpolation_id": interpolation_id,
        "interpolation_uv": interpolation_uv
    }

    output = {k: v for k, v in output.items() if v is not None}

    return output


@contextmanager
def rasterize_triangles_peeling(
    size: Tuple[int, int],
    *,
    vertices: ndarray,
    attributes: ndarray,
    attributes_domain: Literal['vertex', 'face'] = 'vertex',
    faces: Optional[ndarray] = None,
    view: ndarray = None,
    projection: ndarray = None,
    extrinsics: ndarray = None,
    intrinsics: ndarray = None,
    near: float = 0.01,
    far: float = float('inf'),
    cull_backface: bool = False,
    return_depth: bool = False,
    return_interpolation: bool = False,
    ctx: Optional[RastContext] = None,
) -> Iterator[Iterator[Dict[str, ndarray]]]:
    """
    Rasterize triangles with depth peeling.

    Parameters
    ----
    - `size` (Tuple[int, int]): (height, width) of the output image
    - `vertices` (ndarray): (N, 3) or (T, 3, 3)
    - `faces` (Optional[ndarray]): (T, 3) or None. If `None`, the vertices must be an array with shape (T, 3, 3)
    - `attributes` (ndarray): (N, C), (T, 3, C) for vertex domain or (T, C) for face domain
    - `attributes_domain` (Literal['vertex', 'face']): domain of the attributes
    - `view` | `extrinsics` (ndarray): (4, 4) View matrix or extrinsics matrix. Provide either one of them.
    - `projection` | `intrinsics` (ndarray): (4, 4) Projection matrix or (3, 3) Intrinsics matrix. Provide either one of them.
    - `near` (float): near clipping plane. Only used for intrinsics. Ignored if projection matrix is provided.
    - `far` (float): far clipping plane. Only used for intrinsics. Ignored if projection matrix is provided.
    - `cull_backface` (bool): whether to cull backfaces
    - `ctx` (RastContext): rasterization context. Created by `RastContext()`. Default to the thread-local default context.

    Returns
    ----
    A generator that yields dictionaries for each peeling layer containing:
    
    - `mask` (ndarray): (H, W) bool mask of valid pixels in this layer

    if attributes is not None
    - `image` (ndarray): (H, W, C) float32 rendered image corresponding to the input attributes

    if return_depth is True
    - `depth` (ndarray): (H, W) float32 camera space linear depth. Empty pixels are filled with inf.
    
    if return_interpolation is True
    - `interpolation_id` (ndarray): (H, W) int32 triangle ID map
    - `interpolation_uv` (ndarray): (H, W, 2) float32 triangle UV (first two channels of barycentric coordinates)

    The last layer yielded will be empty (no pixels covered), then the generator stops.

    Example
    ----
    ```
    for i, layer_output in enumerate(rasterize_triangles_peeling(
        (512, 512), 
        vertices=vertices, 
        faces=faces, 
        attributes=attributes,
        view=view,
        projection=projection
    )):
        print(f"Layer {i}:")
        for key, value in layer_output.items():
            print(f"  {key}: {value.shape}")
        if i >= 4:  # Stop after 5 layers at most
            break
    ```
    """
    if ctx is None:
        ctx = RastContext.get_default_context()

    height, width = size
    if faces is None:
        assert vertices.ndim == 3 and vertices.shape[1] == vertices.shape[2] == 3, "If faces is None, vertices must be an array with shape (T, 3, 3)"
    else:
        assert faces.ndim == 2 and faces.shape[1] == 3, f"Faces should be a 2D array with shape (T, 3), but got {faces.shape}"
        assert vertices.ndim == 2 and vertices.shape[1] == 3

    assert vertices.dtype == np.float32

    if attributes is not None:
        assert attributes.dtype == np.float32
        if attributes_domain == 'vertex':
            assert (attributes.shape[:-1] == vertices.shape[:-1]), f"Attribute shape {attributes.shape} does not match vertex shape {vertices.shape}"
        elif attributes_domain == 'face':
            if faces is None:
                assert attributes.shape[0] == vertices.shape[0], f"Attribute shape {attributes.shape} does not match vertex shape {vertices.shape}"
            else:
                assert attributes.shape[0] == faces.shape[0], f"Attribute shape {attributes.shape} does not match face shape {faces.shape}"
        else:
            raise ValueError(f"Unknown attributes domain: {attributes_domain}")

        assert attributes.shape[-1] in [1, 2, 3, 4], f'Vertex attribute only supports channels 1, 2, 3, 4, but got {attributes.shape[-1]}'

    assert intrinsics is None or (intrinsics.shape == (3, 3) and intrinsics.dtype == np.float32), f"Intrinsics should be a 3x3 float32 matrix, but got {intrinsics.shape} {intrinsics.dtype}"
    assert extrinsics is None or (extrinsics.shape == (4, 4) and extrinsics.dtype == np.float32), f"Extrinsics should be a 4x4 float32 matrix, but got {extrinsics.shape} {extrinsics.dtype}"
    assert view is None or (view.shape == (4, 4) and view.dtype == np.float32), f"View should be a 4x4 float32 matrix, but got {view.shape} {view.dtype}"
    assert projection is None or (projection.shape == (4, 4) and projection.dtype == np.float32), f"Projection should be a 4x4 float32 matrix, but got {projection.shape} {projection.dtype}"

    if faces is not None:
        vertices = vertices[faces]
    if attributes is not None:
        num_channels = attributes.shape[-1]
        attributes = np.concatenate([attributes, np.zeros((*attributes.shape[:-1], 4 - num_channels,), dtype=attributes.dtype)], axis=-1) if num_channels < 4 else attributes
        if attributes_domain == 'vertex':
            if faces is not None:
                attributes = attributes[faces]
        elif attributes_domain == 'face':
            attributes = attributes[:, None, :].repeat(3, axis=1)
    
    if extrinsics is not None:
        assert view is None, "Cannot specify both extrinsics and view."
        view = extrinsics_to_view(extrinsics)
    if intrinsics is not None:
        assert projection is None, "Cannot specify both intrinsics and projection."
        projection = intrinsics_to_perspective(intrinsics, near, far)
    if view is None:
        view = np.eye(4, dtype=np.float32)
    if projection is None:
        projection = np.eye(4, dtype=np.float32) 
    
    # Get program
    flat = attributes_domain == 'face'
    program_name = f'rasterize_triangles_peeling(flat={flat}, return_interpolation={return_interpolation})'
    if (prog := ctx.programs.get(program_name, None)) is None:
        prog = ctx.mgl_ctx.program(
            vertex_shader=RASTERIZE_TRIANGLES_VS.replace('{flat}', 'flat' if flat else '').replace('{return_interpolation}', '' if return_interpolation else '//'),
            fragment_shader=RASTERIZE_TRIANGLES_PEELING_FS.replace('{flat}', 'flat' if flat else '').replace('{return_interpolation}', '' if return_interpolation else '//')
        )
        prog['prevBufferDepthMap'] = 0
        ctx.programs[program_name] = prog
    
    # Create buffers
    vbo_vertices = ctx.mgl_ctx.buffer(np.ascontiguousarray(vertices, dtype='f4'))
    if attributes is not None:
        vbo_attributes = ctx.mgl_ctx.buffer(np.ascontiguousarray(attributes, dtype='f4'))
    if return_interpolation:
        vbo_uv = ctx.mgl_ctx.buffer(np.ascontiguousarray((np.array([[1., 0.], [0., 1.], [0., 0.]], dtype=np.float32).reshape(1, 3, 2).repeat(vertices.shape[0], axis=0))))
    
    vao = ctx.mgl_ctx.vertex_array(
        prog,
        list(filter(lambda x: x is not None, [
            (vbo_vertices, '3f', 'inVert'),
            (vbo_attributes, f'4f', 'inAttr') if attributes is not None else None,
            (vbo_uv, '2f', 'inUV') if return_interpolation else None,
        ])),
        mode=moderngl.TRIANGLES,
    )

    # Create textures
    image_tex = ctx.mgl_ctx.texture((width, height), 4, dtype='f4')
    mask_tex = ctx.mgl_ctx.texture((width, height), 1, dtype='f1')
    buffer_depth_tex_a = ctx.mgl_ctx.depth_texture((width, height))
    buffer_depth_tex_b = ctx.mgl_ctx.depth_texture((width, height))
    if return_depth:
        linear_depth_tex = ctx.mgl_ctx.texture((width, height), 1, dtype='f4')
    if return_interpolation:
        interpolation_id_tex = ctx.mgl_ctx.texture((width, height), 1, dtype='i4')
        interpolation_uv_tex = ctx.mgl_ctx.texture((width, height), 2, dtype='f4')

    # Create frame buffers
    if return_interpolation:
        color_attachments=[image_tex, interpolation_id_tex, interpolation_uv_tex]
    else:
        color_attachments=[image_tex]
    fbo_curr = ctx.mgl_ctx.framebuffer(
        color_attachments=color_attachments,
        depth_attachment=buffer_depth_tex_a,
    )
    fbo_curr.viewport = (0, 0, width, height)
    fbo_prev = ctx.mgl_ctx.framebuffer(
        color_attachments=color_attachments,
        depth_attachment=buffer_depth_tex_b,
    )
    fbo_prev.viewport = (0, 0, width, height)

    # Set uniforms
    prog['uViewMat'].write(np.ascontiguousarray(view.transpose().astype('f4')))
    prog['uProjectionMat'].write(np.ascontiguousarray(projection.transpose().astype('f4')))
    prog['uScreenSize'].value = (width, height)

    # Rendering settings
    ctx.mgl_ctx.depth_func = '<'
    ctx.mgl_ctx.enable(ctx.mgl_ctx.DEPTH_TEST)
    if cull_backface:
        ctx.mgl_ctx.enable(ctx.mgl_ctx.CULL_FACE)
    else:
        ctx.mgl_ctx.disable(ctx.mgl_ctx.CULL_FACE)
    ctx.mgl_ctx.disable(ctx.mgl_ctx.BLEND)

    # Initialize prev depth texture = 0
    clear_texture(ctx, fbo_prev.depth_attachment, value=0.0)
    
    try:
        while True:
            # Clear
            clear_texture(ctx, image_tex, value=(0.0, 0.0, 0.0, 1.0))
            clear_texture(ctx, mask_tex, value=0.0)
            clear_texture(ctx, fbo_curr.depth_attachment, value=1.0)
            if return_interpolation:
                clear_texture(ctx, interpolation_id_tex, value=(-1,))
                clear_texture(ctx, interpolation_uv_tex, value=(0.0, 0.0))

            # Render
            fbo_prev.depth_attachment.use(location=0) 
            fbo_curr.use()
            vao.render()

            # Read
            mask = np.frombuffer(mask_tex.read(), dtype=bool).reshape((height, width))
            mask = np.flip(mask, axis=0)
            if attributes is not None:
                image = np.frombuffer(image_tex.read(), dtype='f4').reshape((height, width, 4))
                image = np.flip(image, axis=0)
                image = image[:, :, :num_channels]
            else:
                image = None
            if return_depth:
                depth_texture_buffer_to_linear(ctx, fbo_curr.depth_attachment, linear_depth_tex)
                depth = np.frombuffer(linear_depth_tex.read(), dtype='f4').reshape((height, width))
                depth = np.flip(depth, axis=0)
            else:
                depth = None
            if return_interpolation:
                interpolation_id = np.frombuffer(interpolation_id_tex.read(), dtype='i4').reshape((height, width))
                interpolation_id = np.flip(interpolation_id, axis=0)
                interpolation_uv = np.frombuffer(interpolation_uv_tex.read(), dtype='f4').reshape((height, width, 2))
                interpolation_uv = np.flip(interpolation_uv, axis=0)
            else:
                interpolation_id = None
                interpolation_uv = None

            # Yield
            output = {
                "image": image,
                "mask": mask,
                "depth": depth,
                "interpolation_id": interpolation_id,
                "interpolation_uv": interpolation_uv
            }
            output = {k: v for k, v in output.items() if v is not None}
            yield output

            # Check termination
            if not np.any(mask):
                break  # No more layers

            # Swap curr and prev fbos
            fbo_curr, fbo_prev = fbo_prev, fbo_curr
    finally:
        # Release
        vao.release()
        vbo_vertices.release()
        vbo_attributes.release()
        fbo_curr.release()
        fbo_prev.release()
        image_tex.release()
        mask_tex.release()
        buffer_depth_tex_a.release()
        buffer_depth_tex_b.release()
        if return_depth:
            linear_depth_tex.release()
        if return_interpolation:
            interpolation_id_tex.release()
            interpolation_uv_tex.release()


def rasterize_lines(
    size: Tuple[int, int],
    *,
    vertices: ndarray,
    attributes: Optional[ndarray],
    attributes_domain: Literal['vertex', 'line'] = 'vertex',
    lines: ndarray | None = None,
    view: Optional[ndarray] = None,
    projection: Optional[ndarray] = None,
    extrinsics: Optional[ndarray] = None,
    intrinsics: Optional[ndarray] = None,
    near: float = 0.01,
    far: float = float('inf'),
    line_width: float = 1.0,
    return_depth: bool = False,
    return_interpolation: bool = False,
    background: Optional[Dict[str, ndarray]] = None,
    ctx: Optional[RastContext] = None,
) -> Tuple[ndarray, ...]:
    """
    Rasterize lines.

    Parameters
    ----
    - `size` (Tuple[int, int]): (height, width) of the output image
    - `vertices` (ndarray): (N, 3) or (L, 2, 3)
    - `attributes` (ndarray): (N, C), (T, 3, C) for vertex domain or (T, C) for face domain
    - `attributes_domain` (Literal['vertex', 'face']): domain of the attributes
    - `lines` (ndarray): (L, 2) int32 array of line vertex indices. If None, vertices shape must be (L, 2, 3) and will be interpreted as line segments directly.
    - `view` | `extrinsics` (ndarray): (4, 4) View matrix or extrinsics matrix. Provide either one of them.
    - `projection` | `intrinsics` (ndarray): (4, 4) Projection matrix or (3, 3) Intrinsics matrix. Provide either one of them.
    - `near` (float): near clipping plane. Only used for intrinsics. Ignored if projection matrix is provided.
    - `far` (float): far clipping plane. Only used for intrinsics. Ignored if projection matrix is provided.
    - `background` (Optional[Dict[str, ndarray]]): background to composite with. Contains the same keys as the return dictionary, each with shape matching the output.
    - `ctx` (RastContext): rasterization context. Created by `RastContext()`. Defaults to the current default context.

    Returns
    ----
    A dictionary containing
    
    - `mask` (ndarray): (H, W) bool mask of valid pixels

    if attributes is not None
    - `image` (ndarray): (H, W, C) float32 rendered image corresponding to the input attributes

    if return_depth is True
    - `depth` (ndarray): (H, W) float32 camera space linear depth, ranging from 0 to 1.
    
    if return_interpolation is True
    - `interpolation_id` (ndarray): (H, W) int32 triangle ID map
    - `interpolation_uv` (ndarray): (H, W, 2) float32 triangle UV (first two channels of barycentric coordinates)
    """
    if ctx is None:
        ctx = RastContext.get_default_context()

    height, width = size
    if lines is None:
        assert vertices.ndim == 3 and vertices.shape[1] == 2 and vertices.shape[2] == 3, "If lines is None, vertices must be an array with shape (T, 2, 3)"
    else:
        assert lines.ndim == 2 and lines.shape[1] == 2, f"Lines should be a int32 or uint32 array with shape (T, 2), but got {lines.shape} {lines.dtype}"
        assert lines.dtype == np.uint32 or lines.dtype == np.int32
        assert vertices.ndim == 2 and vertices.shape[1] == 3

    assert vertices.dtype == np.float32

    if attributes is not None:
        assert attributes.dtype == np.float32
        if attributes_domain == 'vertex':
            assert (attributes.shape[:-1] == vertices.shape[:-1]), f"Attribute shape {attributes.shape} does not match vertex shape {vertices.shape}"
        elif attributes_domain == 'line':
            if lines is None:
                assert attributes.shape[0] == vertices.shape[0], f"Attribute shape {attributes.shape} does not match vertex shape {vertices.shape}"
            else:
                assert attributes.shape[0] == lines.shape[0], f"Attribute shape {attributes.shape} does not match line shape {lines.shape}"
        else:
            raise ValueError(f"Unknown attributes domain: {attributes_domain}")

        assert attributes.shape[-1] in [1, 2, 3, 4], f'Vertex attribute only supports channels 1, 2, 3, 4, but got {attributes.shape[-1]}'

    assert intrinsics is None or (intrinsics.shape == (3, 3) and intrinsics.dtype == np.float32), f"Intrinsics should be a 3x3 float32 matrix, but got {intrinsics.shape} {intrinsics.dtype}"
    assert extrinsics is None or (extrinsics.shape == (4, 4) and extrinsics.dtype == np.float32), f"Extrinsics should be a 4x4 float32 matrix, but got {extrinsics.shape} {extrinsics.dtype}"
    assert view is None or (view.shape == (4, 4) and view.dtype == np.float32), f"View should be a 4x4 float32 matrix, but got {view.shape} {view.dtype}"
    assert projection is None or (projection.shape == (4, 4) and projection.dtype == np.float32), f"Projection should be a 4x4 float32 matrix, but got {projection.shape} {projection.dtype}"

    background_image, background_mask, background_depth, background_interpolation_id, background_interpolation_uv = [background.get(k, None) if background is not None else None for k in ['image', 'mask', 'depth', 'interpolation_id', 'interpolation_uv']]
    if background_image is not None:
        assert background_image.ndim == 3 and background_image.shape == (height, width, attributes.shape[-1]), f"Image should be a float32 array with shape (H, W, {attributes.shape[1]}), but got {background_image.shape} {background_image.dtype}"
    if background_mask is not None:
        assert background_mask.dtype == np.bool_ and background_mask.ndim == 2 and background_mask.shape == (height, width), f"Mask should be a bool array with shape (H, W), but got {background_mask.shape} {background_mask.dtype}"
    if background_depth is not None:
        assert background_depth.dtype == np.float32 and background_depth.ndim == 2 and background_depth.shape == (height, width), f"Depth should be a float32 array with shape (H, W), but got {background_depth.shape} {background_depth.dtype}"
    if background_interpolation_id is not None:
        assert background_interpolation_id.dtype == np.int32 and background_interpolation_id.ndim == 2 and background_interpolation_id.shape == (height, width), f"Interpolation ID should be a int32 array with shape (H, W), but got {background_interpolation_id.shape} {background_interpolation_id.dtype}"
    if background_interpolation_uv is not None:
        assert background_interpolation_uv.dtype == np.float32 and background_interpolation_uv.ndim == 3 and background_interpolation_uv.shape == (height, width, 2), f"Interpolation UV should be a float32 array with shape (H, W, 2), but got {background_interpolation_uv.shape} {background_interpolation_uv.dtype}"

    if lines is not None:
        vertices = vertices[lines]
    if attributes is not None:
        num_channels = attributes.shape[-1]
        attributes = np.concatenate([attributes, np.zeros((*attributes.shape[:-1], 4 - num_channels,), dtype=attributes.dtype)], axis=-1) if num_channels < 4 else attributes
        if attributes_domain == 'vertex':
            if lines is not None:
                attributes = attributes[lines]
        elif attributes_domain == 'line':
            attributes = attributes[:, None, :].repeat(2, axis=1)

    if extrinsics is not None:
        assert view is None, "Cannot specify both extrinsics and view."
        view = extrinsics_to_view(extrinsics)
    if intrinsics is not None:
        assert projection is None, "Cannot specify both intrinsics and projection."
        projection = intrinsics_to_perspective(intrinsics, near, far)
    if view is None:
        view = np.eye(4, dtype=np.float32)
    if projection is None:
        projection = np.eye(4, dtype=np.float32) 
    if background_image is not None:
        background_image = np.concatenate([background_image, np.ones((height, width, 4 - background_image.shape[-1]), dtype=background_image.dtype)], axis=-1) if background_image.shape[-1] < 4 else background_image

    # Get program (shared with triangle rasterization)
    flat = attributes_domain == 'line'
    program_name = f'rasterize_triangles(flat={flat}, return_interpolation={return_interpolation})'
    if  (prog := ctx.programs.get(program_name, None)) is None:
        prog = ctx.mgl_ctx.program(
            vertex_shader=RASTERIZE_TRIANGLES_VS.replace('{flat}', 'flat' if flat else '').replace('{return_interpolation}', '' if return_interpolation else '//'),
            fragment_shader=RASTERIZE_TRIANGLES_FS.replace('{flat}', 'flat' if flat else '').replace('{return_interpolation}', '' if return_interpolation else '//')
        )
        ctx.programs[program_name] = prog

    # Create buffers
    vbo_vertices = ctx.mgl_ctx.buffer(np.ascontiguousarray(vertices, dtype='f4'))
    if attributes is not None:
        vbo_attributes = ctx.mgl_ctx.buffer(np.ascontiguousarray(attributes, dtype='f4'))
    if return_interpolation:
        vbo_uv = ctx.mgl_ctx.buffer(np.ascontiguousarray((np.array([[1., 0.], [0., 1.]], dtype=np.float32).reshape(1, 2, 2).repeat(vertices.shape[0], axis=0))))
    
    vao = ctx.mgl_ctx.vertex_array(
        prog,
        list(filter(lambda x: x is not None, [
            (vbo_vertices, '3f', 'inVert'),
            (vbo_attributes, f'4f', 'inAttr') if attributes is not None else None,
            (vbo_uv, '2f', 'inUV') if return_interpolation else None,
        ])),
        mode=moderngl.LINES,
    )

    # Create textures
    image_tex = ctx.mgl_ctx.texture((width, height), 4, dtype='f4', data=np.ascontiguousarray(background_image[::-1, :, :]) if background_image is not None else None)
    mask_tex = ctx.mgl_ctx.texture((width, height), 1, dtype='f1', data=np.ascontiguousarray(background_mask[::-1, :]) if background_mask is not None else None)
    buffer_depth_tex = ctx.mgl_ctx.depth_texture((width, height))
    if return_depth is not None:
        linear_depth_tex = ctx.mgl_ctx.texture((width, height), 1, dtype='f4', data=np.ascontiguousarray(background_depth[::-1, :]) if background_depth is not None else None)

    if return_interpolation:
        interpolation_id_tex = ctx.mgl_ctx.texture((width, height), 1, dtype='i4', data=np.ascontiguousarray(background_interpolation_id[::-1, :]) if background_interpolation_id is not None else None)
        interpolation_uv_tex = ctx.mgl_ctx.texture((width, height), 2, dtype='f4', data=np.ascontiguousarray(background_interpolation_uv[::-1, :, :]) if background_interpolation_uv is not None else None)
    
    if background_image is None:
        clear_texture(ctx, image_tex, value=(0.0, 0.0, 0.0, 1.0))
    if background_mask is None:
        clear_texture(ctx, mask_tex, value=0.0)
    if background_depth is None:
        clear_texture(ctx, buffer_depth_tex, value=1.0)
    else:
        depth_texture_linear_to_buffer(ctx, linear_depth_tex, buffer_depth_tex)
    if return_interpolation:
        if background_interpolation_id is None:
            clear_texture(ctx, interpolation_id_tex, value=(-1,))
        if background_interpolation_uv is None:
            clear_texture(ctx, interpolation_uv_tex, value=(0.0, 0.0))

    # Create framebuffer
    if return_interpolation:
        fbo = ctx.mgl_ctx.framebuffer(
            color_attachments=[image_tex, mask_tex, interpolation_id_tex, interpolation_uv_tex],
            depth_attachment=buffer_depth_tex,
        )
    else:
        fbo = ctx.mgl_ctx.framebuffer(
            color_attachments=[image_tex, mask_tex],
            depth_attachment=buffer_depth_tex,
        )
    fbo.viewport = (0, 0, width, height)

    # Set uniforms
    prog['uViewMat'].write(np.ascontiguousarray(view.transpose().astype('f4')))
    prog['uProjectionMat'].write(np.ascontiguousarray(projection.transpose().astype('f4')))

    # Set render states
    ctx.mgl_ctx.depth_func = '<'
    ctx.mgl_ctx.enable(ctx.mgl_ctx.DEPTH_TEST)
    ctx.mgl_ctx.disable(ctx.mgl_ctx.BLEND)
    ctx.mgl_ctx.line_width = line_width

    # Render
    fbo.use()
    vao.render()

    # Read
    mask = np.frombuffer(mask_tex.read(), dtype=bool).reshape((height, width))
    mask = np.flip(mask, axis=0)
    if attributes is not None:
        image = np.frombuffer(image_tex.read(), dtype='f4').reshape((height, width, 4))
        image = np.flip(image, axis=0)
        image = image[:, :, :num_channels]
    else:
        image = None
    if return_depth:
        depth_texture_buffer_to_linear(ctx, buffer_depth_tex, linear_depth_tex)
        depth = np.frombuffer(linear_depth_tex.read(), dtype='f4').reshape((height, width))
        depth = np.flip(depth, axis=0)
    else:
        depth = None
    if return_interpolation:
        interpolation_id = np.frombuffer(interpolation_id_tex.read(), dtype='i4').reshape((height, width))
        interpolation_id = np.flip(interpolation_id, axis=0)
        interpolation_uv = np.frombuffer(interpolation_uv_tex.read(), dtype='f4').reshape((height, width, 2))
        interpolation_uv = np.flip(interpolation_uv, axis=0)
    else:
        interpolation_id = None
        interpolation_uv = None

    # Release
    vao.release()
    vbo_vertices.release()
    if attributes is not None:
        vbo_attributes.release()
    if return_interpolation:
        vbo_uv.release()
    fbo.release()
    image_tex.release()
    mask_tex.release()
    buffer_depth_tex.release()
    if background_depth is not None or return_depth:
        linear_depth_tex.release()
    if return_interpolation:
        interpolation_id_tex.release()
        interpolation_uv_tex.release()

    output = {
        "image": image,
        "mask": mask,
        "depth": depth,
        "interpolation_id": interpolation_id,
        "interpolation_uv": interpolation_uv
    }

    output = {k: v for k, v in output.items() if v is not None}

    return output


def rasterize_point_cloud(
    size: Tuple[int, int],
    *,
    points: ndarray,
    point_sizes: Union[float, ndarray] = 10,
    point_size_in: Literal['2d', '3d'] = '2d',
    point_shape: Literal['triangle', 'square', 'pentagon', 'hexagon', 'circle'] = 'square',
    attributes: Optional[ndarray] = None,
    view: ndarray = None,
    projection: ndarray = None,
    extrinsics: ndarray = None,
    intrinsics: ndarray = None,
    near: float = 0.01,
    far: float = float('inf'),
    return_depth: bool = False,
    return_point_id: bool = False,
    background: Optional[Dict[str, ndarray]] = None,
    ctx: Optional[RastContext] = None,
) -> Dict[str, ndarray]:
    """
    Rasterize point cloud.

    Parameters
    ----
    - `size` (Tuple[int, int]): (height, width) of the output image
    - `points` (ndarray): (N, 3)
    - `point_sizes` (ndarray): (N,) or float
    - `point_size_in`: Literal['2d', '3d'] = '2d'. Whether the point sizes are in 2D (size in pixels) or 3D (size in world units).
    - `point_shape`: Literal['triangle', 'square', 'pentagon', 'hexagon', 'circle'] = 'square'. The visual shape of the points.
    - `attributes` (ndarray): (N, C)
    - `view` | `extrinsics` (ndarray): (4, 4) View matrix or extrinsics matrix. Provide either one of them.
    - `projection` | `intrinsics` (ndarray): (4, 4) Projection matrix or (3, 3) Intrinsics matrix. Provide either one of them.
    - `near` (float): near clipping plane. Only used for intrinsics. Ignored if projection matrix is provided.
    - `far` (float): far clipping plane. Only used for intrinsics. Ignored if projection matrix is provided.
    - `return_depth` (bool): whether to return depth map
    - `return_point_id` (bool): whether to return point ID map
    - `background` (Optional[Dict[str, ndarray]]): background to composite with. Contains the same keys as the return dictionary, each with shape matching the output.
    - `ctx` (RastContext): rasterization context. Created by `RastContext()`. Defaults to the current default context.

    Returns
    ----
    A dictionary containing
    
    if attributes is not None
    - `image` (ndarray): (H, W, C) float32 rendered image corresponding to the input attributes

    if return_depth is True
    - `depth` (ndarray): (H, W) float32 camera space linear depth, ranging from 0 to 1.
    
    if return_point_id is True
    - `point_id` (ndarray): (H, W) int32 point ID map
    """
    if ctx is None:
        ctx = RastContext.get_default_context()

    height, width = size
    assert points.ndim == 2 and points.shape[1] == 3 and points.dtype == np.float32, f"Points should be a float32 array with shape (N, 3), but got {points.shape} {points.dtype}"

    if attributes is not None:
        assert attributes.dtype == np.float32
        assert (attributes.shape[:-1] == points.shape[:-1]), f"Attribute shape {attributes.shape} does not match point shape {points.shape}"
        assert attributes.shape[-1] in [1, 2, 3, 4], f'Vertex attribute only supports channels 1, 2, 3, 4, but got {attributes.shape[-1]}'

    assert intrinsics is None or (intrinsics.shape == (3, 3) and intrinsics.dtype == np.float32), f"Intrinsics should be a 3x3 float32 matrix, but got {intrinsics.shape} {intrinsics.dtype}"
    assert extrinsics is None or (extrinsics.shape == (4, 4) and extrinsics.dtype == np.float32), f"Extrinsics should be a 4x4 float32 matrix, but got {extrinsics.shape} {extrinsics.dtype}"
    assert view is None or (view.shape == (4, 4) and view.dtype == np.float32), f"View should be a 4x4 float32 matrix, but got {view.shape} {view.dtype}"
    assert projection is None or (projection.shape == (4, 4) and projection.dtype == np.float32), f"Projection should be a 4x4 float32 matrix, but got {projection.shape} {projection.dtype}"

    background_image, background_depth, background_mask, background_point_id = [background.get(k, None) if background is not None else None for k in ['image', 'depth', 'mask', 'point_id']]
    if background_image is not None:
        assert background_image.ndim == 3 and background_image.shape == (height, width, attributes.shape[-1]), f"Image should be a float32 array with shape (H, W, {attributes.shape[1]}), but got {background_image.shape} {background_image.dtype}"
    if background_depth is not None:
        assert background_depth.dtype == np.float32 and background_depth.ndim == 2 and background_depth.shape == (height, width), f"Depth should be a float32 array with shape (H, W), but got {background_depth.shape} {background_depth.dtype}"
    if background_mask is not None:
        assert background_mask.dtype == np.bool_ and background_mask.ndim == 2 and background_mask.shape == (height, width), f"Mask should be a bool array with shape (H, W), but got {background_mask.shape} {background_mask.dtype}"
    if background_point_id is not None:
        assert background_point_id.dtype == np.int32 and background_point_id.ndim == 2 and background_point_id.shape == (height, width), f"Point ID should be a int32 array with shape (H, W), but got {background_point_id.shape} {background_point_id.dtype}"

    if not isinstance(point_sizes, ndarray):
        point_sizes = np.full((points.shape[0],), point_sizes, dtype=np.float32)

    if attributes is not None:
        num_channels = attributes.shape[-1]
        attributes = np.concatenate([attributes, np.zeros((*attributes.shape[:-1], 4 - num_channels,), dtype=attributes.dtype)], axis=-1) if num_channels < 4 else attributes
    if background_image is not None:
        background_image = np.concatenate([background_image, np.ones((height, width, 4 - background_image.shape[-1]), dtype=background_image.dtype)], axis=-1) if background_image.shape[-1] < 4 else background_image
    
    if extrinsics is not None:
        assert view is None, "Cannot specify both extrinsics and view."
        view = extrinsics_to_view(extrinsics)
    if intrinsics is not None:
        assert projection is None, "Cannot specify both intrinsics and projection."
        projection = intrinsics_to_perspective(intrinsics, near, far)
    if view is None:
        view = np.eye(4, dtype=np.float32)
    if projection is None:
        projection = np.eye(4, dtype=np.float32) 

    # Get program
    program_name = f'rasterize_point_cloud(return_point_id={return_point_id}, point_shape={point_shape})'
    num_point_shape_vertices = {'triangle': 3, 'square': 4, 'pentagon': 5, 'hexagon': 6, 'circle': 4}.get(point_shape)
    if (prog := ctx.programs.get(program_name, None)) is None:
        prog = ctx.mgl_ctx.program(
            vertex_shader=RASTERIZE_POINT_CLOUD_VS.replace('{return_point_id}', '' if return_point_id else '//'),
            geometry_shader=RASTERIZE_POINT_CLOUD_GS\
                .replace('{return_point_id}', '' if return_point_id else '//')\
                .replace('{point_shape}', point_shape)\
                .replace('{num_point_shape_vertices}', str(num_point_shape_vertices)),
            fragment_shader=RASTERIZE_POINT_CLOUD_FS\
                .replace('{return_point_id}', '' if return_point_id else '//')\
                .replace('{circle_begin}', '' if point_shape == 'circle' else '/*')\
                .replace('{circle_end}', '' if point_shape == 'circle' else '*/')
        )
        ctx.programs[program_name] = prog

    # Create buffers
    vbo_vertices = ctx.mgl_ctx.buffer(np.ascontiguousarray(points, dtype='f4'))
    if attributes is not None:
        vbo_attributes = ctx.mgl_ctx.buffer(np.ascontiguousarray(attributes, dtype='f4'))
    vbo_point_sizes = ctx.mgl_ctx.buffer(np.ascontiguousarray(point_sizes, dtype='f4'))

    vao = ctx.mgl_ctx.vertex_array(
        prog,
        list(filter(lambda x: x is not None, [
            (vbo_vertices, '3f', 'inVert'),
            (vbo_attributes, f'4f', 'inAttr') if attributes is not None else None,
            (vbo_point_sizes, '1f', 'inPointSize') if point_sizes is not None else None,
        ]))
    )

    # Create textures
    image_tex = ctx.mgl_ctx.texture((width, height), 4, dtype='f4', data=np.ascontiguousarray(background_image[::-1, :, :]) if background_image is not None else None)
    mask_tex = ctx.mgl_ctx.texture((width, height), 1, dtype='f1', data=np.ascontiguousarray(background_mask[::-1, :]) if background_mask is not None else None)
    buffer_depth_tex = ctx.mgl_ctx.depth_texture((width, height))
    if return_depth:
        linear_depth_tex = ctx.mgl_ctx.texture((width, height), 1, dtype='f4', data=np.ascontiguousarray(background_depth[::-1, :]) if background_depth is not None else None)

    if return_point_id:
        point_id_tex = ctx.mgl_ctx.texture((width, height), 1, dtype='i4', data=np.ascontiguousarray(background_point_id[::-1, :]) if background_point_id is not None else None)

    if background_image is None:
        clear_texture(ctx, image_tex, value=(0.0, 0.0, 0.0, 1.0))
    if background_mask is None:
        clear_texture(ctx, mask_tex, value=0.0)
    if background_depth is None:
        clear_texture(ctx, buffer_depth_tex, value=1.0)
    else:
        depth_texture_linear_to_buffer(ctx, linear_depth_tex, buffer_depth_tex)
    if return_point_id:
        clear_texture(ctx, point_id_tex, value=(-1,))

    # Create framebuffer
    if return_point_id:
        fbo = ctx.mgl_ctx.framebuffer(
            color_attachments=[image_tex, mask_tex, point_id_tex],
            depth_attachment=buffer_depth_tex,
        )
    else:
        fbo = ctx.mgl_ctx.framebuffer(
            color_attachments=[image_tex, mask_tex],
            depth_attachment=buffer_depth_tex,
        )
    fbo.viewport = (0, 0, width, height)

    # Set uniforms
    prog['uViewMat'].write(np.ascontiguousarray(view.transpose().astype('f4')))
    prog['uProjectionMat'].write(np.ascontiguousarray(projection.transpose().astype('f4')))
    prog['uScreenSize'].value = (width, height)
    prog['uIsPointSize3D'].value = point_size_in == '3d'

    # Set render states
    ctx.mgl_ctx.depth_func = '<'
    ctx.mgl_ctx.enable(ctx.mgl_ctx.DEPTH_TEST)
    ctx.mgl_ctx.disable(ctx.mgl_ctx.CULL_FACE)
    ctx.mgl_ctx.disable(ctx.mgl_ctx.BLEND)
    
    # Render
    fbo.use()
    vao.render(mode=moderngl.POINTS)

    # Read
    mask = np.frombuffer(mask_tex.read(), dtype=bool).reshape((height, width))
    mask = np.flip(mask, axis=0)
    if attributes is not None:
        image = np.frombuffer(image_tex.read(), dtype='f4').reshape((height, width, 4))
        image = np.flip(image, axis=0)
        image = image[:, :, :num_channels]
    else:
        image = None
    if return_depth:
        depth_texture_buffer_to_linear(ctx, buffer_depth_tex, linear_depth_tex)
        depth = np.frombuffer(linear_depth_tex.read(), dtype='f4').reshape((height, width))
        depth = np.flip(depth, axis=0)
    else:
        depth = None
    if return_point_id:
        point_id = np.frombuffer(point_id_tex.read(), dtype='i4').reshape((height, width))
        point_id = np.flip(point_id, axis=0)
    else:
        point_id = None

    # Release
    vao.release()
    vbo_vertices.release()
    if attributes is not None:
        vbo_attributes.release()
    vbo_point_sizes.release()
    fbo.release()
    image_tex.release()
    mask_tex.release()
    buffer_depth_tex.release()
    if background_depth is not None or return_depth:
        linear_depth_tex.release()
    if return_point_id:
        point_id_tex.release()

    output = {
        "image": image,
        "depth": depth,
        "mask": mask,
        "point_id": point_id
    }

    output = {k: v for k, v in output.items() if v is not None}

    return output


def sample_texture(
    uv_map: ndarray,
    texture_map: ndarray,
    interpolation: Literal['linear', 'nearest'] = 'linear',
    mipmap_level: Union[int, Tuple[int, int]] = 0,
    repeat: Union[bool, Tuple[bool, bool]] = False,
    anisotropic: float = 1.0,
    ctx: Optional[RastContext] = None
) -> ndarray:
    """
    Sample from a texture map with a UV map.
    """
    if ctx is None:
        ctx = RastContext.get_default_context()

    assert len(texture_map.shape) == 3 and 1 <= texture_map.shape[2] <= 4
    assert uv_map.shape[2] == 2
    height, width = uv_map.shape[:2]
    texture_dtype = map_np_dtype(texture_map.dtype)

    # Create texture
    texture_tex = ctx.mgl_ctx.texture((texture_map.shape[1], texture_map.shape[0]), texture_map.shape[2], dtype=texture_dtype, data=np.ascontiguousarray(texture_map))
    if isinstance(mipmap_level, tuple):
        base_mipmap_level, max_mipmap_level = mipmap_level
    else:
        base_mipmap_level, max_mipmap_level = 0, mipmap_level
    if base_mipmap_level > 0 or max_mipmap_level > 0:
        use_mipmap = True
        texture_tex.build_mipmaps(base_mipmap_level, max_mipmap_level)
    else:
        use_mipmap = False
    if interpolation == 'linear':
        texture_tex.filter = (moderngl.LINEAR_MIPMAP_LINEAR if use_mipmap else moderngl.LINEAR, moderngl.LINEAR)
    elif interpolation == 'nearest':
        texture_tex.filter = (moderngl.NEAREST_MIPMAP_NEAREST if use_mipmap else moderngl.NEAREST, moderngl.NEAREST)
    if isinstance(repeat, tuple):
        texture_tex.repeat_x, texture_tex.repeat_y = repeat
    else:
        texture_tex.repeat_x = texture_tex.repeat_y = repeat
    texture_tex.anisotropy = anisotropic

    uv_tex = ctx.mgl_ctx.texture((width, height), 2, dtype='f4', data=np.ascontiguousarray(uv_map.astype('f4', copy=False)))
    uv_tex.filter = (moderngl.NEAREST, moderngl.NEAREST)

    # Create render buffer and frame buffer
    output_tex = ctx.mgl_ctx.texture((uv_map.shape[1], uv_map.shape[0]), texture_map.shape[2], dtype=texture_dtype)

    # Get program
    program_name = 'sample_texture'
    if (prog := ctx.programs.get(program_name, None)) is None:
        prog = ctx.mgl_ctx.program(
            vertex_shader=FULL_SCREEN_VS,
            fragment_shader=SAMPLE_TEXTURE_FS
        )
        prog['texMap'] = 0
        prog['uvMap'] = 1
        ctx.programs[program_name] = prog

    # Render
    run_full_screen_program(ctx, prog, [texture_tex, uv_tex], output_tex)

    # Read buffer
    image_buffer = np.frombuffer(output_tex.read(), dtype=texture_dtype).reshape((height, width, texture_tex.shape[2]))

    # Release
    texture_tex.release()
    output_tex.release()
    uv_tex.release()

    return image_buffer


# def warp_image_by_depth(
#     ctx: RastContext,
#     src_depth: ndarray,
#     src_image: ndarray = None,
#     width: int = None,
#     height: int = None,
#     *,
#     extrinsics_src: ndarray = None,
#     extrinsics_tgt: ndarray = None,
#     intrinsics_src: ndarray = None,
#     intrinsics_tgt: ndarray = None,
#     near: float = 0.1,
#     far: float = 100.0,
#     cull_backface: bool = True,
#     ssaa: int = 1,
#     return_depth: bool = False,
# ) -> Tuple[ndarray, ...]:
#     """
#     Warp image by depth map.

#     ## Parameters
#         ctx (RastContext): rasterizer context
#         src_depth (ndarray): [H, W]
#         src_image (ndarray, optional): [H, W, C]. The image to warp. Defaults to None (use uv coordinates).
#         width (int, optional): width of the output image. None to use depth map width. Defaults to None.
#         height (int, optional): height of the output image. None to use depth map height. Defaults to None.
#         extrinsics_src (ndarray, optional): extrinsics matrix of the source camera. Defaults to None (identity).
#         extrinsics_tgt (ndarray, optional): extrinsics matrix of the target camera. Defaults to None (identity).
#         intrinsics_src (ndarray, optional): intrinsics matrix of the source camera. Defaults to None (use the same as intrinsics_tgt).
#         intrinsics_tgt (ndarray, optional): intrinsics matrix of the target camera. Defaults to None (use the same as intrinsics_src).
#         cull_backface (bool, optional): whether to cull backface. Defaults to True.
#         ssaa (int, optional): super sampling anti-aliasing. Defaults to 1.
    
#     ## Returns
#         tgt_image (ndarray): [H, W, C] warped image (or uv coordinates if image is None).
#         tgt_depth (ndarray): [H, W] screen space depth, ranging from 0 to 1. If return_depth is False, it is None.
#     """
#     assert src_depth.ndim == 2

#     if width is None:
#         width = src_depth.shape[1]
#     if height is None:
#         height = src_depth.shape[0]
#     if src_image is not None:
#         assert src_image.shape[-2:] == src_depth.shape[-2:], f'Shape of source image {src_image.shape} does not match shape of source depth {src_depth.shape}'

#     # set up default camera parameters
#     extrinsics_src = np.eye(4) if extrinsics_src is None else extrinsics_src
#     extrinsics_tgt = np.eye(4) if extrinsics_tgt is None else extrinsics_tgt
#     intrinsics_src = intrinsics_tgt if intrinsics_src is None else intrinsics_src
#     intrinsics_tgt = intrinsics_src if intrinsics_tgt is None else intrinsics_tgt
    
#     assert all(x is not None for x in [extrinsics_src, extrinsics_tgt, intrinsics_src, intrinsics_tgt]), "Make sure you have provided all the necessary camera parameters."

#     # check shapes
#     assert extrinsics_src.shape == (4, 4) and extrinsics_tgt.shape == (4, 4)
#     assert intrinsics_src.shape == (3, 3) and intrinsics_tgt.shape == (3, 3) 

#     # convert to view and perspective matrices
#     view_tgt = transforms.extrinsics_to_view(extrinsics_tgt)
#     perspective_tgt = transforms.intrinsics_to_perspective(intrinsics_tgt, near=near, far=far)

#     # unproject depth map
#     uv, faces = utils.image_mesh(*src_depth.shape[-2:])
#     pts = transforms.unproject_cv(uv, src_depth.reshape(-1), extrinsics_src, intrinsics_src)
#     faces = mesh.triangulate_mesh(faces, vertices=pts)

#     # rasterize attributes
#     if src_image is not None:
#         attr = src_image.reshape(-1, src_image.shape[-1])
#     else:
#         attr = uv

#     tgt_image, tgt_depth = rasterize_triangles(
#         ctx,
#         width * ssaa,
#         height * ssaa,
#         vertices=pts,
#         faces=faces,
#         attr=attr,
#         transform=perspective_tgt @ view_tgt,
#         cull_backface=cull_backface,
#         return_depth=return_depth,
#     )

#     if ssaa > 1:
#         tgt_image = tgt_image.reshape(height, ssaa, width, ssaa, -1).mean(axis=(1, 3))
#         tgt_depth = tgt_depth.reshape(height, ssaa, width, ssaa, -1).mean(axis=(1, 3)) if return_depth else None

#     return tgt_image, tgt_depth



def test_rasterization(ctx: Optional[RastContext] = None):
    """
    Test if rasterization works. It will render a cube with random colors and save it as a CHECKME.png file.
    """

    from .mesh import create_cube_mesh
    from .transforms import perspective_from_fov, view_look_at
        
    vertices, faces = create_cube_mesh(tri=True)
    attributes = np.random.rand(len(vertices), 3).astype(np.float32)
    projection = perspective_from_fov(fov_x=np.deg2rad(60), aspect_ratio=1, near=1e-8, far=100000)
    view = view_look_at([1, 2, 2], [0, 0, 0], [0, 1, 0])
    out = rasterize_triangles(
        (512, 512), 
        vertices=vertices, 
        attributes=attributes, 
        faces=faces, 
        view=view, 
        projection=projection,
        cull_backface=True,
        return_depth=True,
        return_interpolation=False,
        ctx=ctx,
    )
    image, mask = out['image'], out['mask']
    image = np.concatenate([image, mask[:, :, None].astype(np.float32)], axis=-1)
    import cv2
    cv2.imwrite('CHECKME.png', cv2.cvtColor((image.clip(0, 1) * 255).astype(np.uint8), cv2.COLOR_RGBA2BGRA))
    print("An image has been saved as ./CHECKME.png")


# def test_rasterization(ctx: RastContext):
#     """
#     Test if rasterization works. It will render a cube with random colors and save it as a CHECKME.png file.
#     """
#     from .mesh import cube
#     from .transforms import perspective_from_fov, view_look_at, project, unproject

#     vertices, faces = cube(tri=True)
#     attributes = np.random.rand(len(vertices), 3).astype(np.float32)
#     projection = perspective_from_fov(fov_x=np.deg2rad(60), aspect_ratio=1, near=0.1, far=100000)
#     view = view_look_at([1, 2, 2], [0, 0, 0], [0, 1, 0])
#     for _ in range(20):
#         with timeit():
#             out = rasterize_point_cloud(
#                 ctx, 
#                 512, 512, 
#                 points=vertices, 
#                 point_sizes=1,
#                 point_size_in='3d',
#                 point_shape='circle',
#                 attributes=attributes, 
#                 view=view, 
#                 projection=projection,
#                 return_depth=True,
#                 return_point_id=True,
#             )
#     image = out['image']
#     print(np.unique(image))
#     image = np.concatenate([image, np.zeros((*image.shape[:2], 4 - image.shape[2]), dtype=image.dtype)], axis=-1)
#     import cv2
#     cv2.imwrite('CHECKME.png', cv2.cvtColor((image.clip(0, 1) * 255).astype(np.uint8), cv2.COLOR_RGB2BGR))
#     print("An image has been saved as ./CHECKME.png")