Auto White Balance Adjustment

WhiteBalanceAdjuster.kt

import android.graphics.Bitmap
import android.graphics.Color
import androidx.annotation.ColorInt
import com.curiouscreature.kotlin.math.Float3
import com.curiouscreature.kotlin.math.Mat3
import com.curiouscreature.kotlin.math.saturate
import com.curiouscreature.kotlin.math.transpose
import kotlin.math.pow


class WhiteBalanceAdjuster {
    private val CIECAT02_to_XYZ = transpose(Mat3.of(
        1.0961200f,  0.4543690f, -0.0096276f,
        -0.2788690f,  0.4735330f, -0.0056980f,
        0.1827450f,  0.0720978f,  1.0153300f
    ))

    private val ILLUMINANT_D65_LMS = Float3(0.949237f, 1.03542f, 1.08728f)

    private val XYZ_to_CIECAT02 = transpose(Mat3.of(
        0.7328000f, -0.7036000f,  0.0030000f,
        0.4296000f,  1.6975000f,  0.0136000f,
        -0.1624000f,  0.0061000f,  0.9834000f
    ))

    private val sRGB_to_XYZ = transpose(Mat3.of(
        0.4124560f,  0.2126730f,  0.0193339f,
        0.3575760f,  0.7151520f,  0.1191920f,
        0.1804380f,  0.0721750f,  0.9503040f
    ))

    private val sRGB_to_LMS = XYZ_to_CIECAT02 * sRGB_to_XYZ

    private val XYZ_to_sRGB = transpose(Mat3.of(
        3.2404542f, -0.9692660f,  0.0556434f,
        -1.5371385f,  1.8760108f, -0.2040259f,
        -0.4985314f,  0.0415560f,  1.0572252f
    ))
    private val LMS_to_sRGB = XYZ_to_sRGB * CIECAT02_to_XYZ

    private fun EOCF_sRGB(inputVector: Float3): Float3 {
        val a = 0.055f
        val a1 = 1.055F
        val b = 1.0f / 12.92f
        val p = 2.4f
        return inputVector.transform { input ->
            if (input <= 0.04045f) {
                input * b
            } else {
                ((input + a) / a1).pow(p)
            }
        }
    }

    private fun OECF_SRGB(inputVector: Float3): Float3 {
        val a = 0.055f
        val a1 = 1.055F
        val b = 12.92f
        val p = 1.0f / 2.4f
        return inputVector.transform { input ->
            if (input <= 0.0031308f) {
                input * b
            } else {
                a1 * input.pow(p) - a
            }
        }
    }

    private fun Int.toFloat3(): Float3 {
        val red = Color.red(this).toFloat() / 255.0f
        val green = Color.green(this).toFloat() / 255.0f
        val blue = Color.blue(this).toFloat() / 255.0f
        return Float3(red, green, blue)
    }

    @ColorInt
    private fun Float3.toColorInt() = Color.rgb(this.r, this.g, this.b)

    fun whiteBalance(bitmap: Bitmap, @ColorInt referenceWhite: Int): Bitmap {
        val mutableBitmap = bitmap.copy(bitmap.config, true)
        val pixelArray = IntArray(mutableBitmap.width * mutableBitmap.height)
        mutableBitmap.getPixels(
            pixelArray,
            0,
            mutableBitmap.width,
            0,
            0,
            mutableBitmap.width,
            mutableBitmap.height
        )

        val lmsReferenceWhite = sRGB_to_LMS * EOCF_sRGB(referenceWhite.toFloat3())
        val adaptation = ILLUMINANT_D65_LMS / lmsReferenceWhite
        val adaptationTransform = LMS_to_sRGB * Mat3.of(
            adaptation.x, 0.0f, 0.0f,
            0.0f, adaptation.y, 0.0f,
            0.0f, 0.0f, adaptation.z
        ) * sRGB_to_LMS

        for (y in 0 until bitmap.height) {
            for (x in 0 until bitmap.width) {
                val pixelArrayIndex = bitmap.width * y + x
                val pixel = pixelArray[pixelArrayIndex]
                val adaptedPixel = adaptationTransform * EOCF_sRGB(pixel.toFloat3())
                val newPixel = OECF_SRGB(adaptedPixel.transform { v -> saturate(v) }).toColorInt()
                pixelArray[pixelArrayIndex] = newPixel
            }
        }

        mutableBitmap.setPixels(
            pixelArray,
            0,
            mutableBitmap.width,
            0,
            0,
            mutableBitmap.width,
            mutableBitmap.height
        )
        return mutableBitmap
    }
}