Extend orientation beyond UP (pytorch#20088)

Differential Revision: D107156015

Pull Request resolved: pytorch#20088

Author: metascroy

extension/apple/ExecuTorch/Exported/ExecuTorch+ImageProcessor.swift

@@ -64,33 +64,50 @@ public extension ImageProcessor {
   /// RGBA, 8-bit NV12, and 10-bit P010. Output is a `Tensor<Float>` with
   /// shape `[1, 3, target_height, target_width]`.
   ///
-  /// The buffer is treated as already upright: orientation correction is not
-  /// applied and cannot be derived from a CVPixelBuffer, so the caller is
-  /// responsible for supplying an upright buffer.
-  func process(_ pixelBuffer: CVPixelBuffer) throws -> Tensor<Float> {
-    let anyTensor = try processPixelBuffer(pixelBuffer)
+  /// `orientation` is the EXIF orientation of the buffer's contents; the
+  /// pipeline rotates it upright before resizing. It cannot be derived from a
+  /// CVPixelBuffer, so the caller supplies it (defaults to `.up`).
+  func process(
+    _ pixelBuffer: CVPixelBuffer,
+    orientation: ImageOrientation = .up
+  ) throws -> Tensor<Float> {
+    let anyTensor = try processPixelBuffer(pixelBuffer, orientation: orientation)
     return Tensor<Float>(anyTensor)
   }
 
   /// Process a CVPixelBuffer into a caller-provided tensor, reusing its storage.
   ///
-  /// Avoids the per-call allocation of `process(_:)`, which matters for
-  /// sustained video. `tensor` must be a `Tensor<Float>` with shape
+  /// Avoids the per-call allocation of `process(_:orientation:)`, which matters
+  /// for sustained video. `tensor` must be a `Tensor<Float>` with shape
   /// `[1, 3, target_height, target_width]`; its storage is overwritten and can
   /// be reused across frames. The contents are valid until the next call that
   /// writes into the same tensor.
   ///
-  /// The buffer is treated as already upright (see `process(_:)`).
-  func process(_ pixelBuffer: CVPixelBuffer, into tensor: Tensor<Float>) throws {
-    try processPixelBuffer(pixelBuffer, into: tensor.anyTensor)
+  /// `orientation` matches `process(_:orientation:)` (defaults to `.up`).
+  func process(
+    _ pixelBuffer: CVPixelBuffer,
+    orientation: ImageOrientation = .up,
+    into tensor: Tensor<Float>
+  ) throws {
+    try processPixelBuffer(
+      pixelBuffer, orientation: orientation, into: tensor.anyTensor)
   }
 
   /// Letterbox padding (per side, in pixels) applied for a source of the given
   /// size: `x` is the left/right pad and `y` the top/bottom pad of the resized
   /// content. Returns `(0, 0)` for the stretch resize mode or the top-left
   /// anchor. Lets callers map the padded output back to the source region.
-  func computeLetterboxPadding(inputWidth: Int, inputHeight: Int) -> (x: Int, y: Int) {
-    let padding = __computeLetterboxPadding(forInputWidth: inputWidth, height: inputHeight)
+  ///
+  /// `orientation` is the EXIF orientation of the source (defaults to `.up`);
+  /// the dimensions are oriented before the padding is computed, matching the
+  /// geometry `process(_:orientation:)` produces.
+  func computeLetterboxPadding(
+    inputWidth: Int,
+    inputHeight: Int,
+    orientation: ImageOrientation = .up
+  ) -> (x: Int, y: Int) {
+    let padding = __computeLetterboxPadding(
+      forInputWidth: inputWidth, height: inputHeight, orientation: orientation)
     return (padding.x, padding.y)
   }
 }

extension/apple/ExecuTorch/Exported/ExecuTorchImageProcessor.h

@@ -30,6 +30,14 @@ typedef struct ExecuTorchImageLetterboxPadding {
   NSInteger y;
 } ExecuTorchImageLetterboxPadding NS_SWIFT_NAME(ImageLetterboxPadding);
 
+/// EXIF orientation of the source image. The pipeline rotates the content
+/// upright before resizing. Only these rotation codes are supported.
+typedef NS_ENUM(uint8_t, ExecuTorchImageOrientation) {
+  ExecuTorchImageOrientationUp = 1,    // no rotation
+  ExecuTorchImageOrientationDown = 3,  // 180 degrees
+  ExecuTorchImageOrientationRight = 6, // 90 degrees clockwise
+  ExecuTorchImageOrientationLeft = 8,  // 90 degrees counter-clockwise
+} NS_SWIFT_NAME(ImageOrientation);
 NS_SWIFT_NAME(ImageNormalization)
 __attribute__((objc_subclassing_restricted))
 @interface ExecuTorchImageNormalization : NSObject
@@ -93,36 +101,52 @@ __attribute__((objc_subclassing_restricted))
 
 - (instancetype)initWithConfig:(ExecuTorchImageProcessorConfig *)config;
 
+/// Process a CVPixelBuffer into a normalized float tensor, treating the buffer
+/// as already upright (orientation `up`). Use
+/// processPixelBuffer:orientation:error: to specify a source orientation.
+- (nullable ExecuTorchTensor *)processPixelBuffer:(_Nullable CVPixelBufferRef)pixelBuffer
+                                            error:(NSError **)error;
+
+/// Reuse-friendly variant of processPixelBuffer:error: that writes into a
+/// caller-provided tensor; treats the buffer as already upright (orientation
+/// `up`). See processPixelBuffer:orientation:intoTensor:error:.
+- (BOOL)processPixelBuffer:(_Nullable CVPixelBufferRef)pixelBuffer
+                intoTensor:(ExecuTorchTensor *)tensor
+                     error:(NSError **)error;
+
 /// Process a CVPixelBuffer into a normalized float tensor.
 ///
 /// Auto-detects pixel format from the buffer's metadata. Supported
 /// formats: BGRA, RGBA, 8-bit NV12, and 10-bit P010 (P010 is narrowed to NV12
 /// internally). Other formats return an error.
 ///
-/// The buffer is treated as already upright. Orientation correction is not
-/// applied and cannot be derived from a CVPixelBuffer, so the caller is
-/// responsible for supplying an upright buffer (e.g. by configuring the
-/// capture connection's orientation).
+/// `orientation` is the EXIF orientation of the buffer's contents; the pipeline
+/// rotates it upright before resizing. It cannot be derived from a
+/// CVPixelBuffer, so the caller supplies it (e.g. from capture metadata).
 ///
 /// @param pixelBuffer The input pixel buffer.
+/// @param orientation The source orientation.
 /// @param error On failure, set to an NSError describing what went wrong.
 /// @return An ExecuTorchTensor with shape [1, 3, H, W] (CHW), or nil on failure.
 - (nullable ExecuTorchTensor *)processPixelBuffer:(_Nullable CVPixelBufferRef)pixelBuffer
+                                      orientation:(ExecuTorchImageOrientation)orientation
                                             error:(NSError **)error;
 
 /// Process a CVPixelBuffer into a caller-provided tensor, reusing its storage.
 ///
-/// Avoids the per-call output allocation of processPixelBuffer:error:, which
-/// matters for sustained video. `tensor` must be a Float tensor shaped
+/// Avoids the per-call output allocation of processPixelBuffer:orientation:error:,
+/// which matters for sustained video. `tensor` must be a Float tensor shaped
 /// [1, 3, targetHeight, targetWidth]; its storage is overwritten and can be
 /// reused across frames. The result aliases `tensor`, so the caller must
 /// finish using the previous result before the next call.
 ///
 /// @param pixelBuffer The input pixel buffer.
+/// @param orientation The source orientation (see processPixelBuffer:orientation:error:).
 /// @param tensor The output tensor to fill.
 /// @param error On failure, set to an NSError describing what went wrong.
 /// @return YES on success, NO on failure.
 - (BOOL)processPixelBuffer:(_Nullable CVPixelBufferRef)pixelBuffer
+               orientation:(ExecuTorchImageOrientation)orientation
                 intoTensor:(ExecuTorchTensor *)tensor
                      error:(NSError **)error;
 
@@ -132,11 +156,31 @@ __attribute__((objc_subclassing_restricted))
 /// top-left anchor. Lets callers map the padded output back to the source
 /// region without replicating the resize geometry.
 ///
+/// Treats the source as already upright (orientation `up`). Use
+/// computeLetterboxPaddingForInputWidth:height:orientation: for a rotated
+/// source.
+///
+/// @param inputWidth The source pixel width.
+/// @param inputHeight The source pixel height.
+/// @return The {x, y} padding in pixels.
+- (ExecuTorchImageLetterboxPadding)computeLetterboxPaddingForInputWidth:(NSInteger)inputWidth
+                                                                height:(NSInteger)inputHeight
+    NS_REFINED_FOR_SWIFT;
+
+/// Letterbox padding (per side, in pixels) the processor applies for a source
+/// of the given size and orientation. The source dimensions are oriented
+/// (width/height swapped for the 90-degree rotations) before the padding is
+/// computed, so the result matches the geometry that
+/// processPixelBuffer:orientation:error: produces. Returns {0, 0} for the
+/// stretch resize mode or the top-left anchor.
+///
 /// @param inputWidth The source pixel width.
 /// @param inputHeight The source pixel height.
+/// @param orientation The source orientation (see processPixelBuffer:orientation:error:).
 /// @return The {x, y} padding in pixels.
 - (ExecuTorchImageLetterboxPadding)computeLetterboxPaddingForInputWidth:(NSInteger)inputWidth
                                                                 height:(NSInteger)inputHeight
+                                     orientation:(ExecuTorchImageOrientation)orientation
     NS_REFINED_FOR_SWIFT;
 
 + (instancetype)new NS_UNAVAILABLE;

extension/apple/ExecuTorch/Exported/ExecuTorchImageProcessor.mm

@@ -31,6 +31,10 @@
 static_assert((int)ExecuTorchImageResizeModeLetterbox == (int)ResizeMode::LETTERBOX, "ExecuTorchImageResizeModeLetterbox must match ResizeMode::LETTERBOX");
 static_assert((int)ExecuTorchImageLetterboxAnchorCenter == (int)LetterboxAnchor::CENTER, "ExecuTorchImageLetterboxAnchorCenter must match LetterboxAnchor::CENTER");
 static_assert((int)ExecuTorchImageLetterboxAnchorTopLeft == (int)LetterboxAnchor::TOP_LEFT, "ExecuTorchImageLetterboxAnchorTopLeft must match LetterboxAnchor::TOP_LEFT");
+static_assert((int)ExecuTorchImageOrientationUp == (int)Orientation::UP, "ExecuTorchImageOrientationUp must match Orientation::UP");
+static_assert((int)ExecuTorchImageOrientationDown == (int)Orientation::DOWN, "ExecuTorchImageOrientationDown must match Orientation::DOWN");
+static_assert((int)ExecuTorchImageOrientationRight == (int)Orientation::RIGHT, "ExecuTorchImageOrientationRight must match Orientation::RIGHT");
+static_assert((int)ExecuTorchImageOrientationLeft == (int)Orientation::LEFT, "ExecuTorchImageOrientationLeft must match Orientation::LEFT");
 
 // MARK: - Private interfaces
 
@@ -178,17 +182,36 @@ - (instancetype)initWithConfig:(ExecuTorchImageProcessorConfig *)config {
 
 - (nullable ExecuTorchTensor *)processPixelBuffer:(_Nullable CVPixelBufferRef)pixelBuffer
                                             error:(NSError **)error {
+  return [self processPixelBuffer:pixelBuffer
+                      orientation:ExecuTorchImageOrientationUp
+                            error:error];
+}
+
+- (BOOL)processPixelBuffer:(_Nullable CVPixelBufferRef)pixelBuffer
+                intoTensor:(ExecuTorchTensor *)tensor
+                     error:(NSError **)error {
+  return [self processPixelBuffer:pixelBuffer
+                      orientation:ExecuTorchImageOrientationUp
+                       intoTensor:tensor
+                            error:error];
+}
+
+- (nullable ExecuTorchTensor *)processPixelBuffer:(_Nullable CVPixelBufferRef)pixelBuffer
+                                      orientation:(ExecuTorchImageOrientation)orientation
+                                            error:(NSError **)error {
   if (!pixelBuffer) {
     if (error) {
       *error = ExecuTorchErrorWithCode(ExecuTorchErrorCodeInvalidArgument);
     }
     return nil;
   }
-  auto result = process_pixelbuffer(*_processor, pixelBuffer);
+  auto result = process_pixelbuffer(
+      *_processor, pixelBuffer, static_cast<Orientation>(orientation));
   return tensorFromResult(result, error);
 }
 
 - (BOOL)processPixelBuffer:(_Nullable CVPixelBufferRef)pixelBuffer
+               orientation:(ExecuTorchImageOrientation)orientation
                 intoTensor:(ExecuTorchTensor *)tensor
                      error:(NSError **)error {
   if (!pixelBuffer || !tensor) {
@@ -199,7 +222,8 @@ - (BOOL)processPixelBuffer:(_Nullable CVPixelBufferRef)pixelBuffer
   }
   auto* tensorPtr = reinterpret_cast<TensorPtr*>(tensor.nativeInstance);
   auto err = process_pixelbuffer_into(
-      *_processor, pixelBuffer, Orientation::UP, **tensorPtr);
+      *_processor, pixelBuffer, static_cast<Orientation>(orientation),
+      **tensorPtr);
   if (err != executorch::runtime::Error::Ok) {
     if (error) {
       *error = ExecuTorchErrorWithCode((ExecuTorchErrorCode)err);
@@ -211,8 +235,17 @@ - (BOOL)processPixelBuffer:(_Nullable CVPixelBufferRef)pixelBuffer
 
 - (ExecuTorchImageLetterboxPadding)computeLetterboxPaddingForInputWidth:(NSInteger)inputWidth
                                                                 height:(NSInteger)inputHeight {
+  return [self computeLetterboxPaddingForInputWidth:inputWidth
+                                             height:inputHeight
+                                        orientation:ExecuTorchImageOrientationUp];
+}
+
+- (ExecuTorchImageLetterboxPadding)computeLetterboxPaddingForInputWidth:(NSInteger)inputWidth
+                                                                height:(NSInteger)inputHeight
+                                     orientation:(ExecuTorchImageOrientation)orientation {
   const auto padding = _processor->compute_letterbox_padding(
-      static_cast<int32_t>(inputWidth), static_cast<int32_t>(inputHeight));
+      static_cast<int32_t>(inputWidth), static_cast<int32_t>(inputHeight),
+      static_cast<Orientation>(orientation));
   return {padding.first, padding.second};
 }