Merge pull request #4924 from google:mrschmidt/gpu
PiperOrigin-RevId: 578991016
This commit is contained in:
commit
35b9453da4
|
@ -14,6 +14,15 @@
|
|||
|
||||
#import <Foundation/Foundation.h>
|
||||
|
||||
/**
|
||||
* The delegate to run MediaPipe. If the delegate is not set, the default
|
||||
* delegate CPU is used.
|
||||
*/
|
||||
typedef NS_ENUM(NSUInteger, MPPDelegate) {
|
||||
MPPDelegateCPU,
|
||||
MPPDelegateGPU,
|
||||
} NS_SWIFT_NAME(Delegate);
|
||||
|
||||
NS_ASSUME_NONNULL_BEGIN
|
||||
|
||||
/**
|
||||
|
@ -26,6 +35,9 @@ NS_SWIFT_NAME(BaseOptions)
|
|||
/** The path to the model asset to open and mmap in memory. */
|
||||
@property(nonatomic, copy) NSString *modelAssetPath;
|
||||
|
||||
/** Overrides the default backend to use for the provided model. */
|
||||
@property(nonatomic) MPPDelegate delegate;
|
||||
|
||||
@end
|
||||
|
||||
NS_ASSUME_NONNULL_END
|
||||
|
|
|
@ -20,6 +20,7 @@
|
|||
self = [super init];
|
||||
if (self) {
|
||||
self.modelAssetPath = [[NSString alloc] init];
|
||||
self.delegate = MPPDelegateCPU;
|
||||
}
|
||||
return self;
|
||||
}
|
||||
|
@ -28,6 +29,7 @@
|
|||
MPPBaseOptions *baseOptions = [[MPPBaseOptions alloc] init];
|
||||
|
||||
baseOptions.modelAssetPath = self.modelAssetPath;
|
||||
baseOptions.delegate = self.delegate;
|
||||
|
||||
return baseOptions;
|
||||
}
|
||||
|
|
|
@ -21,6 +21,7 @@ objc_library(
|
|||
srcs = ["sources/MPPBaseOptions+Helpers.mm"],
|
||||
hdrs = ["sources/MPPBaseOptions+Helpers.h"],
|
||||
deps = [
|
||||
"//mediapipe/calculators/tensor:inference_calculator_cc_proto",
|
||||
"//mediapipe/tasks/cc/core/proto:acceleration_cc_proto",
|
||||
"//mediapipe/tasks/cc/core/proto:base_options_cc_proto",
|
||||
"//mediapipe/tasks/cc/core/proto:external_file_cc_proto",
|
||||
|
|
|
@ -12,12 +12,14 @@
|
|||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
#include "mediapipe/calculators/tensor/inference_calculator.pb.h"
|
||||
#include "mediapipe/tasks/cc/core/proto/acceleration.pb.h"
|
||||
#include "mediapipe/tasks/cc/core/proto/external_file.pb.h"
|
||||
#import "mediapipe/tasks/ios/core/utils/sources/MPPBaseOptions+Helpers.h"
|
||||
|
||||
namespace {
|
||||
using BaseOptionsProto = ::mediapipe::tasks::core::proto::BaseOptions;
|
||||
using InferenceCalculatorOptionsProto = ::mediapipe::InferenceCalculatorOptions;
|
||||
}
|
||||
|
||||
@implementation MPPBaseOptions (Helpers)
|
||||
|
@ -33,6 +35,11 @@ using BaseOptionsProto = ::mediapipe::tasks::core::proto::BaseOptions;
|
|||
if (self.modelAssetPath) {
|
||||
baseOptionsProto->mutable_model_asset()->set_file_name(self.modelAssetPath.UTF8String);
|
||||
}
|
||||
|
||||
if (self.delegate == MPPDelegateGPU) {
|
||||
baseOptionsProto->mutable_acceleration()->mutable_gpu()->MergeFrom(
|
||||
InferenceCalculatorOptionsProto::Delegate::Gpu());
|
||||
}
|
||||
}
|
||||
|
||||
@end
|
||||
|
|
Loading…
Reference in New Issue
Block a user