CIE Lab to CIE XYZ - Mojo

cie_lab.py

from algorithm import parallelize
from benchmark import Benchmark
from math import min
from math.limit import inf
from python import Python
from tensor import Tensor, TensorSpec, TensorShape
from random import rand
from time import now
from utils.index import Index

var colour: PythonObject = None
var np: PythonObject = None


fn load_python_packages():
    try:
        colour = Python.import_module("colour")
    except e:
        print(e)

    try:
        np = Python.import_module("numpy")
    except e:
        print(e)


fn XYZ_to_Lab_python(XYZ: PythonObject) -> PythonObject:
    var Lab: PythonObject = None
    try:
        Lab = colour.XYZ_to_Lab(XYZ)
    except e:
        print(e)

    return Lab


fn intermediate_lightness_function_CIE1976_naive(Y: Float64, Y_n: Float64) -> Float64:
    let Y_Y_n: Float64 = Y / Y_n

    var Y_i: Float64 = Y_Y_n ** (1 / 3)
    if Y_Y_n <= (24 / 116) ** 3:
        Y_i = (841 / 108) * Y_Y_n + 16 / 116

    return Y_i


fn XYZ_to_Lab_mojo_naive(
    XYZ: Tensor[DType.float64], XYZ_n: Tensor[DType.float64]
) -> Tensor[DType.float64]:
    let shape: TensorShape = XYZ.shape()
    let height: Int = shape[0]
    let width: Int = shape[1]
    let channels: Int = shape[1]

    var Lab = Tensor[DType.float64](TensorSpec(DType.float64, height, width, channels))

    for y in range(height):
        for x in range(width):
            let f_X_X_n: Float64 = intermediate_lightness_function_CIE1976_naive(
                XYZ[y, x, 0] / XYZ_n[0], XYZ_n[0]
            )
            let f_Y_Y_n: Float64 = intermediate_lightness_function_CIE1976_naive(
                XYZ[y, x, 1] / XYZ_n[1], XYZ_n[1]
            )
            let f_Z_Z_n: Float64 = intermediate_lightness_function_CIE1976_naive(
                XYZ[y, x, 2] / XYZ_n[2], XYZ_n[2]
            )

            let L: Float64 = 116 * f_Y_Y_n - 16
            let a: Float64 = 500 * (f_X_X_n - f_Y_Y_n)
            let b: Float64 = 200 * (f_Y_Y_n - f_Z_Z_n)

            Lab[Index(y, x, 0)] = L
            Lab[Index(y, x, 1)] = a
            Lab[Index(y, x, 2)] = b

    return Lab


alias SIMD_WIDTH = simdwidthof[DType.float64]()


fn intermediate_lightness_function_CIE1976_optimised[
    SIMD_WIDTH: Int
](Y: SIMD[DType.float64, SIMD_WIDTH], Y_n: SIMD[DType.float64, SIMD_WIDTH]) -> SIMD[
    DType.float64, SIMD_WIDTH
]:
    let Y_Y_n: SIMD[DType.float64, SIMD_WIDTH] = Y / Y_n

    var Y_i: SIMD[DType.float64, SIMD_WIDTH] = Y_Y_n ** (1 / 3)
    if Y_Y_n <= (24 / 116) ** 3:
        Y_i = (841 / 108) * Y_Y_n + 16 / 116

    return Y_i


fn XYZ_to_Lab_mojo_optimised_1(
    XYZ: Tensor[DType.float64], XYZ_n: Tensor[DType.float64]
) -> Tensor[DType.float64]:
    let shape: TensorShape = XYZ.shape()
    let height: Int = shape[0]
    let width: Int = shape[1]
    let channels: Int = shape[1]

    var Lab = Tensor[DType.float64](TensorSpec(DType.float64, height, width, channels))

    for y in range(height):
        for x in range(width):
            let f_X_X_n: SIMD[
                DType.float64, SIMD_WIDTH
            ] = intermediate_lightness_function_CIE1976_optimised(
                XYZ[y, x, 0] / XYZ_n[0], XYZ_n[0]
            )
            let f_Y_Y_n: SIMD[
                DType.float64, SIMD_WIDTH
            ] = intermediate_lightness_function_CIE1976_optimised(
                XYZ[y, x, 1] / XYZ_n[1], XYZ_n[1]
            )
            let f_Z_Z_n: SIMD[
                DType.float64, SIMD_WIDTH
            ] = intermediate_lightness_function_CIE1976_optimised(
                XYZ[y, x, 2] / XYZ_n[2], XYZ_n[2]
            )

            let L: SIMD[DType.float64, SIMD_WIDTH] = 116 * f_Y_Y_n - 16
            let a: SIMD[DType.float64, SIMD_WIDTH] = 500 * (f_X_X_n - f_Y_Y_n)
            let b: SIMD[DType.float64, SIMD_WIDTH] = 200 * (f_Y_Y_n - f_Z_Z_n)

            Lab.simd_store(y * x + 0, L)
            Lab.simd_store(y * x + 1, a)
            Lab.simd_store(y * x + 2, b)

    return Lab


