PyTorch recently released with version 1.3.0 their first Android version which allows you to run model inference on your smartphone. The set of functions is not very complete yet and the implementation has a few breaking bugs, but it’s a great and helpful step for using DNNs on your device.

Photo by Chris Ried on Unsplash

One function that is dearly missing is the option to convert the output of a model such as a UNet back from a tensor to a bitmap to show it to the user. I implemented this function in Kotlin and you can find it as a gist here: https://gist.github.com/phillies/830f52b0bf592c32fc507669694f66d8

The function takes a float array and converts it into an RGBA bitmap with mapping the smallest float value to 0 and the largest float value to 255 or the other way round.

The function expects a floatArray as primary parameter, which can be obtained from a tensor via myTensor.dataAsFloatArray and should be a 2D tensor of shape [height, width] . Make sure your model puts out this shape. The standard PyTorch NCHW shape is fine as long as you make sure that N and C are 1. The alpha parameter sets the value for the alpha channel and reverseScale defines if the smallest float value is set to 0 ( false ) or 255 ( true ).

The resulting Bitmap is in RGBA format and can be shown easily using an ImageView for example. Hope this helps some of you :-)