|
| 1 | +/* |
| 2 | + * Copyright (c) Meta Platforms, Inc. and affiliates. |
| 3 | + * All rights reserved. |
| 4 | + * |
| 5 | + * This source code is licensed under the BSD-style license found in the |
| 6 | + * LICENSE file in the root directory of this source tree. |
| 7 | + */ |
| 8 | + |
| 9 | +#import <CoreVideo/CoreVideo.h> |
| 10 | +#import <Foundation/Foundation.h> |
| 11 | + |
| 12 | +#import "ExecuTorchTensor.h" |
| 13 | + |
| 14 | +NS_ASSUME_NONNULL_BEGIN |
| 15 | + |
| 16 | +typedef NS_ENUM(uint8_t, ExecuTorchImageResizeMode) { |
| 17 | + ExecuTorchImageResizeModeStretch, |
| 18 | + ExecuTorchImageResizeModeLetterbox, |
| 19 | +} NS_SWIFT_NAME(ImageResizeMode); |
| 20 | + |
| 21 | +typedef NS_ENUM(uint8_t, ExecuTorchImageLetterboxAnchor) { |
| 22 | + ExecuTorchImageLetterboxAnchorCenter, |
| 23 | + ExecuTorchImageLetterboxAnchorTopLeft, |
| 24 | +} NS_SWIFT_NAME(ImageLetterboxAnchor); |
| 25 | + |
| 26 | +NS_SWIFT_NAME(ImageNormalization) |
| 27 | +__attribute__((objc_subclassing_restricted)) |
| 28 | +@interface ExecuTorchImageNormalization : NSObject |
| 29 | + |
| 30 | ++ (instancetype)zeroToOne; |
| 31 | ++ (instancetype)imagenet; |
| 32 | + |
| 33 | ++ (instancetype)new NS_UNAVAILABLE; |
| 34 | +- (instancetype)init NS_UNAVAILABLE; |
| 35 | + |
| 36 | +@end |
| 37 | + |
| 38 | +NS_SWIFT_NAME(ImageProcessorConfig) |
| 39 | +__attribute__((objc_subclassing_restricted)) |
| 40 | +@interface ExecuTorchImageProcessorConfig : NSObject |
| 41 | + |
| 42 | +@property(nonatomic, readonly) NSInteger targetWidth; |
| 43 | +@property(nonatomic, readonly) NSInteger targetHeight; |
| 44 | +@property(nonatomic, readonly) ExecuTorchImageResizeMode resizeMode; |
| 45 | +@property(nonatomic, readonly) ExecuTorchImageLetterboxAnchor letterboxAnchor; |
| 46 | +@property(nonatomic, readonly) float padValue; |
| 47 | +@property(nonatomic, readonly) ExecuTorchImageNormalization *normalization; |
| 48 | +// Minimum source pixel count (width * height) at which the GPU path may be |
| 49 | +// used; smaller inputs run on the CPU. 0 forces GPU, NSIntegerMax forces CPU. |
| 50 | +@property(nonatomic, readonly) NSInteger gpuMinInputPixels; |
| 51 | + |
| 52 | +// Default value for gpuMinInputPixels (mirrors the C++ config default). |
| 53 | +@property(class, nonatomic, readonly) NSInteger defaultGpuMinInputPixels; |
| 54 | + |
| 55 | +- (instancetype)initWithTargetWidth:(NSInteger)targetWidth |
| 56 | + targetHeight:(NSInteger)targetHeight |
| 57 | + resizeMode:(ExecuTorchImageResizeMode)resizeMode |
| 58 | + letterboxAnchor:(ExecuTorchImageLetterboxAnchor)letterboxAnchor |
| 59 | + padValue:(float)padValue |
| 60 | + normalization:(ExecuTorchImageNormalization *)normalization |
| 61 | + gpuMinInputPixels:(NSInteger)gpuMinInputPixels NS_REFINED_FOR_SWIFT; |
| 62 | + |
| 63 | ++ (instancetype)new NS_UNAVAILABLE; |
| 64 | +- (instancetype)init NS_UNAVAILABLE; |
| 65 | + |
| 66 | +@end |
| 67 | + |
| 68 | +/// Thread-safety: ExecuTorchImageProcessor is NOT thread-safe per instance. |
| 69 | +/// Internal scratch buffers are mutated during processing. Use one instance |
| 70 | +/// per concurrent caller. Different instances are safe to use concurrently. |
| 71 | +NS_SWIFT_NAME(ImageProcessor) |
| 72 | +__attribute__((objc_subclassing_restricted)) |
| 73 | +@interface ExecuTorchImageProcessor : NSObject |
| 74 | + |
| 75 | +@property(nonatomic, readonly) ExecuTorchImageProcessorConfig *config; |
| 76 | + |
| 77 | +- (instancetype)initWithConfig:(ExecuTorchImageProcessorConfig *)config; |
| 78 | + |
| 79 | +/// Process a CVPixelBuffer into a normalized float tensor. |
| 80 | +/// |
| 81 | +/// Auto-detects pixel format from the buffer's metadata. Supported |
| 82 | +/// formats: BGRA, RGBA, 8-bit NV12, and 10-bit P010 (P010 is narrowed to NV12 |
| 83 | +/// internally). Other formats return an error. |
| 84 | +/// |
| 85 | +/// The buffer is treated as already upright. Orientation correction is not |
| 86 | +/// applied and cannot be derived from a CVPixelBuffer, so the caller is |
| 87 | +/// responsible for supplying an upright buffer (e.g. by configuring the |
| 88 | +/// capture connection's orientation). |
| 89 | +/// |
| 90 | +/// @param pixelBuffer The input pixel buffer. |
| 91 | +/// @param error On failure, set to an NSError describing what went wrong. |
| 92 | +/// @return An ExecuTorchTensor with shape [1, 3, H, W] (CHW), or nil on failure. |
| 93 | +- (nullable ExecuTorchTensor *)processPixelBuffer:(_Nullable CVPixelBufferRef)pixelBuffer |
| 94 | + error:(NSError **)error; |
| 95 | + |
| 96 | +/// Process a CVPixelBuffer into a caller-provided tensor, reusing its storage. |
| 97 | +/// |
| 98 | +/// Avoids the per-call output allocation of processPixelBuffer:error:, which |
| 99 | +/// matters for sustained video. `tensor` must be a Float tensor shaped |
| 100 | +/// [1, 3, targetHeight, targetWidth]; its storage is overwritten and can be |
| 101 | +/// reused across frames. The result aliases `tensor`, so the caller must |
| 102 | +/// finish using the previous result before the next call. |
| 103 | +/// |
| 104 | +/// @param pixelBuffer The input pixel buffer. |
| 105 | +/// @param tensor The output tensor to fill. |
| 106 | +/// @param error On failure, set to an NSError describing what went wrong. |
| 107 | +/// @return YES on success, NO on failure. |
| 108 | +- (BOOL)processPixelBuffer:(_Nullable CVPixelBufferRef)pixelBuffer |
| 109 | + intoTensor:(ExecuTorchTensor *)tensor |
| 110 | + error:(NSError **)error; |
| 111 | + |
| 112 | +/// Letterbox padding (per side, in pixels) the processor applies for a source |
| 113 | +/// of the given size: `x` is the left/right pad and `y` the top/bottom pad of |
| 114 | +/// the resized content. Returns {0, 0} for the stretch resize mode or the |
| 115 | +/// top-left anchor. Lets callers map the padded output back to the source |
| 116 | +/// region without replicating the resize geometry. |
| 117 | +/// |
| 118 | +/// @param inputWidth The source pixel width. |
| 119 | +/// @param inputHeight The source pixel height. |
| 120 | +/// @return The {x, y} padding (integer values stored in a CGPoint). |
| 121 | +- (CGPoint)computeLetterboxPaddingForInputWidth:(NSInteger)inputWidth |
| 122 | + height:(NSInteger)inputHeight |
| 123 | + NS_REFINED_FOR_SWIFT; |
| 124 | + |
| 125 | ++ (instancetype)new NS_UNAVAILABLE; |
| 126 | +- (instancetype)init NS_UNAVAILABLE; |
| 127 | + |
| 128 | +@end |
| 129 | + |
| 130 | +NS_ASSUME_NONNULL_END |
0 commit comments