Custom Glide ModelLoader for Yuv Image [YUV_420_888]

YuvImageDataFetcher.kt

import android.graphics.ImageFormat
import android.graphics.Rect
import android.graphics.YuvImage
import android.media.Image
import com.bumptech.glide.Priority
import com.bumptech.glide.load.DataSource
import com.bumptech.glide.load.data.DataFetcher
import java.io.ByteArrayOutputStream
import java.nio.ByteBuffer

class YuvImageDataFetcher(private val image: Image) : DataFetcher<ByteBuffer> {

    private val outputStream = ByteArrayOutputStream()

    override fun loadData(
        priority: Priority,
        callback: DataFetcher.DataCallback<in ByteBuffer>
    ) {
        try {
            val yBuffer = image.planes[0].buffer // Y
            val uBuffer = image.planes[1].buffer // U
            val vBuffer = image.planes[2].buffer // V

            val ySize = yBuffer.remaining()
            val uSize = uBuffer.remaining()
            val vSize = vBuffer.remaining()

            val nv21 = ByteArray(ySize + uSize + vSize)

            //U and V are swapped
            yBuffer.get(nv21, 0, ySize)
            vBuffer.get(nv21, ySize, vSize)
            uBuffer.get(nv21, ySize + vSize, uSize)

            val yuvImage = YuvImage(nv21, ImageFormat.NV21, image.width, image.height, null)
            yuvImage.compressToJpeg(Rect(0, 0, yuvImage.width, yuvImage.height), 100, outputStream)
            val byteBuffer = ByteBuffer.wrap(outputStream.toByteArray())

            callback.onDataReady(byteBuffer)
        } catch (error: Exception) {
            callback.onLoadFailed(error)
        }
    }

    override fun cleanup() {
        outputStream.close()
    }

    override fun cancel() {
        // Do nothing
    }

    override fun getDataClass(): Class<ByteBuffer> =
        ByteBuffer::class.java

    override fun getDataSource(): DataSource =
        DataSource.LOCAL
}

YuvImageModelLoader.kt

import android.graphics.ImageFormat
import android.media.Image
import com.bumptech.glide.load.Options
import com.bumptech.glide.load.model.ModelLoader
import com.bumptech.glide.load.model.ModelLoader.LoadData
import com.bumptech.glide.signature.ObjectKey
import java.nio.ByteBuffer

class YuvImageModelLoader : ModelLoader<Image, ByteBuffer> {

    override fun buildLoadData(
        model: Image,
        width: Int,
        height: Int,
        options: Options
    ): LoadData<ByteBuffer>? =
        LoadData(ObjectKey(model), YuvImageDataFetcher(model))

    override fun handles(model: Image): Boolean =
        model.format == ImageFormat.YUV_420_888
}

YuvImageModelLoaderFactory.kt

import android.media.Image
import com.bumptech.glide.load.model.ModelLoader
import com.bumptech.glide.load.model.ModelLoaderFactory
import com.bumptech.glide.load.model.MultiModelLoaderFactory
import java.nio.ByteBuffer

class YuvImageModelLoaderFactory : ModelLoaderFactory<Image, ByteBuffer> {

    override fun build(multiFactory: MultiModelLoaderFactory): ModelLoader<Image, ByteBuffer> =
        YuvImageModelLoader()

    override fun teardown() {
        // Do nothing
    }
}