diff --git a/extension/android/CMakeLists.txt b/extension/android/CMakeLists.txt index c94b95286d0..8722444f519 100644 --- a/extension/android/CMakeLists.txt +++ b/extension/android/CMakeLists.txt @@ -161,6 +161,15 @@ if(EXECUTORCH_JNI_CUSTOM_LIBRARY) ) endif() +if(EXECUTORCH_BUILD_EXTENSION_IMAGE) + target_sources(executorch_jni PRIVATE jni/jni_layer_image.cpp) + # jnigraphics provides AndroidBitmap_lockPixels, used to read Bitmap storage. + list(APPEND link_libraries extension_image jnigraphics) + target_compile_definitions( + executorch_jni PUBLIC EXECUTORCH_BUILD_EXTENSION_IMAGE=1 + ) +endif() + if(EXECUTORCH_BUILD_EXTENSION_TRAINING) target_sources(executorch_jni PRIVATE jni/jni_layer_training.cpp) list(APPEND link_libraries extension_training) diff --git a/extension/android/executorch_android/src/androidTest/java/org/pytorch/executorch/ImageProcessorInstrumentationTest.kt b/extension/android/executorch_android/src/androidTest/java/org/pytorch/executorch/ImageProcessorInstrumentationTest.kt new file mode 100644 index 00000000000..71e0186330a --- /dev/null +++ b/extension/android/executorch_android/src/androidTest/java/org/pytorch/executorch/ImageProcessorInstrumentationTest.kt @@ -0,0 +1,306 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ +package org.pytorch.executorch + +import android.graphics.Bitmap +import android.graphics.Color +import androidx.test.ext.junit.runners.AndroidJUnit4 +import org.junit.Assert.assertArrayEquals +import org.junit.Assert.assertEquals +import org.junit.Assert.assertFalse +import org.junit.Assert.assertTrue +import org.junit.Assert.fail +import org.junit.Test +import org.junit.runner.RunWith +import org.pytorch.executorch.extension.image.ImageOrientation +import org.pytorch.executorch.extension.image.ImageProcessor +import org.pytorch.executorch.extension.image.ImageProcessorConfig +import org.pytorch.executorch.extension.image.LetterboxAnchor +import org.pytorch.executorch.extension.image.Normalization +import org.pytorch.executorch.extension.image.ResizeMode + +/** + * Instrumentation tests for [ImageProcessor]. + * + * The processor is pure image math, so these run without a .pte fixture: each test builds a bitmap + * in memory and asserts on the produced tensor. + */ +@RunWith(AndroidJUnit4::class) +class ImageProcessorInstrumentationTest { + + // ─── Config validation ────────────────────────────────────────────────────── + + @Test + fun testNonPositiveTargetSizeThrows() { + try { + ImageProcessorConfig(targetWidth = 0, targetHeight = 224) + fail("Should throw for a non-positive target width") + } catch (_: IllegalArgumentException) {} + } + + @Test + fun testZeroStandardDeviationThrows() { + try { + Normalization(1.0f, floatArrayOf(0f, 0f, 0f), floatArrayOf(1f, 0f, 1f)) + fail("Should throw for a zero standard deviation") + } catch (_: IllegalArgumentException) {} + } + + @Test + fun testWrongChannelCountThrows() { + try { + Normalization(1.0f, floatArrayOf(0f, 0f), floatArrayOf(1f, 1f, 1f)) + fail("Should throw for a mean with the wrong channel count") + } catch (_: IllegalArgumentException) {} + } + + @Test + fun testBuilderMatchesConstructor() { + val built = + ImageProcessorConfig.Builder() + .setTargetSize(320, 240) + .setResizeMode(ResizeMode.LETTERBOX) + .setPadValue(0.5f) + .build() + assertEquals(320, built.targetWidth) + assertEquals(240, built.targetHeight) + assertEquals(ResizeMode.LETTERBOX, built.resizeMode) + assertEquals(0.5f, built.padValue, 0.0f) + } + + // ─── Lifecycle ────────────────────────────────────────────────────────────── + + @Test + fun testCloseIsIdempotent() { + val processor = ImageProcessor(ImageProcessorConfig()) + assertTrue(processor.isValid) + processor.close() + assertFalse(processor.isValid) + processor.close() + assertFalse(processor.isValid) + } + + @Test + fun testUseAfterCloseThrows() { + val processor = ImageProcessor(ImageProcessorConfig()) + processor.close() + try { + processor.process(solidBitmap(4, 4, Color.RED)) + fail("Should throw after close") + } catch (_: IllegalStateException) {} + } + + // ─── Channel order and normalization ──────────────────────────────────────── + + @Test + fun testSolidRedKeepsChannelOrder() { + ImageProcessor(ImageProcessorConfig(targetWidth = 2, targetHeight = 2)).use { processor -> + val tensor = processor.process(solidBitmap(4, 4, Color.RED)) + assertArrayEquals(longArrayOf(1, 3, 2, 2), tensor.shape()) + val data = tensor.dataAsFloatArray + // CHW: the whole R plane is 1.0, G and B planes are 0.0. + for (i in 0 until 4) { + assertEquals("R[$i]", 1.0f, data[i], TOLERANCE) + assertEquals("G[$i]", 0.0f, data[4 + i], TOLERANCE) + assertEquals("B[$i]", 0.0f, data[8 + i], TOLERANCE) + } + } + } + + @Test + fun testImagenetNormalizationIsApplied() { + val config = + ImageProcessorConfig( + targetWidth = 1, + targetHeight = 1, + normalization = Normalization.imagenet(), + ) + ImageProcessor(config).use { processor -> + val data = processor.process(solidBitmap(4, 4, Color.WHITE)).dataAsFloatArray + assertEquals((1.0f - 0.485f) / 0.229f, data[0], TOLERANCE) + assertEquals((1.0f - 0.456f) / 0.224f, data[1], TOLERANCE) + assertEquals((1.0f - 0.406f) / 0.225f, data[2], TOLERANCE) + } + } + + // ─── Geometry ─────────────────────────────────────────────────────────────── + + @Test + fun testStretchOutputShapeIgnoresAspectRatio() { + ImageProcessor(ImageProcessorConfig(targetWidth = 224, targetHeight = 224)).use { processor -> + assertArrayEquals(longArrayOf(1, 3, 224, 224), processor.computeOutputShape(640, 480)) + assertEquals(0, processor.computeLetterboxPadding(640, 480).x) + assertEquals(0, processor.computeLetterboxPadding(640, 480).y) + } + } + + @Test + fun testLetterboxPadsTheShorterAxis() { + val config = + ImageProcessorConfig( + targetWidth = 100, + targetHeight = 100, + resizeMode = ResizeMode.LETTERBOX, + letterboxAnchor = LetterboxAnchor.CENTER, + ) + ImageProcessor(config).use { processor -> + // A 200x100 source scales to 100x50, leaving 25px above and below. + val padding = processor.computeLetterboxPadding(200, 100) + assertEquals(0, padding.x) + assertEquals(25, padding.y) + } + } + + @Test + fun testTopLeftAnchorHasNoPadding() { + val config = + ImageProcessorConfig( + targetWidth = 100, + targetHeight = 100, + resizeMode = ResizeMode.LETTERBOX, + letterboxAnchor = LetterboxAnchor.TOP_LEFT, + ) + ImageProcessor(config).use { processor -> + val padding = processor.computeLetterboxPadding(200, 100) + assertEquals(0, padding.x) + assertEquals(0, padding.y) + } + } + + @Test + fun testLetterboxFillsPaddingWithPadValue() { + val config = + ImageProcessorConfig( + targetWidth = 4, + targetHeight = 4, + resizeMode = ResizeMode.LETTERBOX, + padValue = -1.0f, + ) + ImageProcessor(config).use { processor -> + // An 8x4 source scales to 4x2, so rows 0 and 3 are padding. + val data = processor.process(solidBitmap(8, 4, Color.WHITE)).dataAsFloatArray + for (channel in 0 until 3) { + val plane = channel * 16 + for (col in 0 until 4) { + assertEquals("top pad", -1.0f, data[plane + col], TOLERANCE) + assertEquals("bottom pad", -1.0f, data[plane + 12 + col], TOLERANCE) + } + } + } + } + + // ─── Orientation ──────────────────────────────────────────────────────────── + + @Test + fun testOrientationMovesLetterboxPaddingAxis() { + ImageProcessor(ImageProcessorConfig()).use { processor -> + // The output is always the target size, whatever the source orientation. + assertArrayEquals( + longArrayOf(1, 3, 224, 224), + processor.computeOutputShape(640, 480, ImageOrientation.RIGHT), + ) + } + val letterbox = + ImageProcessorConfig( + targetWidth = 100, + targetHeight = 100, + resizeMode = ResizeMode.LETTERBOX, + ) + ImageProcessor(letterbox).use { processor -> + // Upright, a 200x100 source pads vertically; rotated 90 degrees it is + // 100x200, so the padding moves to the horizontal axis. + assertEquals(25, processor.computeLetterboxPadding(200, 100, ImageOrientation.UP).y) + assertEquals(25, processor.computeLetterboxPadding(200, 100, ImageOrientation.RIGHT).x) + } + } + + @Test + fun testRightOrientationRotatesClockwise() { + val config = ImageProcessorConfig(targetWidth = 1, targetHeight = 2) + ImageProcessor(config).use { processor -> + // Source is red on the left, blue on the right. Rotating 90 degrees + // clockwise puts red on top. + val bitmap = Bitmap.createBitmap(2, 1, Bitmap.Config.ARGB_8888) + bitmap.setPixel(0, 0, Color.RED) + bitmap.setPixel(1, 0, Color.BLUE) + val data = processor.process(bitmap, ImageOrientation.RIGHT).dataAsFloatArray + // CHW over a 1x2 output: R plane is data[0..1], B plane is data[4..5]. + assertTrue("red should land on top, got R=${data[0]}", data[0] > 0.5f) + assertTrue("red should not be at the bottom, got R=${data[1]}", data[1] < 0.5f) + assertTrue("blue should land at the bottom, got B=${data[5]}", data[5] > 0.5f) + } + } + + @Test + fun testUnsupportedOrientationCodeIsRejected() { + // Mirrored EXIF codes are not supported; the enum only exposes rotations, so + // this guards the native validation against a future enum addition. + assertEquals(4, ImageOrientation.values().size) + assertArrayEquals( + intArrayOf(1, 3, 6, 8), + ImageOrientation.values().map { it.exifCode }.toIntArray(), + ) + } + + // ─── Reuse ────────────────────────────────────────────────────────────────── + + @Test + fun testProcessIntoMatchesProcess() { + ImageProcessor(ImageProcessorConfig(targetWidth = 8, targetHeight = 8)).use { processor -> + val bitmap = gradientBitmap(16, 16) + val allocated = processor.process(bitmap) + val reused = Tensor.fromBlob(Tensor.allocateFloatBuffer(3 * 8 * 8), longArrayOf(1, 3, 8, 8)) + processor.processInto(bitmap, reused) + assertArrayEquals(allocated.dataAsFloatArray, reused.dataAsFloatArray, TOLERANCE) + } + } + + @Test + fun testProcessIntoRejectsWrongShape() { + ImageProcessor(ImageProcessorConfig(targetWidth = 8, targetHeight = 8)).use { processor -> + val wrong = Tensor.fromBlob(Tensor.allocateFloatBuffer(3 * 4 * 4), longArrayOf(1, 3, 4, 4)) + try { + processor.processInto(solidBitmap(16, 16, Color.RED), wrong) + fail("Should throw for a tensor with the wrong shape") + } catch (_: IllegalArgumentException) {} + } + } + + @Test + fun testProcessIntoRejectsWrongDtype() { + ImageProcessor(ImageProcessorConfig(targetWidth = 8, targetHeight = 8)).use { processor -> + val wrong = Tensor.fromBlob(IntArray(3 * 8 * 8), longArrayOf(1, 3, 8, 8)) + try { + processor.processInto(solidBitmap(16, 16, Color.RED), wrong) + fail("Should throw for a tensor with the wrong dtype") + } catch (_: IllegalArgumentException) {} + } + } + + // ─── Helpers ──────────────────────────────────────────────────────────────── + + private fun solidBitmap(width: Int, height: Int, color: Int): Bitmap { + val bitmap = Bitmap.createBitmap(width, height, Bitmap.Config.ARGB_8888) + bitmap.eraseColor(color) + return bitmap + } + + private fun gradientBitmap(width: Int, height: Int): Bitmap { + val bitmap = Bitmap.createBitmap(width, height, Bitmap.Config.ARGB_8888) + for (y in 0 until height) { + for (x in 0 until width) { + bitmap.setPixel(x, y, Color.rgb(x * 255 / width, y * 255 / height, 128)) + } + } + return bitmap + } + + private companion object { + const val TOLERANCE = 1e-3f + } +} diff --git a/extension/android/executorch_android/src/main/java/org/pytorch/executorch/extension/image/ImageProcessor.kt b/extension/android/executorch_android/src/main/java/org/pytorch/executorch/extension/image/ImageProcessor.kt new file mode 100644 index 00000000000..7b69dbe3fa2 --- /dev/null +++ b/extension/android/executorch_android/src/main/java/org/pytorch/executorch/extension/image/ImageProcessor.kt @@ -0,0 +1,351 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +package org.pytorch.executorch.extension.image + +import android.graphics.Bitmap +import java.io.Closeable +import java.nio.ByteBuffer +import java.nio.FloatBuffer +import java.util.concurrent.atomic.AtomicLong +import org.pytorch.executorch.DType +import org.pytorch.executorch.ExecutorchRuntimeException +import org.pytorch.executorch.Tensor +import org.pytorch.executorch.annotations.Experimental + +/** + * Converts camera frames and bitmaps into the normalized `[1, 3, H, W]` float tensors vision models + * expect, replacing the hand-rolled resize/normalize loops apps otherwise write per model. + * + * The pipeline rotates the source upright according to the supplied [ImageOrientation], resizes it + * per [ImageProcessorConfig.resizeMode], and normalizes it per + * [ImageProcessorConfig.normalization]. The resize and normalize steps use NEON where available. + * + * Warning: These APIs are experimental and subject to change without notice + * + * Thread-safety: an ImageProcessor is NOT thread-safe. Internal scratch buffers are reused across + * calls, so concurrent calls on one instance are unsafe. Use one instance per thread; separate + * instances are independent. + * + * @param config The target size, resize mode, and normalization to apply. + */ +@Experimental +class ImageProcessor(val config: ImageProcessorConfig) : Closeable { + + private val nativeHandle = AtomicLong(0L) + + init { + val handle = + nativeCreate( + config.targetWidth, + config.targetHeight, + config.resizeMode.ordinal, + config.letterboxAnchor.ordinal, + config.padValue, + config.normalization.scaleFactor, + config.normalization.mean, + config.normalization.standardDeviation, + ) + if (handle == 0L) { + throw ExecutorchRuntimeException( + ExecutorchRuntimeException.INTERNAL, + "Failed to create native ImageProcessor", + ) + } + nativeHandle.set(handle) + } + + companion object { + init { + System.loadLibrary("executorch") + } + + private const val OUTPUT_CHANNELS = 3 + + @JvmStatic + private external fun nativeCreate( + targetWidth: Int, + targetHeight: Int, + resizeMode: Int, + letterboxAnchor: Int, + padValue: Float, + scaleFactor: Float, + mean: FloatArray, + standardDeviation: FloatArray, + ): Long + + @JvmStatic private external fun nativeDestroy(nativeHandle: Long) + + @JvmStatic + private external fun nativeProcessBitmap( + nativeHandle: Long, + bitmap: Bitmap, + orientationCode: Int, + outBuffer: FloatBuffer, + outCapacity: Int, + ) + + @JvmStatic + private external fun nativeProcessYuv( + nativeHandle: Long, + yPlane: ByteBuffer, + yStride: Int, + yCapacity: Int, + uvPlane: ByteBuffer, + uvStride: Int, + uvCapacity: Int, + width: Int, + height: Int, + format: Int, + range: Int, + orientationCode: Int, + outBuffer: FloatBuffer, + outCapacity: Int, + ) + + @JvmStatic + private external fun nativeComputeOutputShape( + nativeHandle: Long, + inputWidth: Int, + inputHeight: Int, + orientationCode: Int, + ): IntArray + + @JvmStatic + private external fun nativeComputeLetterboxPadding( + nativeHandle: Long, + inputWidth: Int, + inputHeight: Int, + orientationCode: Int, + ): IntArray + } + + /** Check if the native handle is valid (not yet closed). */ + val isValid: Boolean + get() = nativeHandle.get() != 0L + + /** Releases native resources. Call this when done with the processor. */ + override fun close() { + val handle = nativeHandle.getAndSet(0L) + if (handle != 0L) { + nativeDestroy(handle) + } + } + + /** + * Process an ARGB_8888 bitmap into a normalized float tensor. + * + * @param bitmap The source bitmap. Must be [Bitmap.Config.ARGB_8888]. + * @param orientation The EXIF orientation of the bitmap's contents. + * @return A float Tensor shaped `[1, 3, targetHeight, targetWidth]`. + * @throws IllegalStateException if the processor has been closed + * @throws ExecutorchRuntimeException if processing fails + */ + @JvmOverloads + fun process(bitmap: Bitmap, orientation: ImageOrientation = ImageOrientation.UP): Tensor { + val handle = requireHandle() + val buffer = Tensor.allocateFloatBuffer(outputNumel()) + nativeProcessBitmap(handle, bitmap, orientation.exifCode, buffer, buffer.capacity()) + return Tensor.fromBlob(buffer, outputShape()) + } + + /** + * Process an ARGB_8888 bitmap into a caller-provided tensor, reusing its storage. + * + * Avoids the per-call allocation of [process], which matters for sustained video. `tensor` must + * be a float Tensor shaped `[1, 3, targetHeight, targetWidth]`; its storage is overwritten, so + * the caller must finish using the previous contents before calling again. + * + * @param bitmap The source bitmap. Must be [Bitmap.Config.ARGB_8888]. + * @param tensor The output tensor to fill. + * @param orientation The EXIF orientation of the bitmap's contents. + * @throws IllegalArgumentException if `tensor` has the wrong dtype or shape + * @throws IllegalStateException if the processor has been closed + * @throws ExecutorchRuntimeException if processing fails + */ + @JvmOverloads + fun processInto( + bitmap: Bitmap, + tensor: Tensor, + orientation: ImageOrientation = ImageOrientation.UP, + ) { + val handle = requireHandle() + val buffer = outputBufferOf(tensor) + nativeProcessBitmap(handle, bitmap, orientation.exifCode, buffer, buffer.capacity()) + } + + /** + * Process semi-planar YUV (NV12/NV21) camera planes into a normalized float tensor. + * + * For a CameraX `ImageProxy` in `YUV_420_888`, pass `planes[0].buffer` as [yPlane] and the + * interleaved chroma plane as [uvPlane]: `planes[1].buffer` for [YuvFormat.NV12], or + * `planes[2].buffer` for [YuvFormat.NV21]. Both buffers must be direct. + * + * Only semi-planar chroma is supported. Check `planes[1].pixelStride == 2` before calling; + * a stride of 1 means fully planar I420, which this path cannot consume. + * + * Both buffers are bounds-checked against the strides and dimensions given here. The decode + * reads the interleaved chroma plane through `uvStride * (height / 2 - 1) + width`, so a plane + * buffer that a camera HAL trimmed below that is rejected rather than read past its end. + * + * @param yPlane Direct buffer holding the luma plane. + * @param yStride Row stride of the luma plane, in bytes. + * @param uvPlane Direct buffer holding the interleaved chroma plane. + * @param uvStride Row stride of the chroma plane, in bytes. + * @param width Source width in pixels. + * @param height Source height in pixels. + * @param format Chroma order of the interleaved plane. + * @param orientation The EXIF orientation of the frame's contents. + * @param range Quantization range of the samples. + * @return A float Tensor shaped `[1, 3, targetHeight, targetWidth]`. + * @throws IllegalStateException if the processor has been closed + * @throws ExecutorchRuntimeException if processing fails + */ + @JvmOverloads + fun processYuv( + yPlane: ByteBuffer, + yStride: Int, + uvPlane: ByteBuffer, + uvStride: Int, + width: Int, + height: Int, + format: YuvFormat, + orientation: ImageOrientation = ImageOrientation.UP, + range: YuvRange = YuvRange.VIDEO, + ): Tensor { + val handle = requireHandle() + val buffer = Tensor.allocateFloatBuffer(outputNumel()) + nativeProcessYuv( + handle, + yPlane, + yStride, + yPlane.capacity(), + uvPlane, + uvStride, + uvPlane.capacity(), + width, + height, + format.ordinal, + range.ordinal, + orientation.exifCode, + buffer, + buffer.capacity(), + ) + return Tensor.fromBlob(buffer, outputShape()) + } + + /** + * Process semi-planar YUV camera planes into a caller-provided tensor, reusing its storage. + * + * See [processYuv] for the plane contract and [processInto] for the reuse contract. + * + * @throws IllegalArgumentException if `tensor` has the wrong dtype or shape + */ + @JvmOverloads + fun processYuvInto( + yPlane: ByteBuffer, + yStride: Int, + uvPlane: ByteBuffer, + uvStride: Int, + width: Int, + height: Int, + format: YuvFormat, + tensor: Tensor, + orientation: ImageOrientation = ImageOrientation.UP, + range: YuvRange = YuvRange.VIDEO, + ) { + val handle = requireHandle() + val buffer = outputBufferOf(tensor) + nativeProcessYuv( + handle, + yPlane, + yStride, + yPlane.capacity(), + uvPlane, + uvStride, + uvPlane.capacity(), + width, + height, + format.ordinal, + range.ordinal, + orientation.exifCode, + buffer, + buffer.capacity(), + ) + } + + /** + * Shape of the tensor this processor produces for the given source. + * + * @param inputWidth Source width in pixels. + * @param inputHeight Source height in pixels. + * @param orientation The EXIF orientation of the source. + */ + @JvmOverloads + fun computeOutputShape( + inputWidth: Int, + inputHeight: Int, + orientation: ImageOrientation = ImageOrientation.UP, + ): LongArray { + val handle = requireHandle() + return nativeComputeOutputShape(handle, inputWidth, inputHeight, orientation.exifCode).map { + it.toLong() + } + .toLongArray() + } + + /** + * Letterbox padding (per side, in pixels) applied for the given source, letting callers map model + * output back to source coordinates without replicating the resize geometry. Returns `(0, 0)` for + * [ResizeMode.STRETCH] or [LetterboxAnchor.TOP_LEFT]. + * + * @param inputWidth Source width in pixels. + * @param inputHeight Source height in pixels. + * @param orientation The EXIF orientation of the source. + */ + @JvmOverloads + fun computeLetterboxPadding( + inputWidth: Int, + inputHeight: Int, + orientation: ImageOrientation = ImageOrientation.UP, + ): LetterboxPadding { + val handle = requireHandle() + val padding = + nativeComputeLetterboxPadding(handle, inputWidth, inputHeight, orientation.exifCode) + return LetterboxPadding(padding[0], padding[1]) + } + + private fun requireHandle(): Long { + val handle = nativeHandle.get() + check(handle != 0L) { "ImageProcessor has been closed" } + return handle + } + + private fun outputNumel(): Int = OUTPUT_CHANNELS * config.targetHeight * config.targetWidth + + private fun outputShape(): LongArray = + longArrayOf( + 1, + OUTPUT_CHANNELS.toLong(), + config.targetHeight.toLong(), + config.targetWidth.toLong(), + ) + + private fun outputBufferOf(tensor: Tensor): FloatBuffer { + require(tensor.dtype() == DType.FLOAT) { "Output tensor must be float, got ${tensor.dtype()}" } + require(tensor.shape().contentEquals(outputShape())) { + "Output tensor must be shaped ${outputShape().contentToString()}, got " + + tensor.shape().contentToString() + } + val buffer = tensor.getRawDataBuffer() + require(buffer is FloatBuffer && buffer.isDirect) { + "Output tensor must be backed by a direct FloatBuffer" + } + return buffer + } +} diff --git a/extension/android/executorch_android/src/main/java/org/pytorch/executorch/extension/image/ImageProcessorConfig.kt b/extension/android/executorch_android/src/main/java/org/pytorch/executorch/extension/image/ImageProcessorConfig.kt new file mode 100644 index 00000000000..27100d10e06 --- /dev/null +++ b/extension/android/executorch_android/src/main/java/org/pytorch/executorch/extension/image/ImageProcessorConfig.kt @@ -0,0 +1,175 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +package org.pytorch.executorch.extension.image + +import org.pytorch.executorch.annotations.Experimental + +/** How the source image is fitted to the target dimensions. */ +@Experimental +enum class ResizeMode { + /** Scale to the target dimensions directly, ignoring aspect ratio. */ + STRETCH, + /** Scale to fit inside the target dimensions, then pad with `padValue`. */ + LETTERBOX, +} + +/** Where letterboxed content sits within the padded canvas. */ +@Experimental +enum class LetterboxAnchor { + CENTER, + TOP_LEFT, +} + +/** + * EXIF orientation of the source image. The pipeline rotates the content upright before resizing. + * Only the four rotation codes are supported (no mirrored variants). + * + * @property exifCode The EXIF orientation code this entry represents. + */ +@Experimental +enum class ImageOrientation(val exifCode: Int) { + UP(1), + DOWN(3), + RIGHT(6), + LEFT(8), +} + +/** Chroma layout of semi-planar YUV input. */ +@Experimental +enum class YuvFormat { + NV12, + NV21, +} + +/** + * Quantization range of YUV samples. VIDEO is studio/limited range (Y in [16, 235], chroma in + * [16, 240]); FULL spans [0, 255]. Decoding with the wrong range shifts contrast and color. + */ +@Experimental +enum class YuvRange { + VIDEO, + FULL, +} + +/** + * Per-channel RGB normalization, applied as + * `(pixel * scaleFactor - mean[c]) / standardDeviation[c]`. + * + * @property scaleFactor Scale applied to the raw 0-255 sample before mean subtraction. + * @property mean Per-channel mean, exactly 3 entries (R, G, B). + * @property standardDeviation Per-channel standard deviation, exactly 3 nonzero entries (R, G, B). + */ +@Experimental +class Normalization( + val scaleFactor: Float, + mean: FloatArray, + standardDeviation: FloatArray, +) { + val mean: FloatArray = mean.copyOf() + val standardDeviation: FloatArray = standardDeviation.copyOf() + + init { + require(mean.size == CHANNELS) { "mean must have $CHANNELS entries, got ${mean.size}" } + require(standardDeviation.size == CHANNELS) { + "standardDeviation must have $CHANNELS entries, got ${standardDeviation.size}" + } + require(standardDeviation.all { it != 0.0f }) { "standardDeviation entries must be nonzero" } + } + + companion object { + private const val CHANNELS = 3 + + /** Maps 0-255 samples to [0, 1] with no mean subtraction. */ + @JvmStatic + fun zeroToOne(): Normalization = + Normalization(1.0f / 255.0f, floatArrayOf(0.0f, 0.0f, 0.0f), floatArrayOf(1.0f, 1.0f, 1.0f)) + + /** The standard ImageNet mean and standard deviation over [0, 1] samples. */ + @JvmStatic + fun imagenet(): Normalization = + Normalization( + 1.0f / 255.0f, + floatArrayOf(0.485f, 0.456f, 0.406f), + floatArrayOf(0.229f, 0.224f, 0.225f), + ) + } +} + +/** + * Configuration for [ImageProcessor]. + * + * Warning: These APIs are experimental and subject to change without notice + * + * @property targetWidth Width of the produced tensor, in pixels. + * @property targetHeight Height of the produced tensor, in pixels. + * @property resizeMode How the source is fitted to the target dimensions. + * @property letterboxAnchor Where letterboxed content sits; ignored for [ResizeMode.STRETCH]. + * @property padValue Value written to letterbox padding, in normalized output units. + * @property normalization Per-channel normalization applied to the output. + */ +@Experimental +data class ImageProcessorConfig( + val targetWidth: Int = 224, + val targetHeight: Int = 224, + val resizeMode: ResizeMode = ResizeMode.STRETCH, + val letterboxAnchor: LetterboxAnchor = LetterboxAnchor.CENTER, + val padValue: Float = 0.0f, + val normalization: Normalization = Normalization.zeroToOne(), +) { + init { + require(targetWidth > 0) { "targetWidth must be positive" } + require(targetHeight > 0) { "targetHeight must be positive" } + } + + /** Builder class for ImageProcessorConfig for Java interoperability. */ + class Builder { + private var targetWidth: Int = 224 + private var targetHeight: Int = 224 + private var resizeMode: ResizeMode = ResizeMode.STRETCH + private var letterboxAnchor: LetterboxAnchor = LetterboxAnchor.CENTER + private var padValue: Float = 0.0f + private var normalization: Normalization = Normalization.zeroToOne() + + fun setTargetSize(width: Int, height: Int) = apply { + require(width > 0 && height > 0) { "Target dimensions must be positive" } + this.targetWidth = width + this.targetHeight = height + } + + fun setResizeMode(resizeMode: ResizeMode) = apply { this.resizeMode = resizeMode } + + fun setLetterboxAnchor(letterboxAnchor: LetterboxAnchor) = apply { + this.letterboxAnchor = letterboxAnchor + } + + fun setPadValue(padValue: Float) = apply { this.padValue = padValue } + + fun setNormalization(normalization: Normalization) = apply { + this.normalization = normalization + } + + fun build() = + ImageProcessorConfig( + targetWidth = targetWidth, + targetHeight = targetHeight, + resizeMode = resizeMode, + letterboxAnchor = letterboxAnchor, + padValue = padValue, + normalization = normalization, + ) + } +} + +/** + * Per-side letterbox padding in pixels. + * + * @property x The left/right pad of the resized content. + * @property y The top/bottom pad of the resized content. + */ +@Experimental data class LetterboxPadding(val x: Int, val y: Int) diff --git a/extension/android/jni/jni_layer_image.cpp b/extension/android/jni/jni_layer_image.cpp new file mode 100644 index 00000000000..e3260c7bb67 --- /dev/null +++ b/extension/android/jni/jni_layer_image.cpp @@ -0,0 +1,439 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include + +#include + +#include +#include +#include + +#include +#include +#include + +#include + +namespace image = ::executorch::extension::image; +using ::executorch::extension::from_blob; +using ::executorch::extension::TensorPtr; +using ::executorch::jni_helper::setExecutorchPendingException; +using ::executorch::runtime::Error; + +namespace { + +image::ImageProcessor* toProcessor(JNIEnv* env, jlong handle) { + if (handle == 0) { + setExecutorchPendingException( + env, + static_cast(Error::InvalidState), + "ImageProcessor has been closed"); + return nullptr; + } + return reinterpret_cast(handle); +} + +bool toOrientation( + JNIEnv* env, + jint exifCode, + image::Orientation& orientation) { + orientation = static_cast(exifCode); + if (!image::is_supported_orientation(orientation)) { + setExecutorchPendingException( + env, + static_cast(Error::InvalidArgument), + "Unsupported EXIF orientation: " + std::to_string(exifCode)); + return false; + } + return true; +} + +// Wraps a caller-supplied direct FloatBuffer as the [1, 3, H, W] output tensor +// the processor writes into. `floatCapacity` is passed from Kotlin rather than +// read via GetDirectBufferCapacity, whose unit is unspecified for view buffers. +TensorPtr outputTensor( + JNIEnv* env, + jobject outBuffer, + jint floatCapacity, + const image::ImageProcessorConfig& config) { + void* data = env->GetDirectBufferAddress(outBuffer); + if (data == nullptr) { + setExecutorchPendingException( + env, + static_cast(Error::InvalidArgument), + "Output buffer must be a direct java.nio.FloatBuffer"); + return nullptr; + } + const int32_t required = image::ImageProcessorConfig::kOutputChannels * + config.target_height * config.target_width; + if (floatCapacity < required) { + setExecutorchPendingException( + env, + static_cast(Error::InvalidArgument), + "Output buffer holds " + std::to_string(floatCapacity) + + " floats, need " + std::to_string(required)); + return nullptr; + } + return from_blob( + static_cast(data), + {1, + image::ImageProcessorConfig::kOutputChannels, + config.target_height, + config.target_width}, + ::executorch::aten::ScalarType::Float); +} + +// The decoder takes raw pointers, so the plane bounds can only be enforced +// here. `capacity` comes from the Java buffer; GetDirectBufferAddress returns +// the region base, so capacity (not remaining) is the matching bound. +const uint8_t* directBytes( + JNIEnv* env, + jobject buffer, + jint capacity, + int64_t required, + const char* name) { + void* data = env->GetDirectBufferAddress(buffer); + if (data == nullptr) { + setExecutorchPendingException( + env, + static_cast(Error::InvalidArgument), + std::string(name) + " must be a direct java.nio.ByteBuffer"); + return nullptr; + } + if (capacity < required) { + setExecutorchPendingException( + env, + static_cast(Error::InvalidArgument), + std::string(name) + " holds " + std::to_string(capacity) + + " bytes, need " + std::to_string(required)); + return nullptr; + } + return static_cast(data); +} + +// Locks an ARGB_8888 bitmap for the lifetime of the scope. Android stores that +// format as RGBA bytes in memory, so it maps to ColorFormat::RGBA. +class ScopedBitmapLock { + public: + ScopedBitmapLock(JNIEnv* env, jobject bitmap) : env_(env), bitmap_(bitmap) { + if (AndroidBitmap_getInfo(env, bitmap, &info_) != + ANDROID_BITMAP_RESULT_SUCCESS) { + setExecutorchPendingException( + env, + static_cast(Error::InvalidArgument), + "Failed to read Bitmap info"); + return; + } + if (info_.format != ANDROID_BITMAP_FORMAT_RGBA_8888) { + setExecutorchPendingException( + env, + static_cast(Error::InvalidArgument), + "Bitmap must be ARGB_8888"); + return; + } + if (AndroidBitmap_lockPixels(env, bitmap, &pixels_) != + ANDROID_BITMAP_RESULT_SUCCESS) { + setExecutorchPendingException( + env, + static_cast(Error::AccessFailed), + "Failed to lock Bitmap pixels"); + pixels_ = nullptr; + return; + } + locked_ = true; + } + + ~ScopedBitmapLock() { + if (locked_) { + AndroidBitmap_unlockPixels(env_, bitmap_); + } + } + + ScopedBitmapLock(const ScopedBitmapLock&) = delete; + ScopedBitmapLock& operator=(const ScopedBitmapLock&) = delete; + + bool ok() const { + return locked_; + } + const uint8_t* pixels() const { + return static_cast(pixels_); + } + int32_t width() const { + return static_cast(info_.width); + } + int32_t height() const { + return static_cast(info_.height); + } + int32_t stride() const { + return static_cast(info_.stride); + } + + private: + JNIEnv* env_; + jobject bitmap_; + AndroidBitmapInfo info_{}; + void* pixels_ = nullptr; + bool locked_ = false; +}; + +} // namespace + +extern "C" { + +JNIEXPORT jlong JNICALL +Java_org_pytorch_executorch_extension_image_ImageProcessor_nativeCreate( + JNIEnv* env, + jclass /* clazz */, + jint targetWidth, + jint targetHeight, + jint resizeMode, + jint letterboxAnchor, + jfloat padValue, + jfloat scaleFactor, + jfloatArray mean, + jfloatArray stdDev) { + if (targetWidth <= 0 || targetHeight <= 0) { + setExecutorchPendingException( + env, + static_cast(Error::InvalidArgument), + "Target dimensions must be positive"); + return 0; + } + if (env->GetArrayLength(mean) != + image::ImageProcessorConfig::kOutputChannels || + env->GetArrayLength(stdDev) != + image::ImageProcessorConfig::kOutputChannels) { + setExecutorchPendingException( + env, + static_cast(Error::InvalidArgument), + "mean and standardDeviation must each have 3 entries"); + return 0; + } + + image::ImageProcessorConfig config; + config.target_width = targetWidth; + config.target_height = targetHeight; + config.resize_mode = static_cast(resizeMode); + config.letterbox_anchor = + static_cast(letterboxAnchor); + config.pad_value = padValue; + config.normalization.scale_factor = scaleFactor; + // The 4th mean/std slot is reserved for a future RGBA output; keep it an + // identity normalization so it stays divide-safe. + config.normalization.mean[3] = 0.0f; + config.normalization.std_dev[3] = 1.0f; + env->GetFloatArrayRegion( + mean, + 0, + image::ImageProcessorConfig::kOutputChannels, + config.normalization.mean); + env->GetFloatArrayRegion( + stdDev, + 0, + image::ImageProcessorConfig::kOutputChannels, + config.normalization.std_dev); + for (int32_t i = 0; i < image::ImageProcessorConfig::kOutputChannels; ++i) { + if (config.normalization.std_dev[i] == 0.0f) { + setExecutorchPendingException( + env, + static_cast(Error::InvalidArgument), + "standardDeviation entries must be nonzero"); + return 0; + } + } + // The portable implementation has no GPU path; keep the CPU sentinel so the + // config never reports a GPU decision that cannot happen here. + config.gpu_min_input_pixels = image::ImageProcessorConfig::kGpuNever; + + try { + return reinterpret_cast( + new image::ImageProcessor(std::move(config))); + } catch (const std::exception& e) { + ET_LOG(Error, "Failed to create ImageProcessor: %s", e.what()); + setExecutorchPendingException( + env, + static_cast(Error::Internal), + "Failed to create ImageProcessor: " + std::string(e.what())); + return 0; + } +} + +JNIEXPORT void JNICALL +Java_org_pytorch_executorch_extension_image_ImageProcessor_nativeDestroy( + JNIEnv* /* env */, + jclass /* clazz */, + jlong nativeHandle) { + delete reinterpret_cast(nativeHandle); +} + +JNIEXPORT void JNICALL +Java_org_pytorch_executorch_extension_image_ImageProcessor_nativeProcessBitmap( + JNIEnv* env, + jclass /* clazz */, + jlong nativeHandle, + jobject bitmap, + jint orientationCode, + jobject outBuffer, + jint outCapacity) { + auto* processor = toProcessor(env, nativeHandle); + if (processor == nullptr) { + return; + } + image::Orientation orientation; + if (!toOrientation(env, orientationCode, orientation)) { + return; + } + auto out = outputTensor(env, outBuffer, outCapacity, processor->config()); + if (out == nullptr) { + return; + } + + ScopedBitmapLock lock(env, bitmap); + if (!lock.ok()) { + return; + } + + const Error error = processor->process_into( + lock.pixels(), + lock.width(), + lock.height(), + lock.stride(), + image::ColorFormat::RGBA, + *out, + orientation); + if (error != Error::Ok) { + setExecutorchPendingException( + env, static_cast(error), "Failed to process Bitmap"); + } +} + +JNIEXPORT void JNICALL +Java_org_pytorch_executorch_extension_image_ImageProcessor_nativeProcessYuv( + JNIEnv* env, + jclass /* clazz */, + jlong nativeHandle, + jobject yPlane, + jint yStride, + jint yCapacity, + jobject uvPlane, + jint uvStride, + jint uvCapacity, + jint width, + jint height, + jint format, + jint range, + jint orientationCode, + jobject outBuffer, + jint outCapacity) { + auto* processor = toProcessor(env, nativeHandle); + if (processor == nullptr) { + return; + } + image::Orientation orientation; + if (!toOrientation(env, orientationCode, orientation)) { + return; + } + auto out = outputTensor(env, outBuffer, outCapacity, processor->config()); + if (out == nullptr) { + return; + } + // The last row of each plane is only read up to `width`, not a full stride. + // Leave the dimension and stride checks themselves to process_yuv_into; + // clamp to 0 here so a bad input cannot produce a negative bound. + const int64_t yRequired = width > 0 && height > 0 + ? static_cast(yStride) * (height - 1) + width + : 0; + const int64_t uvRequired = width > 0 && height > 0 + ? static_cast(uvStride) * (height / 2 - 1) + width + : 0; + const uint8_t* y = directBytes(env, yPlane, yCapacity, yRequired, "yPlane"); + if (y == nullptr) { + return; + } + const uint8_t* uv = + directBytes(env, uvPlane, uvCapacity, uvRequired, "uvPlane"); + if (uv == nullptr) { + return; + } + + const Error error = processor->process_yuv_into( + y, + yStride, + uv, + uvStride, + width, + height, + static_cast(format), + *out, + orientation, + image::kFullImage, + static_cast(range)); + if (error != Error::Ok) { + setExecutorchPendingException( + env, static_cast(error), "Failed to process YUV planes"); + } +} + +JNIEXPORT jintArray JNICALL +Java_org_pytorch_executorch_extension_image_ImageProcessor_nativeComputeOutputShape( + JNIEnv* env, + jclass /* clazz */, + jlong nativeHandle, + jint inputWidth, + jint inputHeight, + jint orientationCode) { + auto* processor = toProcessor(env, nativeHandle); + if (processor == nullptr) { + return nullptr; + } + image::Orientation orientation; + if (!toOrientation(env, orientationCode, orientation)) { + return nullptr; + } + + const auto shape = + processor->compute_output_shape(inputWidth, inputHeight, orientation); + jintArray result = env->NewIntArray(static_cast(shape.size())); + if (result == nullptr) { + return nullptr; + } + env->SetIntArrayRegion( + result, 0, static_cast(shape.size()), shape.data()); + return result; +} + +JNIEXPORT jintArray JNICALL +Java_org_pytorch_executorch_extension_image_ImageProcessor_nativeComputeLetterboxPadding( + JNIEnv* env, + jclass /* clazz */, + jlong nativeHandle, + jint inputWidth, + jint inputHeight, + jint orientationCode) { + auto* processor = toProcessor(env, nativeHandle); + if (processor == nullptr) { + return nullptr; + } + image::Orientation orientation; + if (!toOrientation(env, orientationCode, orientation)) { + return nullptr; + } + + const auto padding = processor->compute_letterbox_padding( + inputWidth, inputHeight, orientation); + const jint values[2] = {padding.first, padding.second}; + jintArray result = env->NewIntArray(2); + if (result == nullptr) { + return nullptr; + } + env->SetIntArrayRegion(result, 0, 2, values); + return result; +} + +} // extern "C" diff --git a/tools/cmake/preset/android.cmake b/tools/cmake/preset/android.cmake index 5c9bc97e3ef..ea8bcacc3b1 100644 --- a/tools/cmake/preset/android.cmake +++ b/tools/cmake/preset/android.cmake @@ -20,6 +20,7 @@ set_overridable_option(EXECUTORCH_XNNPACK_SHARED_WORKSPACE ON) set_overridable_option(EXECUTORCH_BUILD_EXTENSION_DATA_LOADER ON) set_overridable_option(EXECUTORCH_BUILD_EXTENSION_FLAT_TENSOR ON) +set_overridable_option(EXECUTORCH_BUILD_EXTENSION_IMAGE ON) set_overridable_option(EXECUTORCH_BUILD_EXTENSION_LLM ON) set_overridable_option(EXECUTORCH_BUILD_EXTENSION_LLM_RUNNER ON) set_overridable_option(EXECUTORCH_BUILD_EXTENSION_MODULE ON)