How can I convert Tensor into Bitmap on PyTorch Mobile?

Question

I found that solution (https://itnext.io/converting-pytorch-float-tensor-to-android-rgba-bitmap-with-kotlin-ffd4602a16b6) but when I tried to convert that way I found that the size of inputTensor.dataAsFloatArray is more than bitmap.width*bitmap.height. How works converting tensor to float array or is there any other possible method to convert pytorch tensor to bitmap?

val inputTensor = TensorImageUtils.bitmapToFloat32Tensor(
    bitmap,
    TensorImageUtils.TORCHVISION_NORM_MEAN_RGB, TensorImageUtils.TORCHVISION_NORM_STD_RGB
)

// Float array size is 196608 when width and height are 256x256 = 65536

val res = floatArrayToGrayscaleBitmap(inputTensor.dataAsFloatArray, bitmap.width, bitmap.height)


fun floatArrayToGrayscaleBitmap (
    floatArray: FloatArray,
    width: Int,
    height: Int,
    alpha :Byte = (255).toByte(),
    reverseScale :Boolean = false
) : Bitmap {

    // Create empty bitmap in RGBA format (even though it says ARGB but channels are RGBA)
    val bmp = Bitmap.createBitmap(width, height, Bitmap.Config.ARGB_8888)
    val byteBuffer = ByteBuffer.allocate(width*height*4)
    Log.d("App", floatArray.size.toString() + " " + (width * height * 4).toString())

    // mapping smallest value to 0 and largest value to 255
    val maxValue = floatArray.max() ?: 1.0f
    val minValue = floatArray.min() ?: 0.0f
    val delta = maxValue-minValue
    var tempValue :Byte

    // Define if float min..max will be mapped to 0..255 or 255..0
    val conversion = when(reverseScale) {
        false -> { v: Float -> ((v-minValue)/delta*255).toByte() }
        true -> { v: Float -> (255-(v-minValue)/delta*255).toByte() }
    }

    // copy each value from float array to RGB channels and set alpha channel
    floatArray.forEachIndexed { i, value ->
        tempValue = conversion(value)
        byteBuffer.put(4*i, tempValue)
        byteBuffer.put(4*i+1, tempValue)
        byteBuffer.put(4*i+2, tempValue)
        byteBuffer.put(4*i+3, alpha)
    }

    bmp.copyPixelsFromBuffer(byteBuffer)

    return bmp
}

Answer 1

None of the answers were able to produce the output I wanted, so this is what I came up with - it is basically only reverse engineered version of what happenes in TensorImageUtils.bitmapToFloat32Tensor().

Please note that this function only works if you are using MemoryFormat.CONTIGUOUS (which is default) in TensorImageUtils.bitmapToFloat32Tensor().

fun tensor2Bitmap(input: FloatArray, width: Int, height: Int, normMeanRGB: FloatArray, normStdRGB: FloatArray): Bitmap? {
    val pixelsCount = height * width
    val pixels = IntArray(pixelsCount)
    val output = Bitmap.createBitmap(width, height, Bitmap.Config.ARGB_8888)
    
    val conversion = { v: Float -> ((v.coerceIn(0.0f, 1.0f))*255.0f).roundToInt()}

    val offset_g = pixelsCount
    val offset_b = 2 * pixelsCount
    for (i in 0 until pixelsCount) {
        val r = conversion(input[i] * normStdRGB[0] + normMeanRGB[0])
        val g = conversion(input[i + offset_g] * normStdRGB[1] + normMeanRGB[1])
        val b = conversion(input[i + offset_b] * normStdRGB[2] + normMeanRGB[2])
        pixels[i] = 255 shl 24 or (r.toInt() and 0xff shl 16) or (g.toInt() and 0xff shl 8) or (b.toInt() and 0xff)
    }
    output.setPixels(pixels, 0, width, 0, 0, width, height)
    return output
}

Example usage then could be as follows:

tensor2Bitmap(outputTensor.dataAsFloatArray, bitmap.width, bitmap.height, TensorImageUtils.TORCHVISION_NORM_MEAN_RGB, TensorImageUtils.TORCHVISION_NORM_STD_RGB)

Answer 2

// I faced the same problem, and I found the function itself

TensorImageUtils.bitmapToFloat32Tensor()

tortures the RGB colorspace. You should try to convert yuv to a bitmap and use

TensorImageUtils.bitmapToFloat32Tensor

instead for NOW.

// I modified the code from phillies (up) to get the coloful bitmap. Note that the format of an output tensor is typically NCHW.

// Here's my function in Kotlin. Hopefully it works in your case:

private fun floatArrayToBitmap(floatArray: FloatArray, width: Int, height: Int) : Bitmap {

        // Create empty bitmap in ARGB format
        val bmp: Bitmap = Bitmap.createBitmap(width, height, Bitmap.Config.ARGB_8888)
        val pixels = IntArray(width * height * 4)

        // mapping smallest value to 0 and largest value to 255
        val maxValue = floatArray.max() ?: 1.0f
        val minValue = floatArray.min() ?: -1.0f
        val delta = maxValue-minValue

        // Define if float min..max will be mapped to 0..255 or 255..0
        val conversion = { v: Float -> ((v-minValue)/delta*255.0f).roundToInt()}

        // copy each value from float array to RGB channels
        for (i in 0 until width * height) {
            val r = conversion(floatArray[i])
            val g = conversion(floatArray[i+width*height])
            val b = conversion(floatArray[i+2*width*height])
            pixels[i] = rgb(r, g, b) // you might need to import for rgb()
        }
        bmp.setPixels(pixels, 0, width, 0, 0, width, height)

        return bmp
    }

Hopefully future releases of PyTorch Mobile will fix this bug.