frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 21:14:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

Add methods to write image tensor content to buffer (#27359)

Browse Source 
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/27359

Adding methods  to TensorImageUtils:
```
bitmapToFloatBuffer(..., FloatBuffer outBuffer, int outBufferOffset)
imageYUV420CenterCropToFloat32Tensor(..., FloatBuffer outBuffer, int outBufferOffset)
```
To be able to
 - reuse FloatBuffer for inference
 - to create batch-Tensor (contains several images/bitmaps)

As we reuse FloatBuffer for example demo app - image classification,
profiler shows less memory allocations (before that for every run we created new input tensor with newly allocated FloatBuffer) and ~-20ms on my PixelXL

Known open question:
At the moment every tensor element is written separatly calling `outBuffer.put()`, which is native call crossing lang boundaries
As an alternative - to allocation `float[]` on java side and fill it and put it in `outBuffer` with one call, reducing native calls, but increasing memory allocation on java side.
Tested locally just eyeballing durations - have not noticed big difference - decided to go with less memory allocations.

Will be good to merge into 1.3.0, but if not - demo app can use snapshot dependencies with this change.

PR with integration to demo app:
https://github.com/pytorch/android-demo-app/pull/6

Test Plan: Imported from OSS

Differential Revision: D17758621

Pulled By: IvanKobzarev

fbshipit-source-id: b4f1a068789279002d7ecc0bc680111f781bf980
This commit is contained in:
Ivan Kobzarev
2019-10-04 16:31:51 -07:00
committed by Facebook Github Bot
parent ac0f18437f
commit afbbe16f49
1 changed files with 111 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files

134
android/pytorch_android_torchvision/src/main/java/org/pytorch/torchvision/TensorImageUtils.java
View File

@ -7,6 +7,7 @@ import android.media.Image;
import org.pytorch.Tensor;
import java.nio.ByteBuffer;
import java.nio.FloatBuffer;
import java.util.Locale;
/**
@ -26,7 +27,7 @@ public final class TensorImageUtils {
* @param normStdRGB standard deviation for RGB channels normalization, length must equal 3, RGB order
*/
public static Tensor bitmapToFloat32Tensor(
final Bitmap bitmap, float[] normMeanRGB, float normStdRGB[]) {
final Bitmap bitmap, final float[] normMeanRGB, final float normStdRGB[]) {
checkNormMeanArg(normMeanRGB);
checkNormStdArg(normStdRGB);
@ -34,6 +35,52 @@ public final class TensorImageUtils {
bitmap, 0, 0, bitmap.getWidth(), bitmap.getHeight(), normMeanRGB, normStdRGB);
}
/**
* Writes tensor content from specified {@link android.graphics.Bitmap},
* normalized with specified in parameters mean and std to specified {@link java.nio.FloatBuffer}
* with specified offset.
*
* @param bitmap {@link android.graphics.Bitmap} as a source for Tensor data
* @param x - x coordinate of top left corner of bitmap's area
* @param y - y coordinate of top left corner of bitmap's area
* @param width - width of bitmap's area
* @param height - height of bitmap's area
* @param normMeanRGB means for RGB channels normalization, length must equal 3, RGB order
* @param normStdRGB standard deviation for RGB channels normalization, length must equal 3, RGB order
*/
public static void bitmapToFloatBuffer(
final Bitmap bitmap,
final int x,
final int y,
final int width,
final int height,
final float[] normMeanRGB,
final float[] normStdRGB,
final FloatBuffer outBuffer,
final int outBufferOffset) {
checkOutBufferCapacity(outBuffer, outBufferOffset, width, height);
checkNormMeanArg(normMeanRGB);
checkNormStdArg(normStdRGB);
final int pixelsCount = height * width;
final int[] pixels = new int[pixelsCount];
bitmap.getPixels(pixels, 0, width, x, y, width, height);
final int offset_g = pixelsCount;
final int offset_b = 2 * pixelsCount;
for (int i = 0; i < pixelsCount; i++) {
final int c = pixels[i];
float r = ((c >> 16) & 0xff) / 255.0f;
float g = ((c >> 8) & 0xff) / 255.0f;
float b = ((c) & 0xff) / 255.0f;
float rF = (r - normMeanRGB[0]) / normStdRGB[0];
float gF = (g - normMeanRGB[1]) / normStdRGB[1];
float bF = (b - normMeanRGB[2]) / normStdRGB[2];
outBuffer.put(outBufferOffset + i, rF);
outBuffer.put(outBufferOffset + offset_g + i, gF);
outBuffer.put(outBufferOffset + offset_b + i, bF);
}
}
/**
* Creates new {@link org.pytorch.Tensor} from specified area of {@link android.graphics.Bitmap},
* normalized with specified in parameters mean and std.
@ -57,22 +104,9 @@ public final class TensorImageUtils {
checkNormMeanArg(normMeanRGB);
checkNormStdArg(normStdRGB);
final int pixelsCount = height * width;
final int[] pixels = new int[pixelsCount];
bitmap.getPixels(pixels, 0, width, x, y, width, height);
final float[] floatArray = new float[3 * pixelsCount];
final int offset_g = pixelsCount;
final int offset_b = 2 * pixelsCount;
for (int i = 0; i < pixelsCount; i++) {
final int c = pixels[i];
float r = ((c >> 16) & 0xff) / 255.0f;
float g = ((c >> 8) & 0xff) / 255.0f;
float b = ((c) & 0xff) / 255.0f;
floatArray[i] = (r - normMeanRGB[0]) / normStdRGB[0];
floatArray[offset_g + i] = (g - normMeanRGB[1]) / normStdRGB[1];
floatArray[offset_b + i] = (b - normMeanRGB[2]) / normStdRGB[2];
}
return Tensor.newFloat32Tensor(new long[]{1, 3, height, width}, floatArray);
final FloatBuffer floatBuffer = Tensor.allocateFloatBuffer(3 * width * height);
bitmapToFloatBuffer(bitmap, x, y, width, height, normMeanRGB, normStdRGB, floatBuffer, 0);
return Tensor.newFloat32Tensor(new long[]{1, 3, height, width}, floatBuffer);
}
/**
@ -105,6 +139,52 @@ public final class TensorImageUtils {
checkRotateCWDegrees(rotateCWDegrees);
checkTensorSize(tensorWidth, tensorHeight);
final FloatBuffer floatBuffer = Tensor.allocateFloatBuffer(3 * tensorWidth * tensorHeight);
imageYUV420CenterCropToFloatBuffer(
image,
rotateCWDegrees,
tensorWidth,
tensorHeight,
normMeanRGB, normStdRGB, floatBuffer, 0);
return Tensor.newFloat32Tensor(new long[]{1, 3, tensorHeight, tensorWidth}, floatBuffer);
}
/**
* Writes tensor content from specified {@link android.media.Image}, doing optional rotation,
* scaling (nearest) and center cropping to specified {@link java.nio.FloatBuffer} with specified offset.
*
* @param image {@link android.media.Image} as a source for Tensor data
* @param rotateCWDegrees Clockwise angle through which the input image needs to be rotated to be
* upright. Range of valid values: 0, 90, 180, 270
* @param tensorWidth return tensor width, must be positive
* @param tensorHeight return tensor height, must be positive
* @param normMeanRGB means for RGB channels normalization, length must equal 3, RGB order
* @param normStdRGB standard deviation for RGB channels normalization, length must equal 3, RGB order
* @param outBuffer Output buffer, where tensor content will be written
* @param outBufferOffset Output buffer offset with which tensor content will be written
*/
public static void imageYUV420CenterCropToFloatBuffer(
final Image image,
int rotateCWDegrees,
final int tensorWidth,
final int tensorHeight,
float[] normMeanRGB,
float[] normStdRGB,
final FloatBuffer outBuffer,
final int outBufferOffset) {
checkOutBufferCapacity(outBuffer, outBufferOffset, tensorWidth, tensorHeight);
if (image.getFormat() != ImageFormat.YUV_420_888) {
throw new IllegalArgumentException(
String.format(
Locale.US, "Image format %d != ImageFormat.YUV_420_888", image.getFormat()));
}
checkNormMeanArg(normMeanRGB);
checkNormStdArg(normStdRGB);
checkRotateCWDegrees(rotateCWDegrees);
checkTensorSize(tensorWidth, tensorHeight);
final int widthBeforeRotation = image.getWidth();
final int heightBeforeRotation = image.getHeight();
@ -158,7 +238,6 @@ public final class TensorImageUtils {
final int channelSize = tensorHeight * tensorWidth;
final int tensorInputOffsetG = channelSize;
final int tensorInputOffsetB = 2 * channelSize;
final float[] floatArray = new float[3 * channelSize];
for (int x = 0; x < tensorWidth; x++) {
for (int y = 0; y < tensorHeight; y++) {
@ -198,13 +277,22 @@ public final class TensorImageUtils {
int r = clamp((a0 + a1) >> 10, 0, 255);
int g = clamp((a0 - a2 - a3) >> 10, 0, 255);
int b = clamp((a0 + a4) >> 10, 0, 255);
final int offset = y * tensorWidth + x;
floatArray[offset] = ((r / 255.f) - normMeanRGB[0]) / normStdRGB[0];
floatArray[tensorInputOffsetG + offset] = ((g / 255.f) - normMeanRGB[1]) / normStdRGB[1];
floatArray[tensorInputOffsetB + offset] = ((b / 255.f) - normMeanRGB[2]) / normStdRGB[2];
final int offset = outBufferOffset + y * tensorWidth + x;
float rF = ((r / 255.f) - normMeanRGB[0]) / normStdRGB[0];
float gF = ((g / 255.f) - normMeanRGB[1]) / normStdRGB[1];
float bF = ((b / 255.f) - normMeanRGB[2]) / normStdRGB[2];
outBuffer.put(offset, rF);
outBuffer.put(offset + tensorInputOffsetG, gF);
outBuffer.put(offset + tensorInputOffsetB, bF);
}
}
return Tensor.newFloat32Tensor(new long[]{1, 3, tensorHeight, tensorWidth}, floatArray);
}
private static void checkOutBufferCapacity(FloatBuffer outBuffer, int outBufferOffset, int tensorWidth, int tensorHeight) {
if (outBufferOffset + 3 * tensorWidth * tensorHeight > outBuffer.capacity()) {
throw new IllegalStateException("Buffer underflow");
}
}
private static void checkTensorSize(int tensorWidth, int tensorHeight) {