Skip to content

Commit 6e8f02a

Browse files
authored
Add JNI/Kotlin bindings for the ImageProcessor (pytorch#21830)
`extension/image` has shipped since pytorch#19967 and picked up ObjC/Swift bindings in pytorch#20051, but nothing exposes it to Android — the extension isn't built into the AAR at all. This adds the JNI layer and a Kotlin surface mirroring the ObjC one, as `org.pytorch.executorch.extension.image.ImageProcessor`. All the real work is @metascroy's; this only carries it across. Android callers write the resize-and-normalize step themselves today, because there is no API for it. That step is where letterbox geometry and channel order go wrong. The library already handles both. Its interleave kernel is NEON (pytorch#20137), so arm64 gets that path for free. I hit this converting vision models to `.pte`: on iOS it was a few lines, on Android a loop per model. **Review order** 1. `ImageProcessorConfig.kt`, `ImageProcessor.kt` — the public surface. Config, orientation, letterbox padding and the reuse variant match `ExecuTorchImageProcessor.h` one for one, minus `gpuMinInputPixels`. Input differs where the platforms do: there is no `CVPixelBuffer` to auto-detect, so a `Bitmap` overload and an explicit semi-planar YUV overload take its place. 2. `jni_layer_image.cpp` — the glue. 3. `tools/cmake/preset/android.cmake`, `extension/android/CMakeLists.txt` — one line to turn the extension on (mirroring `apple_common.cmake:26`), then the source and the link. 4. `ImageProcessorInstrumentationTest.kt` — pure image math, so no `.pte` fixture is needed. **Three calls I made, and I'm happy to change any of them** *The output tensor is allocated Java-side.* Kotlin allocates a direct `FloatBuffer` and native fills it in place through `process_into`. That makes `process()` and `processInto()` one native entry point. Capacity is passed as an explicit `int` rather than read with `GetDirectBufferCapacity`, whose unit is unspecified for view buffers. *`gpuMinInputPixels` is not exposed.* The portable implementation has no GPU path, so the field would report a decision that cannot happen on Android. The JNI pins it to `kGpuNever`. *Plane bounds are checked in the JNI.* `process_yuv_into` takes raw pointers, so it cannot see how large the Java buffers are. The decode reads the interleaved chroma plane through `uvStride * (height / 2 - 1) + width`. The binding rejects a shorter buffer rather than reading past its end. **Two things worth your eye** Turning on `EXECUTORCH_BUILD_EXTENSION_IMAGE` for Android makes the non-Apple branch of `extension/image/CMakeLists.txt` run `FetchContent` for stb at configure time. That is a new network fetch in the AAR build. If it doesn't suit CI, I can default the flag off for Android, or vendor the two stb headers — whichever you prefer. The chroma bounds check is the part I am least sure of. Some camera HALs return a `planes[2]` buffer whose capacity stops one byte short of the full interleaved plane. Strict checking turns that into an `InvalidArgument` instead of an out-of-bounds read, but it will reject input that happens to work today. Tolerating the short buffer is a one-line change if you'd rather. **Verification** Built and run on a Pixel 8a (arm64, NDK 28.2.13676358, `--preset android-arm64-v8a`, Release). - `extension_image` builds for `arm64-v8a`. `libexecutorch_jni.so` links with `jni_layer_image.cpp`, `extension_image` and `jnigraphics`, and all six `Java_..._ImageProcessor_*` entry points export through the existing `Java_*` version script. - `ImageProcessorInstrumentationTest`: 18 tests, 0 failures, 0 skipped, through the assembled AAR — so the `Bitmap` lock, the direct-`FloatBuffer` round trip and the exception mapping run as an app would hit them. - The same geometry and normalization cases also ran as a native arm64 binary against `extension_image`, putting the NEON path under test: channel order, ImageNet normalization, letterbox padding upright and rotated, the pad-value fill, and the clockwise direction of `Orientation::RIGHT`. 14/14. Also checked: `-Wall -Wextra` clean for `aarch64-linux-android26`, no `clang-format` diff, the whole `org/pytorch/executorch` Kotlin source set compiles with the new package, and `javap -s` confirms every native descriptor matches its C++ signature including argument order. Thanks for building the library in a shape that made this mostly wiring — happy to reshape any of it into whatever fits the tree. This PR was authored with Claude. cc @kirklandsign @cbilgin
1 parent 9b558d9 commit 6e8f02a

10 files changed

Lines changed: 1820 additions & 13 deletions

File tree

‎extension/android/CMakeLists.txt‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -161,6 +161,15 @@ if(EXECUTORCH_JNI_CUSTOM_LIBRARY)
161161
)
162162
endif()
163163

164+
if(EXECUTORCH_BUILD_EXTENSION_IMAGE)
165+
target_sources(executorch_jni PRIVATE jni/jni_layer_image.cpp)
166+
# jnigraphics provides AndroidBitmap_lockPixels, used to read Bitmap storage.
167+
list(APPEND link_libraries extension_image jnigraphics)
168+
target_compile_definitions(
169+
executorch_jni PUBLIC EXECUTORCH_BUILD_EXTENSION_IMAGE=1
170+
)
171+
endif()
172+
164173
if(EXECUTORCH_BUILD_EXTENSION_TRAINING)
165174
target_sources(executorch_jni PRIVATE jni/jni_layer_training.cpp)
166175
list(APPEND link_libraries extension_training)

0 commit comments

Comments
 (0)