98 lines
2.7 KiB
C++
98 lines
2.7 KiB
C++
#include "mediapipe/framework/tool/container_util.h"
|
|
|
|
#include "mediapipe/framework/tool/switch_container.pb.h"
|
|
|
|
namespace mediapipe {
|
|
namespace tool {
|
|
|
|
std::string ChannelTag(const std::string& tag, int channel) {
|
|
return absl::StrCat("C", channel, "__", tag);
|
|
}
|
|
|
|
// Parses a tag name starting with a channel prefix, like "C2__".
|
|
bool ParseChannelTag(const std::string& channel_name, std::string* name,
|
|
std::string* num) {
|
|
int pos = channel_name.find('C');
|
|
int sep = channel_name.find("__");
|
|
if (pos != 0 || sep == std::string::npos) {
|
|
return false;
|
|
}
|
|
*num = channel_name.substr(pos + 1, sep - (pos + 1));
|
|
*name = channel_name.substr(sep + 2);
|
|
return true;
|
|
}
|
|
|
|
std::set<std::string> ChannelTags(const std::shared_ptr<tool::TagMap>& map) {
|
|
std::set<std::string> result;
|
|
for (const std::string& tag : map->GetTags()) {
|
|
std::string name, num;
|
|
if (ParseChannelTag(tag, &name, &num)) {
|
|
result.insert(name);
|
|
}
|
|
}
|
|
return result;
|
|
}
|
|
|
|
int ChannelCount(const std::shared_ptr<tool::TagMap>& map) {
|
|
int count = 0;
|
|
for (const std::string& tag : map->GetTags()) {
|
|
std::string name, num;
|
|
int channel = -1;
|
|
if (ParseChannelTag(tag, &name, &num)) {
|
|
if (absl::SimpleAtoi(num, &channel)) {
|
|
count = std::max(count, channel + 1);
|
|
}
|
|
}
|
|
}
|
|
return count;
|
|
}
|
|
|
|
void Relay(const InputStreamShard& input, OutputStreamShard* output) {
|
|
if (input.IsEmpty()) {
|
|
Timestamp input_bound = input.Value().Timestamp().NextAllowedInStream();
|
|
if (output->NextTimestampBound() < input_bound) {
|
|
output->SetNextTimestampBound(input_bound);
|
|
}
|
|
} else {
|
|
output->AddPacket(input.Value());
|
|
}
|
|
}
|
|
|
|
int GetChannelIndex(const CalculatorContext& cc, int previous_index) {
|
|
int result = previous_index;
|
|
Packet select_packet;
|
|
Packet enable_packet;
|
|
if (cc.InputTimestamp() == Timestamp::Unstarted()) {
|
|
auto& options = cc.Options<mediapipe::SwitchContainerOptions>();
|
|
if (options.has_enable()) {
|
|
result = options.enable() ? 1 : 0;
|
|
}
|
|
if (options.has_select()) {
|
|
result = options.select();
|
|
}
|
|
if (cc.InputSidePackets().HasTag("ENABLE")) {
|
|
enable_packet = cc.InputSidePackets().Tag("ENABLE");
|
|
}
|
|
if (cc.InputSidePackets().HasTag("SELECT")) {
|
|
select_packet = cc.InputSidePackets().Tag("SELECT");
|
|
}
|
|
} else {
|
|
if (cc.Inputs().HasTag("ENABLE")) {
|
|
enable_packet = cc.Inputs().Tag("ENABLE").Value();
|
|
}
|
|
if (cc.Inputs().HasTag("SELECT")) {
|
|
select_packet = cc.Inputs().Tag("SELECT").Value();
|
|
}
|
|
}
|
|
if (!enable_packet.IsEmpty()) {
|
|
result = enable_packet.Get<bool>() ? 1 : 0;
|
|
}
|
|
if (!select_packet.IsEmpty()) {
|
|
result = select_packet.Get<int>();
|
|
}
|
|
return result;
|
|
}
|
|
|
|
} // namespace tool
|
|
} // namespace mediapipe
|