fn XYZ_to_Lab_mojo_optimised_2(
    XYZ: Tensor[DType.float64], XYZ_n: Tensor[DType.float64]
) -> Tensor[DType.float64]:
    let shape: TensorShape = XYZ.shape()
    let height: Int = shape[0]
    let width: Int = shape[1]
    let channels: Int = shape[1]

    var Lab = Tensor[DType.float64](TensorSpec(DType.float64, height, width, channels))

    @parameter
    fn compute_height(y: Int):
        for x in range(width):
            let f_X_X_n: SIMD[
                DType.float64, SIMD_WIDTH
            ] = intermediate_lightness_function_CIE1976_optimised(
                XYZ[y, x, 0] / XYZ_n[0], XYZ_n[0]
            )
            let f_Y_Y_n: SIMD[
                DType.float64, SIMD_WIDTH
            ] = intermediate_lightness_function_CIE1976_optimised(
                XYZ[y, x, 1] / XYZ_n[1], XYZ_n[1]
            )
            let f_Z_Z_n: SIMD[
                DType.float64, SIMD_WIDTH
            ] = intermediate_lightness_function_CIE1976_optimised(
                XYZ[y, x, 2] / XYZ_n[2], XYZ_n[2]
            )

            let L: SIMD[DType.float64, SIMD_WIDTH] = 116 * f_Y_Y_n - 16
            let a: SIMD[DType.float64, SIMD_WIDTH] = 500 * (f_X_X_n - f_Y_Y_n)
            let b: SIMD[DType.float64, SIMD_WIDTH] = 200 * (f_Y_Y_n - f_Z_Z_n)

            Lab.simd_store(y * x + 0, L)
            Lab.simd_store(y * x + 1, a)
            Lab.simd_store(y * x + 2, b)

    parallelize[compute_height](height, 128)

    return Lab


fn main():
    load_python_packages()

    let width: Int = 1920
    let height: Int = 1080
    let channels: Int = 3

    let XYZ_mojo = rand[DType.float64](width, height, channels)
    var illuminant_mojo = Tensor[DType.float64](TensorSpec(DType.float64, 3))
    illuminant_mojo[0] = 95.04 / 100.0
    illuminant_mojo[1] = 100.0 / 100.0
    illuminant_mojo[2] = 108.88 / 100.0

    var best_mojo_naive: Float64 = inf[DType.float64]()
    var best_mojo_optimised: Float64 = inf[DType.float64]()
    var best_python: Float64 = inf[DType.float64]()

    var i = 0
    while True:
        if i == 20:
            break

        var XYZ_python: PythonObject = None
        try:
            XYZ_python = np.random.random([width, height, channels])
        except e:
            print(e)

        var time_then: Float64 = now()
        XYZ_to_Lab_mojo_naive(XYZ_mojo, illuminant_mojo)
        var time_now: Float64 = now()
        best_mojo_naive = min(time_now - time_then, best_mojo_naive)

        time_then = now()
        XYZ_to_Lab_mojo_optimised_1(XYZ_mojo, illuminant_mojo)
        time_now = now()
        best_mojo_optimised = min(time_now - time_then, best_mojo_optimised)

        time_then = now()
        XYZ_to_Lab_python(XYZ_python)
        time_now = now()
        best_python = min(time_now - time_then, best_python)
        i += 1

    print("Mojo Naive     : ", best_mojo_naive * 1e-6, "ms")
    print("Mojo Optimised : ", best_mojo_optimised * 1e-6, "ms")
    print("Python         : ", best_python * 1e-6, "ms")