Added task info

This commit is contained in:
Prianka Liz Kariat 2022-12-21 01:23:29 +05:30
parent c56ef735d7
commit 6d02108bf5
2 changed files with 205 additions and 0 deletions

View File

@ -0,0 +1,69 @@
// Copyright 2022 The MediaPipe Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#import <Foundation/Foundation.h>
#include "mediapipe/framework/calculator.pb.h"
#import "mediapipe/tasks/ios/core/sources/MPPTaskOptions.h"
#import "mediapipe/tasks/ios/core/sources/MPPTaskOptionsProtocol.h"
NS_ASSUME_NONNULL_BEGIN
/**
* Holds all needed informaton to initialize a MediaPipe Task.
*/
@interface MPPTaskInfo : NSObject <NSCopying>
@property(nonatomic, copy, nonnull) NSString *taskGraphName;
/**
* A task-specific options that is derived from MPPTaskOptions and confirms to
* MPPTaskOptionsProtocol.
*/
@property(nonatomic, copy) id<MPPTaskOptionsProtocol> taskOptions;
/**
* List of task graph input stream info strings in the form TAG:name.
*/
@property(nonatomic, copy) NSArray *inputStreams;
/**
* List of task graph output stream info in the form TAG:name.
*/
@property(nonatomic, copy) NSArray *outputStreams;
/**
* If the task requires a flow limiter.
*/
@property(nonatomic) BOOL enableFlowLimiting;
+ (instancetype)new NS_UNAVAILABLE;
- (instancetype)initWithTaskGraphName:(NSString *)taskGraphName
inputStreams:(NSArray<NSString *> *)inputStreams
outputStreams:(NSArray<NSString *> *)outputStreams
taskOptions:(id<MPPTaskOptionsProtocol>)taskOptions
enableFlowLimiting:(BOOL)enableFlowLimiting
error:(NSError **)error;
/**
* Creates a MediaPipe Task protobuf message from the MPPTaskInfo instance.
*/
- (mediapipe::CalculatorGraphConfig)generateGraphConfig;
- (instancetype)init NS_UNAVAILABLE;
@end
NS_ASSUME_NONNULL_END

View File

@ -0,0 +1,136 @@
// Copyright 2022 The MediaPipe Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#import "mediapipe/tasks/ios/core/sources/MPPTaskInfo.h"
#import "mediapipe/tasks/ios/common/sources/MPPCommon.h"
#import "mediapipe/tasks/ios/common/utils/sources/MPPCommonUtils.h"
#import "mediapipe/tasks/ios/common/utils/sources/NSString+Helpers.h"
#include "mediapipe/calculators/core/flow_limiter_calculator.pb.h"
#include "mediapipe/framework/calculator.pb.h"
#include "mediapipe/framework/calculator_options.pb.h"
namespace {
using CalculatorGraphConfig = ::mediapipe::CalculatorGraphConfig;
using Node = ::mediapipe::CalculatorGraphConfig::Node;
using ::mediapipe::InputStreamInfo;
using ::mediapipe::CalculatorOptions;
using ::mediapipe::FlowLimiterCalculatorOptions;
} // namespace
@implementation MPPTaskInfo
- (instancetype)initWithTaskGraphName:(NSString *)taskGraphName
inputStreams:(NSArray<NSString *> *)inputStreams
outputStreams:(NSArray<NSString *> *)outputStreams
taskOptions:(id<MPPTaskOptionsProtocol>)taskOptions
enableFlowLimiting:(BOOL)enableFlowLimiting
error:(NSError **)error {
if (!taskGraphName || !inputStreams.count || !outputStreams.count) {
[MPPCommonUtils
createCustomError:error
withCode:MPPTasksErrorCodeInvalidArgumentError
description:
@"Task graph's name, input streams, and output streams should be non-empty."];
}
self = [super init];
if (self) {
_taskGraphName = taskGraphName;
_inputStreams = inputStreams;
_outputStreams = outputStreams;
_taskOptions = taskOptions;
_enableFlowLimiting = enableFlowLimiting;
}
return self;
}
- (id)copyWithZone:(NSZone *)zone {
MPPTaskInfo *taskInfo = [[MPPTaskInfo alloc] init];
taskInfo.taskGraphName = self.taskGraphName;
taskInfo.inputStreams = self.inputStreams;
taskInfo.outputStreams = self.outputStreams;
taskInfo.taskOptions = self.taskOptions;
taskInfo.enableFlowLimiting = self.enableFlowLimiting;
return taskInfo;
}
- (CalculatorGraphConfig)generateGraphConfig {
CalculatorGraphConfig graph_config;
Node *task_subgraph_node = graph_config.add_node();
task_subgraph_node->set_calculator(self.taskGraphName.cppString);
[self.taskOptions copyToProto:task_subgraph_node->mutable_options()];
for (NSString *outputStream in self.outputStreams) {
auto cpp_output_stream = std::string(outputStream.cppString);
task_subgraph_node->add_output_stream(cpp_output_stream);
graph_config.add_output_stream(cpp_output_stream);
}
if (self.enableFlowLimiting) {
Node *flow_limit_calculator_node = graph_config.add_node();
flow_limit_calculator_node->set_calculator("FlowLimiterCalculator");
InputStreamInfo *input_stream_info = flow_limit_calculator_node->add_input_stream_info();
input_stream_info->set_tag_index("FINISHED");
input_stream_info->set_back_edge(true);
FlowLimiterCalculatorOptions *flow_limit_calculator_options =
flow_limit_calculator_node->mutable_options()->MutableExtension(
FlowLimiterCalculatorOptions::ext);
flow_limit_calculator_options->set_max_in_flight(1);
flow_limit_calculator_options->set_max_in_queue(1);
for (NSString *inputStream in self.inputStreams) {
graph_config.add_input_stream(inputStream.cppString);
NSString *strippedInputStream = [MPPTaskInfo stripTagIndex:inputStream];
flow_limit_calculator_node->add_input_stream(strippedInputStream.cppString);
NSString *taskInputStream = [MPPTaskInfo addStreamNamePrefix:inputStream];
task_subgraph_node->add_input_stream(taskInputStream.cppString);
NSString *strippedTaskInputStream = [MPPTaskInfo stripTagIndex:taskInputStream];
flow_limit_calculator_node->add_output_stream(strippedTaskInputStream.cppString);
}
NSString *firstOutputStream = self.outputStreams[0];
auto finished_output_stream = "FINISHED:" + firstOutputStream.cppString;
flow_limit_calculator_node->add_input_stream(finished_output_stream);
} else {
for (NSString *inputStream in self.inputStreams) {
auto cpp_input_stream = inputStream.cppString;
task_subgraph_node->add_input_stream(cpp_input_stream);
graph_config.add_input_stream(cpp_input_stream);
}
}
return graph_config;
}
+ (NSString *)stripTagIndex:(NSString *)tagIndexName {
return [tagIndexName componentsSeparatedByString:@":"][1];
}
+ (NSString *)addStreamNamePrefix:(NSString *)tagIndexName {
NSArray *splits = [tagIndexName componentsSeparatedByString:@":"];
return [NSString stringWithFormat:@"%@:throttled_%@", splits[0], splits[1]];
}
@end