RsYuvToRgb.kt

class RsYuvToRgb(val rs: RenderScript) {

    private var rotateAllocation: Allocation? = null
    private var mOutputAllocation: Allocation? = null
    private var mInputAllocation: Allocation? = null

    private val rsYuvToRgb = ScriptIntrinsicYuvToRGB.create(rs, Element.U8_4(rs))

    private var yArray: ByteArray? = null
    private var uArray: ByteArray? = null
    private var vArray: ByteArray? = null

    private fun ByteBuffer.initByteArray(): ByteArray {
        return ByteArray(capacity())
    }

    fun yuvToRgb(yBuffer: ByteBuffer, uBuffer: ByteBuffer, vBuffer: ByteBuffer, width: Int, height: Int): ByteBuffer {
        val currentTimestamp = System.currentTimeMillis()

        if (mInputAllocation == null){
            val totalSize = yBuffer.capacity() + uBuffer.capacity() + vBuffer.capacity()
            val yuvType = Type.Builder(rs, Element.U8(rs)).apply {
                setX(totalSize)
            }
            mInputAllocation = Allocation.createTyped(rs, yuvType.create(), Allocation.USAGE_SCRIPT)
        }
        if (mOutputAllocation == null){
            val rgbType = Type.createXY(rs, Element.RGBA_8888(rs), width, height)
            mOutputAllocation = Allocation.createTyped(rs, rgbType, Allocation.USAGE_SCRIPT)
        }

        if (rotateAllocation == null){
            val rotatorType = Type.createXY(rs, Element.RGBA_8888(rs), height, width)
            rotateAllocation = Allocation.createTyped(rs, rotatorType, Allocation.USAGE_SCRIPT)
        } else {
            rotateAllocation!!.byteBuffer.clear()
//            rotateAllocation!!.byteBuffer.position(0)
        }

        if (yArray == null) {
            yArray = yBuffer.initByteArray()
        }
        if (uArray == null) {
            uArray = uBuffer.initByteArray()
        }
        if (vArray == null) {
            vArray = vBuffer.initByteArray()
        }

        yBuffer.get(yArray)
        uBuffer.get(uArray)
        vBuffer.get(vArray)

        mInputAllocation!!.copy1DRangeFrom(0, yArray!!.size, yArray)
        mInputAllocation!!.copy1DRangeFrom(yArray!!.size, uArray!!.size, uArray)
        mInputAllocation!!.copy1DRangeFrom(yArray!!.size + uArray!!.size, vArray!!.size, vArray)

//        mInputAllocation.copyFrom(yuvBuffer.array())
        rsYuvToRgb.setInput(mInputAllocation)
        rsYuvToRgb.forEach(mOutputAllocation)
//
        val lastAnalyzedTimestamp = System.currentTimeMillis()
        val perf = lastAnalyzedTimestamp - currentTimestamp
        println("PERFORMANCE : $perf ms")
        return mOutputAllocation!!.byteBuffer
    }
}