File plugin.cpp
Go to the documentation of this file
#include "plugin.hpp"
#include "absl/flags/parse.h"
#include "illixr/data_format/misc.hpp"
#include "illixr/data_format/zed_cam.hpp"
#include "mediapipe/calculators/util/image_data.pb.h"
#include "mediapipe/framework/deps/file_helpers.h"
#include "mediapipe/framework/formats/image_frame.h"
#include "mediapipe/framework/formats/image_frame_opencv.h"
#include "mediapipe/framework/port/parse_text_proto.h"
#include <cmath>
#if !MEDIAPIPE_DISABLE_GPU
#include "mediapipe/gpu/gpu_buffer.h"
#include "mediapipe/gpu/gpu_shared_data_internal.h"
#endif
using namespace ILLIXR;
namespace idf = ILLIXR::data_format;
namespace iht = ILLIXR::data_format::ht;
constexpr char kInputStream[] = "input_video";
constexpr char kOutputStream[] = "illixr_data";
constexpr char kImageDataTag[] = "image_data";
void img_convert(cv::Mat& img, bool flip = false) {
switch (img.type()) {
case CV_8UC1:
cv::cvtColor(img, img, cv::COLOR_GRAY2RGB);
break;
case CV_8UC3:
break;
case CV_8UC4:
if (flip) {
cv::cvtColor(img, img, cv::COLOR_BGRA2RGB);
} else {
cv::cvtColor(img, img, cv::COLOR_RGBA2RGB);
}
break;
}
}
ht::cam_type hand_tracking::get_cam_type() {
std::string input_src = switchboard_->get_env("HT_INPUT");
if (input_src.empty())
throw std::runtime_error("HT_INPUT did not have a valid type");
if (ILLIXR::data_format::compare(input_src, "zed")) {
return ht::ZED;
} else if (ILLIXR::data_format::compare(input_src, "cam")) {
return ht::CAM;
} else if (ILLIXR::data_format::compare(input_src, "webcam")) {
return ht::WEBCAM;
} else {
throw std::runtime_error("HT_INPUT did not have a valid type");
}
}
[[maybe_unused]] hand_tracking::hand_tracking(const std::string& name_, phonebook* pb_)
: plugin{name_, pb_}
, switchboard_{pb_->lookup_impl<switchboard>()}
, graph_{{idf::image::LEFT_EYE, nullptr}, {idf::image::RIGHT_EYE, nullptr}, {idf::image::RGB, nullptr}}
, cam_type_{get_cam_type()}
, publisher_{"hand_tracking_publisher", pb_} {
if (cam_type_ == ht::WEBCAM) {
input_type_ = ht::RGB;
first_person_ = false;
} else if (const char *in_type = switchboard_->get_env_char("HT_INPUT_TYPE")) {
if (strcmp(in_type, "LEFT") == 0 || strcmp(in_type, "SINGLE") == 0) {
input_type_ = ht::LEFT;
} else if (strcmp(in_type, "RIGHT") == 0) {
input_type_ = ht::RIGHT;
} else if (strcmp(in_type, "MULTI") == 0 || strcmp(in_type, "BOTH") == 0) {
input_type_ = ht::BOTH;
} else if (strcmp(in_type, "RGB") == 0) {
input_type_ = ht::RGB;
}
} else {
input_type_ = ht::BOTH;
}
if (input_type_ == ht::RGB) {
graph_.at(idf::image::RGB) = new mediapipe::CalculatorGraph();
} else {
if (input_type_ != ht::RIGHT)
graph_.at(idf::image::LEFT_EYE) = new mediapipe::CalculatorGraph();
if (input_type_ == ht::RIGHT || input_type_ == ht::BOTH)
graph_.at(idf::image::RIGHT_EYE) = new mediapipe::CalculatorGraph();
}
publisher_.set_frame_count(input_type_);
}
void hand_tracking::start() {
plugin::start();
const std::string calculator_graph_config_contents =
#if !MEDIAPIPE_DISABLE_GPU
#include "mediapipe/hand_tracking_desktop_live_gpu.pbtxt"
#else
#include "mediapipe/hand_tracking_desktop_live.pbtxt"
#endif
; // NOLINT(whitespace/semicolon)
auto config = mediapipe::ParseTextProtoOrDie<mediapipe::CalculatorGraphConfig>(calculator_graph_config_contents);
absl::Status status;
if (input_type_ == ht::RGB) {
status = graph_[idf::image::RGB]->Initialize(config);
if (!status.ok())
throw std::runtime_error(std::string(status.message()));
} else {
if (input_type_ != ht::RIGHT) {
status = graph_[idf::image::LEFT_EYE]->Initialize(config);
if (!status.ok())
throw std::runtime_error(std::string(status.message()));
}
if (input_type_ == ht::RIGHT || input_type_ == ht::BOTH) {
status = graph_[idf::image::RIGHT_EYE]->Initialize(config);
if (!status.ok())
throw std::runtime_error(std::string(status.message()));
}
}
#if !MEDIAPIPE_DISABLE_GPU
ABSL_LOG(INFO) << "Initialize the GPU.";
#endif
for (auto& item : graph_) {
if (item.second != nullptr) {
#if !MEDIAPIPE_DISABLE_GPU
auto gpu_resources = mediapipe::GpuResources::Create();
if (!gpu_resources.ok())
throw std::runtime_error("");
status = graph_.at(item.first)->SetGpuResources(std::move(gpu_resources).value());
if (!status.ok())
throw std::runtime_error(std::string(status.message()));
gpu_helper_[item.first].InitializeForTest(graph_.at(item.first)->GetGpuResources().get());
#endif
if (graph_.at(item.first) != nullptr) {
auto status_or_poller = graph_.at(item.first)->AddOutputStreamPoller(kOutputStream);
if (!status_or_poller.ok())
throw std::runtime_error("Error with output poller");
publisher_.set_poller(item.first, new mediapipe::OutputStreamPoller(std::move(status_or_poller).value()));
}
}
}
for (auto& item : graph_) {
if (item.second != nullptr) {
status = graph_.at(item.first)->StartRun({});
if (!status.ok())
throw std::runtime_error("Error starting graph");
}
}
// subscribe to the expected type
switch (cam_type_) {
case ht::WEBCAM:
switchboard_->schedule<idf::monocular_cam_type>(
id_, "webcam", [this](const switchboard::ptr<const idf::monocular_cam_type>& frame, std::size_t) {
this->process(frame);
});
break;
case ht::CAM:
switchboard_->schedule<idf::binocular_cam_type>(
id_, "cam", [this](const switchboard::ptr<const idf::binocular_cam_type>& img, std::size_t) {
this->process(img);
});
break;
case ht::ZED:
#ifdef HAVE_ZED
switchboard_->schedule<idf::cam_type_zed>(id_, "cam_zed",
[this](const switchboard::ptr<const idf::cam_type_zed>& img, std::size_t) {
this->process(img);
});
#else
throw std::runtime_error("No support for zed camera");
#endif
break;
default:
throw std::runtime_error("Unexpected HT_INPUT entry. It must be one of webcam, cam, or zed (if supported).");
}
publisher_.start();
}
void hand_tracking::stop() {
for (auto& item : graph_) {
if (item.second != nullptr) {
item.second->Cancel();
auto status = item.second->CloseAllPacketSources();
delete item.second;
item.second = nullptr;
}
}
publisher_.stop();
plugin::stop();
}
hand_tracking::~hand_tracking() {
for (auto& i : graph_)
delete i.second;
}
void hand_tracking::process(const switchboard::ptr<const idf::cam_base_type>& frame) {
current_images_.clear();
pose_image pose_img;
switch (frame->type) {
case idf::camera::BINOCULAR:
switch (input_type_) {
case ht::BOTH:
// If we have both eye data
if (!frame->at(idf::image::LEFT_EYE).empty() && !frame->at(idf::image::RIGHT_EYE).empty()) {
current_images_ = {{idf::image::LEFT_EYE, frame->at(idf::image::LEFT_EYE).clone()},
{idf::image::RIGHT_EYE, frame->at(idf::image::RIGHT_EYE).clone()}};
pose_img.eye_count = 2;
// if we have neither eye data
} else if (frame->at(idf::image::LEFT_EYE).empty() && frame->at(idf::image::RIGHT_EYE).empty()) {
spdlog::get("illixr")->info("[hand_tracking.plugin] Received empty frame, skipping");
return;
// if there is no left eye data, then just use the right
} else if (frame->at(idf::image::LEFT_EYE).empty()) {
current_images_ = {{idf::image::RIGHT_EYE, frame->at(idf::image::RIGHT_EYE).clone()}};
pose_img.eye_count = 1;
pose_img.primary = idf::units::RIGHT_EYE;
// if we have just left eye data, then just use it.
} else {
current_images_ = {{idf::image::LEFT_EYE, frame->at(idf::image::LEFT_EYE).clone()}};
pose_img.eye_count = 1;
pose_img.primary = idf::units::LEFT_EYE;
}
break;
case ht::LEFT:
if (frame->at(idf::image::LEFT_EYE).empty()) {
spdlog::get("illixr")->info("[hand_tracking.plugin] Received empty frame, skipping");
return;
}
current_images_ = {{idf::image::RIGHT_EYE, frame->at(idf::image::RIGHT_EYE).clone()}};
pose_img.eye_count = 1;
pose_img.primary = idf::units::RIGHT_EYE;
break;
case ht::RIGHT:
if (frame->at(idf::image::RIGHT_EYE).empty()) {
spdlog::get("illixr")->info("[hand_tracking.plugin] Received empty frame, skipping");
return;
}
current_images_ = {{idf::image::LEFT_EYE, frame->at(idf::image::LEFT_EYE).clone()}};
pose_img.eye_count = 1;
pose_img.primary = idf::units::LEFT_EYE;
break;
case ht::RGB:
spdlog::get("illixr")->warn("[hand_tracking.plugin] RGB not provided by binocular view");
break;
}
pose_img.insert(current_images_.begin(), current_images_.end());
break;
case idf::camera::MONOCULAR: {
if (frame->at(idf::image::RGB).empty()) {
spdlog::get("illixr")->info("[hand_tracking.plugin] Received empty frame, skipping");
return;
}
cv::Mat temp_img(frame->at(idf::image::RGB).clone());
cv::flip(temp_img, temp_img, 1);
img_convert(temp_img);
current_images_ = {{idf::image::LEFT_EYE, temp_img}};
pose_img.eye_count = 1;
pose_img.primary = idf::units::LEFT_EYE;
pose_img.insert(current_images_.begin(), current_images_.end());
break;
}
case idf::camera::RGB_DEPTH: {
if (frame->at(idf::image::LEFT_EYE).empty() && frame->at(idf::image::RGB).empty()) {
spdlog::get("illixr")->info("[hand_tracking.plugin] Received empty frame, skipping");
return;
}
current_images_ = {{idf::image::LEFT_EYE, frame->at(idf::image::RGB).clone()}};
pose_img.eye_count = 1;
pose_img.primary = idf::units::LEFT_EYE;
pose_img.insert(current_images_.begin(), current_images_.end());
break;
}
case idf::camera::ZED: {
#ifdef HAVE_ZED
switch (input_type_) {
case ht::BOTH: {
if (!frame->at(idf::image::LEFT_EYE).empty() && !frame->at(idf::image::RIGHT_EYE).empty()) {
cv::Mat tempL(frame->at(idf::image::LEFT_EYE).clone());
cv::Mat tempR(frame->at(idf::image::RIGHT_EYE).clone());
img_convert(tempL, true);
img_convert(tempR, true);
current_images_ = {{idf::image::LEFT_EYE, tempL}, {idf::image::RIGHT_EYE, tempR}};
pose_img.eye_count = 2;
pose_img.primary = idf::units::LEFT_EYE;
} else if (frame->at(idf::image::LEFT_EYE).empty() && frame->at(idf::image::RIGHT_EYE).empty()) {
spdlog::get("illixr")->info("[hand_tracking.plugin] Received empty frame, skipping");
return;
} else if (frame->at(idf::image::LEFT_EYE).empty()) {
cv::Mat temp(frame->at(idf::image::RIGHT_EYE).clone());
img_convert(temp, true);
current_images_ = {{idf::image::RIGHT_EYE, temp}};
pose_img.eye_count = 1;
pose_img.primary = idf::units::RIGHT_EYE;
} else {
cv::Mat temp(frame->at(idf::image::LEFT_EYE).clone());
img_convert(temp, true);
current_images_ = {{idf::image::LEFT_EYE, temp}};
pose_img.eye_count = 1;
pose_img.primary = idf::units::LEFT_EYE;
}
break;
}
case ht::LEFT: {
if (frame->at(idf::image::LEFT_EYE).empty()) {
spdlog::get("illixr")->info("[hand_tracking.plugin] Received empty frame, skipping");
return;
}
cv::Mat temp(frame->at(idf::image::LEFT_EYE).clone());
img_convert(temp, true);
current_images_ = {{idf::image::LEFT_EYE, temp}};
pose_img.eye_count = 1;
pose_img.primary = idf::units::LEFT_EYE;
break;
}
case ht::RIGHT: {
if (frame->at(idf::image::RIGHT_EYE).empty()) {
spdlog::get("illixr")->info("[hand_tracking.plugin] Received empty frame, skipping");
return;
}
cv::Mat temp(frame->at(idf::image::RIGHT_EYE).clone());
img_convert(temp, true);
current_images_ = {{idf::image::RIGHT_EYE, temp}};
pose_img.eye_count = 1;
pose_img.primary = idf::units::RIGHT_EYE;
break;
}
case ht::RGB: {
if (frame->at(idf::image::LEFT_EYE).empty() && frame->at(idf::image::RGB).empty()) {
spdlog::get("illixr")->info("[hand_tracking.plugin] Received empty frame, skipping");
return;
}
cv::Mat temp(frame->at(idf::image::RGB).clone());
img_convert(temp, true);
current_images_ = {{idf::image::LEFT_EYE, temp}};
pose_img.eye_count = 1;
pose_img.primary = idf::units::LEFT_EYE;
break;
}
}
if (frame->find(idf::image::DEPTH) != frame->end()) {
if (!pose_img.images[idf::image::DEPTH].empty()) {
pose_img.images[idf::image::DEPTH] = frame->at(idf::image::DEPTH).clone();
pose_img.depth_valid = true;
}
}
if (frame->find(idf::image::CONFIDENCE) != frame->end()) {
if (!pose_img.images[idf::image::CONFIDENCE].empty()) {
pose_img.images[idf::image::CONFIDENCE] = frame->at(idf::image::CONFIDENCE).clone();
pose_img.confidence_valid_ = true;
}
}
pose_img.poses = dynamic_cast<const idf::cam_type_zed*>(frame.get())->poses;
pose_img.pose_valid = true;
break;
#endif
}
default:
throw std::runtime_error("Unexpected frame type");
}
pose_img.insert(current_images_.begin(), current_images_.end());
size_t frame_id = std::chrono::high_resolution_clock::now().time_since_epoch().count();
for (const auto& input : current_images_) {
auto input_frame =
absl::make_unique<mediapipe::ImageFrame>(mediapipe::ImageFormat::SRGB, input.second.cols, input.second.rows,
#if !MEDIAPIPE_DISABLE_GPU
mediapipe::ImageFrame::kGlDefaultAlignmentBoundary
#else
mediapipe::ImageFrame::kDefaultAlignmentBoundary
#endif
);
cv::Mat input_frame_mat = mediapipe::formats::MatView(input_frame.get());
input.second.copyTo(input_frame_mat);
// Send image packet into the graph.
int64_t frame_timestamp_us = std::chrono::high_resolution_clock::now().time_since_epoch().count() / 1000;
// make sure image data makes it into stream first
mediapipe::ImageData image_data;
image_data.set_width(input_frame.get()->Width());
image_data.set_height(input_frame.get()->Height());
image_data.set_frame_id(frame_id);
image_data.set_image_type(input.first);
image_data.set_first_person(first_person_);
auto img_ptr = absl::make_unique<mediapipe::ImageData>(image_data);
auto img_status =
graph_.at(input.first)
->AddPacketToInputStream(kImageDataTag,
mediapipe::Adopt(img_ptr.release()).At(mediapipe::Timestamp(frame_timestamp_us)));
if (!img_status.ok())
throw std::runtime_error(std::string(img_status.message()));
#if !MEDIAPIPE_DISABLE_GPU
auto gl_status =
gpu_helper_.at(input.first)
.RunInGlContext([&input_frame, &frame_timestamp_us, &graph = graph_.at(input.first),
&gpu_helper = gpu_helper_.at(input.first)]() -> absl::Status {
// Convert ImageFrame to GpuBuffer.
auto texture = gpu_helper.CreateSourceTexture(*input_frame.get());
auto gpu_frame = texture.GetFrame<mediapipe::GpuBuffer>();
glFlush();
texture.Release();
// Send GPU image packet into the graph.
auto submit_status = graph->AddPacketToInputStream(
kInputStream, mediapipe::Adopt(gpu_frame.release()).At(mediapipe::Timestamp(frame_timestamp_us)));
if (!submit_status.ok())
throw std::runtime_error(std::string(submit_status.message()));
return absl::OkStatus();
});
if (!gl_status.ok())
throw std::runtime_error(std::string(gl_status.message()));
#else
auto submit_status =
graph_.at(input.first)
->AddPacketToInputStream(kInputStream,
mediapipe::Adopt(input_frame.release()).At(mediapipe::Timestamp(frame_timestamp_us)));
if (!submit_status.ok())
throw std::runtime_error(std::string(submit_status.message()));
#endif
}
publisher_.add_raw(frame_id, pose_img);
}
// ###############################################################################################
PLUGIN_MAIN(hand_tracking)