talksik/mediapipe
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Project import generated by Copybara.

Browse Source 
GitOrigin-RevId: 8e1da4611d93ccb7d9674713157d43be0348d98f
This commit is contained in:
MediaPipe Team
2021-07-27 22:36:23 -04:00
committed by chuoling
parent 50c92c6623
commit b899d17f18
79 changed files with 1808 additions and 946 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mediapipe/calculators/video/box_detector_calculator.cc
+75 -58
View File
@@ -47,6 +47,21 @@
namespace mediapipe {
constexpr char kFrameAlignmentTag[] = "FRAME_ALIGNMENT";
constexpr char kOutputIndexFilenameTag[] = "OUTPUT_INDEX_FILENAME";
constexpr char kIndexProtoStringTag[] = "INDEX_PROTO_STRING";
constexpr char kVizTag[] = "VIZ";
constexpr char kBoxesTag[] = "BOXES";
constexpr char kReacqSwitchTag[] = "REACQ_SWITCH";
constexpr char kCancelObjectIdTag[] = "CANCEL_OBJECT_ID";
constexpr char kAddIndexTag[] = "ADD_INDEX";
constexpr char kImageSizeTag[] = "IMAGE_SIZE";
constexpr char kDescriptorsTag[] = "DESCRIPTORS";
constexpr char kFeaturesTag[] = "FEATURES";
constexpr char kVideoTag[] = "VIDEO";
constexpr char kTrackedBoxesTag[] = "TRACKED_BOXES";
constexpr char kTrackingTag[] = "TRACKING";
// A calculator to detect reappeared box positions from single frame.
//
// Input stream:
@@ -110,66 +125,66 @@ class BoxDetectorCalculator : public CalculatorBase {
REGISTER_CALCULATOR(BoxDetectorCalculator);
absl::Status BoxDetectorCalculator::GetContract(CalculatorContract* cc) {
if (cc->Inputs().HasTag("TRACKING")) {
cc->Inputs().Tag("TRACKING").Set<TrackingData>();
if (cc->Inputs().HasTag(kTrackingTag)) {
cc->Inputs().Tag(kTrackingTag).Set<TrackingData>();
}
if (cc->Inputs().HasTag("TRACKED_BOXES")) {
cc->Inputs().Tag("TRACKED_BOXES").Set<TimedBoxProtoList>();
if (cc->Inputs().HasTag(kTrackedBoxesTag)) {
cc->Inputs().Tag(kTrackedBoxesTag).Set<TimedBoxProtoList>();
}
if (cc->Inputs().HasTag("VIDEO")) {
cc->Inputs().Tag("VIDEO").Set<ImageFrame>();
if (cc->Inputs().HasTag(kVideoTag)) {
cc->Inputs().Tag(kVideoTag).Set<ImageFrame>();
}
if (cc->Inputs().HasTag("FEATURES")) {
RET_CHECK(cc->Inputs().HasTag("DESCRIPTORS"))
if (cc->Inputs().HasTag(kFeaturesTag)) {
RET_CHECK(cc->Inputs().HasTag(kDescriptorsTag))
<< "FEATURES and DESCRIPTORS need to be specified together.";
cc->Inputs().Tag("FEATURES").Set<std::vector<cv::KeyPoint>>();
cc->Inputs().Tag(kFeaturesTag).Set<std::vector<cv::KeyPoint>>();
}
if (cc->Inputs().HasTag("DESCRIPTORS")) {
RET_CHECK(cc->Inputs().HasTag("FEATURES"))
if (cc->Inputs().HasTag(kDescriptorsTag)) {
RET_CHECK(cc->Inputs().HasTag(kFeaturesTag))
<< "FEATURES and DESCRIPTORS need to be specified together.";
cc->Inputs().Tag("DESCRIPTORS").Set<std::vector<float>>();
cc->Inputs().Tag(kDescriptorsTag).Set<std::vector<float>>();
}
if (cc->Inputs().HasTag("IMAGE_SIZE")) {
cc->Inputs().Tag("IMAGE_SIZE").Set<std::pair<int, int>>();
if (cc->Inputs().HasTag(kImageSizeTag)) {
cc->Inputs().Tag(kImageSizeTag).Set<std::pair<int, int>>();
}
if (cc->Inputs().HasTag("ADD_INDEX")) {
cc->Inputs().Tag("ADD_INDEX").Set<std::string>();
if (cc->Inputs().HasTag(kAddIndexTag)) {
cc->Inputs().Tag(kAddIndexTag).Set<std::string>();
}
if (cc->Inputs().HasTag("CANCEL_OBJECT_ID")) {
cc->Inputs().Tag("CANCEL_OBJECT_ID").Set<int>();
if (cc->Inputs().HasTag(kCancelObjectIdTag)) {
cc->Inputs().Tag(kCancelObjectIdTag).Set<int>();
}
if (cc->Inputs().HasTag("REACQ_SWITCH")) {
cc->Inputs().Tag("REACQ_SWITCH").Set<bool>();
if (cc->Inputs().HasTag(kReacqSwitchTag)) {
cc->Inputs().Tag(kReacqSwitchTag).Set<bool>();
}
if (cc->Outputs().HasTag("BOXES")) {
cc->Outputs().Tag("BOXES").Set<TimedBoxProtoList>();
if (cc->Outputs().HasTag(kBoxesTag)) {
cc->Outputs().Tag(kBoxesTag).Set<TimedBoxProtoList>();
}
if (cc->Outputs().HasTag("VIZ")) {
RET_CHECK(cc->Inputs().HasTag("VIDEO"))
if (cc->Outputs().HasTag(kVizTag)) {
RET_CHECK(cc->Inputs().HasTag(kVideoTag))
<< "Output stream VIZ requires VIDEO to be present.";
cc->Outputs().Tag("VIZ").Set<ImageFrame>();
cc->Outputs().Tag(kVizTag).Set<ImageFrame>();
}
if (cc->InputSidePackets().HasTag("INDEX_PROTO_STRING")) {
cc->InputSidePackets().Tag("INDEX_PROTO_STRING").Set<std::string>();
if (cc->InputSidePackets().HasTag(kIndexProtoStringTag)) {
cc->InputSidePackets().Tag(kIndexProtoStringTag).Set<std::string>();
}
if (cc->InputSidePackets().HasTag("OUTPUT_INDEX_FILENAME")) {
cc->InputSidePackets().Tag("OUTPUT_INDEX_FILENAME").Set<std::string>();
if (cc->InputSidePackets().HasTag(kOutputIndexFilenameTag)) {
cc->InputSidePackets().Tag(kOutputIndexFilenameTag).Set<std::string>();
}
if (cc->InputSidePackets().HasTag("FRAME_ALIGNMENT")) {
cc->InputSidePackets().Tag("FRAME_ALIGNMENT").Set<int>();
if (cc->InputSidePackets().HasTag(kFrameAlignmentTag)) {
cc->InputSidePackets().Tag(kFrameAlignmentTag).Set<int>();
}
return absl::OkStatus();
@@ -179,10 +194,10 @@ absl::Status BoxDetectorCalculator::Open(CalculatorContext* cc) {
options_ = cc->Options<BoxDetectorCalculatorOptions>();
box_detector_ = BoxDetectorInterface::Create(options_.detector_options());
if (cc->InputSidePackets().HasTag("INDEX_PROTO_STRING")) {
if (cc->InputSidePackets().HasTag(kIndexProtoStringTag)) {
BoxDetectorIndex predefined_index;
if (!predefined_index.ParseFromString(cc->InputSidePackets()
.Tag("INDEX_PROTO_STRING")
.Tag(kIndexProtoStringTag)
.Get<std::string>())) {
LOG(FATAL) << "failed to parse BoxDetectorIndex from INDEX_PROTO_STRING";
}
@@ -202,12 +217,13 @@ absl::Status BoxDetectorCalculator::Open(CalculatorContext* cc) {
box_detector_->AddBoxDetectorIndex(predefined_index);
}
if (cc->InputSidePackets().HasTag("OUTPUT_INDEX_FILENAME")) {
if (cc->InputSidePackets().HasTag(kOutputIndexFilenameTag)) {
write_index_ = true;
}
if (cc->InputSidePackets().HasTag("FRAME_ALIGNMENT")) {
frame_alignment_ = cc->InputSidePackets().Tag("FRAME_ALIGNMENT").Get<int>();
if (cc->InputSidePackets().HasTag(kFrameAlignmentTag)) {
frame_alignment_ =
cc->InputSidePackets().Tag(kFrameAlignmentTag).Get<int>();
}
return absl::OkStatus();
@@ -218,16 +234,16 @@ absl::Status BoxDetectorCalculator::Process(CalculatorContext* cc) {
const int64 timestamp_msec = timestamp.Value() / 1000;
InputStream* cancel_object_id_stream =
cc->Inputs().HasTag("CANCEL_OBJECT_ID")
? &(cc->Inputs().Tag("CANCEL_OBJECT_ID"))
cc->Inputs().HasTag(kCancelObjectIdTag)
? &(cc->Inputs().Tag(kCancelObjectIdTag))
: nullptr;
if (cancel_object_id_stream && !cancel_object_id_stream->IsEmpty()) {
const int cancel_object_id = cancel_object_id_stream->Get<int>();
box_detector_->CancelBoxDetection(cancel_object_id);
}
InputStream* add_index_stream = cc->Inputs().HasTag("ADD_INDEX")
? &(cc->Inputs().Tag("ADD_INDEX"))
InputStream* add_index_stream = cc->Inputs().HasTag(kAddIndexTag)
? &(cc->Inputs().Tag(kAddIndexTag))
: nullptr;
if (add_index_stream && !add_index_stream->IsEmpty()) {
BoxDetectorIndex predefined_index;
@@ -238,8 +254,8 @@ absl::Status BoxDetectorCalculator::Process(CalculatorContext* cc) {
box_detector_->AddBoxDetectorIndex(predefined_index);
}
InputStream* reacq_switch_stream = cc->Inputs().HasTag("REACQ_SWITCH")
? &(cc->Inputs().Tag("REACQ_SWITCH"))
InputStream* reacq_switch_stream = cc->Inputs().HasTag(kReacqSwitchTag)
? &(cc->Inputs().Tag(kReacqSwitchTag))
: nullptr;
if (reacq_switch_stream && !reacq_switch_stream->IsEmpty()) {
detector_switch_ = reacq_switch_stream->Get<bool>();
@@ -249,16 +265,16 @@ absl::Status BoxDetectorCalculator::Process(CalculatorContext* cc) {
return absl::OkStatus();
}
InputStream* track_stream = cc->Inputs().HasTag("TRACKING")
? &(cc->Inputs().Tag("TRACKING"))
InputStream* track_stream = cc->Inputs().HasTag(kTrackingTag)
? &(cc->Inputs().Tag(kTrackingTag))
: nullptr;
InputStream* video_stream =
cc->Inputs().HasTag("VIDEO") ? &(cc->Inputs().Tag("VIDEO")) : nullptr;
InputStream* feature_stream = cc->Inputs().HasTag("FEATURES")
? &(cc->Inputs().Tag("FEATURES"))
cc->Inputs().HasTag(kVideoTag) ? &(cc->Inputs().Tag(kVideoTag)) : nullptr;
InputStream* feature_stream = cc->Inputs().HasTag(kFeaturesTag)
? &(cc->Inputs().Tag(kFeaturesTag))
: nullptr;
InputStream* descriptor_stream = cc->Inputs().HasTag("DESCRIPTORS")
? &(cc->Inputs().Tag("DESCRIPTORS"))
InputStream* descriptor_stream = cc->Inputs().HasTag(kDescriptorsTag)
? &(cc->Inputs().Tag(kDescriptorsTag))
: nullptr;
CHECK(track_stream != nullptr || video_stream != nullptr ||
@@ -266,9 +282,10 @@ absl::Status BoxDetectorCalculator::Process(CalculatorContext* cc) {
<< "One and only one of {tracking_data, input image frame, "
"feature/descriptor} need to be valid.";
InputStream* tracked_boxes_stream = cc->Inputs().HasTag("TRACKED_BOXES")
? &(cc->Inputs().Tag("TRACKED_BOXES"))
: nullptr;
InputStream* tracked_boxes_stream =
cc->Inputs().HasTag(kTrackedBoxesTag)
? &(cc->Inputs().Tag(kTrackedBoxesTag))
: nullptr;
std::unique_ptr<TimedBoxProtoList> detected_boxes(new TimedBoxProtoList());
if (track_stream != nullptr) {
@@ -309,7 +326,7 @@ absl::Status BoxDetectorCalculator::Process(CalculatorContext* cc) {
}
const auto& image_size =
cc->Inputs().Tag("IMAGE_SIZE").Get<std::pair<int, int>>();
cc->Inputs().Tag(kImageSizeTag).Get<std::pair<int, int>>();
float inv_scale = 1.0f / std::max(image_size.first, image_size.second);
TimedBoxProtoList tracked_boxes;
@@ -359,7 +376,7 @@ absl::Status BoxDetectorCalculator::Process(CalculatorContext* cc) {
detected_boxes.get());
}
if (cc->Outputs().HasTag("VIZ")) {
if (cc->Outputs().HasTag(kVizTag)) {
cv::Mat viz_view;
std::unique_ptr<ImageFrame> viz_frame;
if (video_stream != nullptr && !video_stream->IsEmpty()) {
@@ -370,11 +387,11 @@ absl::Status BoxDetectorCalculator::Process(CalculatorContext* cc) {
for (const auto& box : detected_boxes->box()) {
RenderBox(box, &viz_view);
}
cc->Outputs().Tag("VIZ").Add(viz_frame.release(), timestamp);
cc->Outputs().Tag(kVizTag).Add(viz_frame.release(), timestamp);
}
if (cc->Outputs().HasTag("BOXES")) {
cc->Outputs().Tag("BOXES").Add(detected_boxes.release(), timestamp);
if (cc->Outputs().HasTag(kBoxesTag)) {
cc->Outputs().Tag(kBoxesTag).Add(detected_boxes.release(), timestamp);
}
return absl::OkStatus();
@@ -384,7 +401,7 @@ absl::Status BoxDetectorCalculator::Close(CalculatorContext* cc) {
if (write_index_) {
BoxDetectorIndex index = box_detector_->ObtainBoxDetectorIndex();
MEDIAPIPE_CHECK_OK(mediapipe::file::SetContents(
cc->InputSidePackets().Tag("OUTPUT_INDEX_FILENAME").Get<std::string>(),
cc->InputSidePackets().Tag(kOutputIndexFilenameTag).Get<std::string>(),
index.SerializeAsString()));
}
return absl::OkStatus();
mediapipe/calculators/video/box_tracker_calculator.cc
+86 -69
View File
@@ -293,6 +293,22 @@ const int BoxTrackerCalculator::kMotionBoxPathMinQueueSize = 2;
namespace {
constexpr char kCacheDirTag[] = "CACHE_DIR";
constexpr char kInitialPosTag[] = "INITIAL_POS";
constexpr char kRaBoxesTag[] = "RA_BOXES";
constexpr char kBoxesTag[] = "BOXES";
constexpr char kVizTag[] = "VIZ";
constexpr char kRaTrackProtoStringTag[] = "RA_TRACK_PROTO_STRING";
constexpr char kRaTrackTag[] = "RA_TRACK";
constexpr char kCancelObjectIdTag[] = "CANCEL_OBJECT_ID";
constexpr char kRestartPosTag[] = "RESTART_POS";
constexpr char kStartPosProtoStringTag[] = "START_POS_PROTO_STRING";
constexpr char kStartPosTag[] = "START_POS";
constexpr char kStartTag[] = "START";
constexpr char kVideoTag[] = "VIDEO";
constexpr char kTrackTimeTag[] = "TRACK_TIME";
constexpr char kTrackingTag[] = "TRACKING";
// Convert box position according to rotation angle in degrees.
void ConvertCoordinateForRotation(float in_top, float in_left, float in_bottom,
float in_right, int rotation, float* out_top,
@@ -374,78 +390,78 @@ void AddStateToPath(const MotionBoxState& state, int64 time_msec,
} // namespace.
absl::Status BoxTrackerCalculator::GetContract(CalculatorContract* cc) {
if (cc->Inputs().HasTag("TRACKING")) {
cc->Inputs().Tag("TRACKING").Set<TrackingData>();
if (cc->Inputs().HasTag(kTrackingTag)) {
cc->Inputs().Tag(kTrackingTag).Set<TrackingData>();
}
if (cc->Inputs().HasTag("TRACK_TIME")) {
RET_CHECK(cc->Inputs().HasTag("TRACKING"))
if (cc->Inputs().HasTag(kTrackTimeTag)) {
RET_CHECK(cc->Inputs().HasTag(kTrackingTag))
<< "TRACK_TIME needs TRACKING input";
cc->Inputs().Tag("TRACK_TIME").SetAny();
cc->Inputs().Tag(kTrackTimeTag).SetAny();
}
if (cc->Inputs().HasTag("VIDEO")) {
cc->Inputs().Tag("VIDEO").Set<ImageFrame>();
if (cc->Inputs().HasTag(kVideoTag)) {
cc->Inputs().Tag(kVideoTag).Set<ImageFrame>();
}
if (cc->Inputs().HasTag("START")) {
if (cc->Inputs().HasTag(kStartTag)) {
// Actual packet content does not matter.
cc->Inputs().Tag("START").SetAny();
cc->Inputs().Tag(kStartTag).SetAny();
}
if (cc->Inputs().HasTag("START_POS")) {
cc->Inputs().Tag("START_POS").Set<TimedBoxProtoList>();
if (cc->Inputs().HasTag(kStartPosTag)) {
cc->Inputs().Tag(kStartPosTag).Set<TimedBoxProtoList>();
}
if (cc->Inputs().HasTag("START_POS_PROTO_STRING")) {
cc->Inputs().Tag("START_POS_PROTO_STRING").Set<std::string>();
if (cc->Inputs().HasTag(kStartPosProtoStringTag)) {
cc->Inputs().Tag(kStartPosProtoStringTag).Set<std::string>();
}
if (cc->Inputs().HasTag("RESTART_POS")) {
cc->Inputs().Tag("RESTART_POS").Set<TimedBoxProtoList>();
if (cc->Inputs().HasTag(kRestartPosTag)) {
cc->Inputs().Tag(kRestartPosTag).Set<TimedBoxProtoList>();
}
if (cc->Inputs().HasTag("CANCEL_OBJECT_ID")) {
cc->Inputs().Tag("CANCEL_OBJECT_ID").Set<int>();
if (cc->Inputs().HasTag(kCancelObjectIdTag)) {
cc->Inputs().Tag(kCancelObjectIdTag).Set<int>();
}
if (cc->Inputs().HasTag("RA_TRACK")) {
cc->Inputs().Tag("RA_TRACK").Set<TimedBoxProtoList>();
if (cc->Inputs().HasTag(kRaTrackTag)) {
cc->Inputs().Tag(kRaTrackTag).Set<TimedBoxProtoList>();
}
if (cc->Inputs().HasTag("RA_TRACK_PROTO_STRING")) {
cc->Inputs().Tag("RA_TRACK_PROTO_STRING").Set<std::string>();
if (cc->Inputs().HasTag(kRaTrackProtoStringTag)) {
cc->Inputs().Tag(kRaTrackProtoStringTag).Set<std::string>();
}
if (cc->Outputs().HasTag("VIZ")) {
RET_CHECK(cc->Inputs().HasTag("VIDEO"))
if (cc->Outputs().HasTag(kVizTag)) {
RET_CHECK(cc->Inputs().HasTag(kVideoTag))
<< "Output stream VIZ requires VIDEO to be present.";
cc->Outputs().Tag("VIZ").Set<ImageFrame>();
cc->Outputs().Tag(kVizTag).Set<ImageFrame>();
}
if (cc->Outputs().HasTag("BOXES")) {
cc->Outputs().Tag("BOXES").Set<TimedBoxProtoList>();
if (cc->Outputs().HasTag(kBoxesTag)) {
cc->Outputs().Tag(kBoxesTag).Set<TimedBoxProtoList>();
}
if (cc->Outputs().HasTag("RA_BOXES")) {
cc->Outputs().Tag("RA_BOXES").Set<TimedBoxProtoList>();
if (cc->Outputs().HasTag(kRaBoxesTag)) {
cc->Outputs().Tag(kRaBoxesTag).Set<TimedBoxProtoList>();
}
#if defined(__ANDROID__) || defined(__APPLE__) || defined(__EMSCRIPTEN__)
RET_CHECK(!cc->InputSidePackets().HasTag("INITIAL_POS"))
RET_CHECK(!cc->InputSidePackets().HasTag(kInitialPosTag))
<< "Unsupported on mobile";
#else
if (cc->InputSidePackets().HasTag("INITIAL_POS")) {
cc->InputSidePackets().Tag("INITIAL_POS").Set<std::string>();
if (cc->InputSidePackets().HasTag(kInitialPosTag)) {
cc->InputSidePackets().Tag(kInitialPosTag).Set<std::string>();
}
#endif // defined(__ANDROID__) || defined(__APPLE__) || defined(__EMSCRIPTEN__)
if (cc->InputSidePackets().HasTag("CACHE_DIR")) {
cc->InputSidePackets().Tag("CACHE_DIR").Set<std::string>();
if (cc->InputSidePackets().HasTag(kCacheDirTag)) {
cc->InputSidePackets().Tag(kCacheDirTag).Set<std::string>();
}
RET_CHECK(cc->Inputs().HasTag("TRACKING") !=
cc->InputSidePackets().HasTag("CACHE_DIR"))
RET_CHECK(cc->Inputs().HasTag(kTrackingTag) !=
cc->InputSidePackets().HasTag(kCacheDirTag))
<< "Either TRACKING or CACHE_DIR needs to be specified.";
if (cc->InputSidePackets().HasTag(kOptionsTag)) {
@@ -459,7 +475,7 @@ absl::Status BoxTrackerCalculator::Open(CalculatorContext* cc) {
options_ = tool::RetrieveOptions(cc->Options<BoxTrackerCalculatorOptions>(),
cc->InputSidePackets(), kOptionsTag);
RET_CHECK(!cc->InputSidePackets().HasTag("INITIAL_POS") ||
RET_CHECK(!cc->InputSidePackets().HasTag(kInitialPosTag) ||
!options_.has_initial_position())
<< "Can not specify initial position as side packet and via options";
@@ -468,11 +484,11 @@ absl::Status BoxTrackerCalculator::Open(CalculatorContext* cc) {
}
#if !defined(__ANDROID__) && !defined(__APPLE__) && !defined(__EMSCRIPTEN__)
if (cc->InputSidePackets().HasTag("INITIAL_POS")) {
if (cc->InputSidePackets().HasTag(kInitialPosTag)) {
LOG(INFO) << "Parsing: "
<< cc->InputSidePackets().Tag("INITIAL_POS").Get<std::string>();
<< cc->InputSidePackets().Tag(kInitialPosTag).Get<std::string>();
initial_pos_ = ParseTextProtoOrDie<TimedBoxProtoList>(
cc->InputSidePackets().Tag("INITIAL_POS").Get<std::string>());
cc->InputSidePackets().Tag(kInitialPosTag).Get<std::string>());
}
#endif // !defined(__ANDROID__) && !defined(__APPLE__) &&
// !defined(__EMSCRIPTEN__)
@@ -484,10 +500,11 @@ absl::Status BoxTrackerCalculator::Open(CalculatorContext* cc) {
}
visualize_tracking_data_ =
options_.visualize_tracking_data() && cc->Outputs().HasTag("VIZ");
visualize_state_ = options_.visualize_state() && cc->Outputs().HasTag("VIZ");
options_.visualize_tracking_data() && cc->Outputs().HasTag(kVizTag);
visualize_state_ =
options_.visualize_state() && cc->Outputs().HasTag(kVizTag);
visualize_internal_state_ =
options_.visualize_internal_state() && cc->Outputs().HasTag("VIZ");
options_.visualize_internal_state() && cc->Outputs().HasTag(kVizTag);
// Force recording of internal state for rendering.
if (visualize_internal_state_) {
@@ -500,8 +517,8 @@ absl::Status BoxTrackerCalculator::Open(CalculatorContext* cc) {
options_.mutable_tracker_options()->set_record_path_states(true);
}
if (cc->InputSidePackets().HasTag("CACHE_DIR")) {
cache_dir_ = cc->InputSidePackets().Tag("CACHE_DIR").Get<std::string>();
if (cc->InputSidePackets().HasTag(kCacheDirTag)) {
cache_dir_ = cc->InputSidePackets().Tag(kCacheDirTag).Get<std::string>();
RET_CHECK(!cache_dir_.empty());
box_tracker_.reset(new BoxTracker(cache_dir_, options_.tracker_options()));
} else {
@@ -511,7 +528,7 @@ absl::Status BoxTrackerCalculator::Open(CalculatorContext* cc) {
}
if (options_.streaming_track_data_cache_size() > 0) {
RET_CHECK(!cc->InputSidePackets().HasTag("CACHE_DIR"))
RET_CHECK(!cc->InputSidePackets().HasTag(kCacheDirTag))
<< "Streaming mode not compatible with cache dir.";
}
@@ -533,11 +550,11 @@ absl::Status BoxTrackerCalculator::Process(CalculatorContext* cc) {
return absl::OkStatus();
}
InputStream* track_stream = cc->Inputs().HasTag("TRACKING")
? &(cc->Inputs().Tag("TRACKING"))
InputStream* track_stream = cc->Inputs().HasTag(kTrackingTag)
? &(cc->Inputs().Tag(kTrackingTag))
: nullptr;
InputStream* track_time_stream = cc->Inputs().HasTag("TRACK_TIME")
? &(cc->Inputs().Tag("TRACK_TIME"))
InputStream* track_time_stream = cc->Inputs().HasTag(kTrackTimeTag)
? &(cc->Inputs().Tag(kTrackTimeTag))
: nullptr;
// Cache tracking data if possible.
@@ -562,8 +579,8 @@ absl::Status BoxTrackerCalculator::Process(CalculatorContext* cc) {
}
}
InputStream* start_pos_stream = cc->Inputs().HasTag("START_POS")
? &(cc->Inputs().Tag("START_POS"))
InputStream* start_pos_stream = cc->Inputs().HasTag(kStartPosTag)
? &(cc->Inputs().Tag(kStartPosTag))
: nullptr;
MotionBoxMap fast_forward_boxes;
@@ -575,8 +592,8 @@ absl::Status BoxTrackerCalculator::Process(CalculatorContext* cc) {
}
InputStream* start_pos_proto_string_stream =
cc->Inputs().HasTag("START_POS_PROTO_STRING")
? &(cc->Inputs().Tag("START_POS_PROTO_STRING"))
cc->Inputs().HasTag(kStartPosProtoStringTag)
? &(cc->Inputs().Tag(kStartPosProtoStringTag))
: nullptr;
if (start_pos_stream == nullptr || start_pos_stream->IsEmpty()) {
if (start_pos_proto_string_stream &&
@@ -589,8 +606,8 @@ absl::Status BoxTrackerCalculator::Process(CalculatorContext* cc) {
}
}
InputStream* restart_pos_stream = cc->Inputs().HasTag("RESTART_POS")
? &(cc->Inputs().Tag("RESTART_POS"))
InputStream* restart_pos_stream = cc->Inputs().HasTag(kRestartPosTag)
? &(cc->Inputs().Tag(kRestartPosTag))
: nullptr;
if (restart_pos_stream && !restart_pos_stream->IsEmpty()) {
@@ -600,8 +617,8 @@ absl::Status BoxTrackerCalculator::Process(CalculatorContext* cc) {
}
InputStream* cancel_object_id_stream =
cc->Inputs().HasTag("CANCEL_OBJECT_ID")
? &(cc->Inputs().Tag("CANCEL_OBJECT_ID"))
cc->Inputs().HasTag(kCancelObjectIdTag)
? &(cc->Inputs().Tag(kCancelObjectIdTag))
: nullptr;
if (cancel_object_id_stream && !cancel_object_id_stream->IsEmpty()) {
const int cancel_object_id = cancel_object_id_stream->Get<int>();
@@ -616,8 +633,8 @@ absl::Status BoxTrackerCalculator::Process(CalculatorContext* cc) {
TrackingData track_data_to_render;
if (cc->Outputs().HasTag("VIZ")) {
InputStream* video_stream = &(cc->Inputs().Tag("VIDEO"));
if (cc->Outputs().HasTag(kVizTag)) {
InputStream* video_stream = &(cc->Inputs().Tag(kVideoTag));
if (!video_stream->IsEmpty()) {
input_view = formats::MatView(&video_stream->Get<ImageFrame>());
@@ -745,7 +762,7 @@ absl::Status BoxTrackerCalculator::Process(CalculatorContext* cc) {
++frame_num_since_reset_;
// Generate results for queued up request.
if (cc->Outputs().HasTag("BOXES") && !queued_track_requests_.empty()) {
if (cc->Outputs().HasTag(kBoxesTag) && !queued_track_requests_.empty()) {
for (int j = 0; j < queued_track_requests_.size(); ++j) {
const Timestamp& past_time = queued_track_requests_[j];
RET_CHECK(past_time.Value() < timestamp.Value())
@@ -770,7 +787,7 @@ absl::Status BoxTrackerCalculator::Process(CalculatorContext* cc) {
}
// Output for every time.
cc->Outputs().Tag("BOXES").Add(past_box_list.release(), past_time);
cc->Outputs().Tag(kBoxesTag).Add(past_box_list.release(), past_time);
}
queued_track_requests_.clear();
@@ -845,8 +862,8 @@ absl::Status BoxTrackerCalculator::Process(CalculatorContext* cc) {
}
// Handle random access track requests.
InputStream* ra_track_stream = cc->Inputs().HasTag("RA_TRACK")
? &(cc->Inputs().Tag("RA_TRACK"))
InputStream* ra_track_stream = cc->Inputs().HasTag(kRaTrackTag)
? &(cc->Inputs().Tag(kRaTrackTag))
: nullptr;
if (ra_track_stream && !ra_track_stream->IsEmpty()) {
@@ -861,8 +878,8 @@ absl::Status BoxTrackerCalculator::Process(CalculatorContext* cc) {
}
InputStream* ra_track_proto_string_stream =
cc->Inputs().HasTag("RA_TRACK_PROTO_STRING")
? &(cc->Inputs().Tag("RA_TRACK_PROTO_STRING"))
cc->Inputs().HasTag(kRaTrackProtoStringTag)
? &(cc->Inputs().Tag(kRaTrackProtoStringTag))
: nullptr;
if (ra_track_stream == nullptr || ra_track_stream->IsEmpty()) {
if (ra_track_proto_string_stream &&
@@ -881,15 +898,15 @@ absl::Status BoxTrackerCalculator::Process(CalculatorContext* cc) {
// Always output in batch, only output in streaming if tracking data
// is present (might be in fast forward mode instead).
if (cc->Outputs().HasTag("BOXES") &&
if (cc->Outputs().HasTag(kBoxesTag) &&
(box_tracker_ || !track_stream->IsEmpty())) {
std::unique_ptr<TimedBoxProtoList> boxes(new TimedBoxProtoList());
*boxes = std::move(box_track_list);
cc->Outputs().Tag("BOXES").Add(boxes.release(), timestamp);
cc->Outputs().Tag(kBoxesTag).Add(boxes.release(), timestamp);
}
if (viz_frame) {
cc->Outputs().Tag("VIZ").Add(viz_frame.release(), timestamp);
cc->Outputs().Tag(kVizTag).Add(viz_frame.release(), timestamp);
}
return absl::OkStatus();
@@ -1001,7 +1018,7 @@ void BoxTrackerCalculator::OutputRandomAccessTrack(
}
cc->Outputs()
.Tag("RA_BOXES")
.Tag(kRaBoxesTag)
.Add(result_list.release(), cc->InputTimestamp());
}
mediapipe/calculators/video/flow_packager_calculator.cc
+34 -27
View File
@@ -29,6 +29,13 @@
namespace mediapipe {
constexpr char kCacheDirTag[] = "CACHE_DIR";
constexpr char kCompleteTag[] = "COMPLETE";
constexpr char kTrackingChunkTag[] = "TRACKING_CHUNK";
constexpr char kTrackingTag[] = "TRACKING";
constexpr char kCameraTag[] = "CAMERA";
constexpr char kFlowTag[] = "FLOW";
using mediapipe::CameraMotion;
using mediapipe::FlowPackager;
using mediapipe::RegionFlowFeatureList;
@@ -91,27 +98,27 @@ class FlowPackagerCalculator : public CalculatorBase {
REGISTER_CALCULATOR(FlowPackagerCalculator);
absl::Status FlowPackagerCalculator::GetContract(CalculatorContract* cc) {
if (!cc->Inputs().HasTag("FLOW")) {
if (!cc->Inputs().HasTag(kFlowTag)) {
return tool::StatusFail("No input flow was specified.");
}
cc->Inputs().Tag("FLOW").Set<RegionFlowFeatureList>();
cc->Inputs().Tag(kFlowTag).Set<RegionFlowFeatureList>();
if (cc->Inputs().HasTag("CAMERA")) {
cc->Inputs().Tag("CAMERA").Set<CameraMotion>();
if (cc->Inputs().HasTag(kCameraTag)) {
cc->Inputs().Tag(kCameraTag).Set<CameraMotion>();
}
if (cc->Outputs().HasTag("TRACKING")) {
cc->Outputs().Tag("TRACKING").Set<TrackingData>();
if (cc->Outputs().HasTag(kTrackingTag)) {
cc->Outputs().Tag(kTrackingTag).Set<TrackingData>();
}
if (cc->Outputs().HasTag("TRACKING_CHUNK")) {
cc->Outputs().Tag("TRACKING_CHUNK").Set<TrackingDataChunk>();
if (cc->Outputs().HasTag(kTrackingChunkTag)) {
cc->Outputs().Tag(kTrackingChunkTag).Set<TrackingDataChunk>();
}
if (cc->Outputs().HasTag("COMPLETE")) {
cc->Outputs().Tag("COMPLETE").Set<bool>();
if (cc->Outputs().HasTag(kCompleteTag)) {
cc->Outputs().Tag(kCompleteTag).Set<bool>();
}
if (cc->InputSidePackets().HasTag("CACHE_DIR")) {
cc->InputSidePackets().Tag("CACHE_DIR").Set<std::string>();
if (cc->InputSidePackets().HasTag(kCacheDirTag)) {
cc->InputSidePackets().Tag(kCacheDirTag).Set<std::string>();
}
return absl::OkStatus();
@@ -122,24 +129,24 @@ absl::Status FlowPackagerCalculator::Open(CalculatorContext* cc) {
flow_packager_.reset(new FlowPackager(options_.flow_packager_options()));
use_caching_ = cc->InputSidePackets().HasTag("CACHE_DIR");
build_chunk_ = use_caching_ || cc->Outputs().HasTag("TRACKING_CHUNK");
use_caching_ = cc->InputSidePackets().HasTag(kCacheDirTag);
build_chunk_ = use_caching_ || cc->Outputs().HasTag(kTrackingChunkTag);
if (use_caching_) {
cache_dir_ = cc->InputSidePackets().Tag("CACHE_DIR").Get<std::string>();
cache_dir_ = cc->InputSidePackets().Tag(kCacheDirTag).Get<std::string>();
}
return absl::OkStatus();
}
absl::Status FlowPackagerCalculator::Process(CalculatorContext* cc) {
InputStream* flow_stream = &(cc->Inputs().Tag("FLOW"));
InputStream* flow_stream = &(cc->Inputs().Tag(kFlowTag));
const RegionFlowFeatureList& flow = flow_stream->Get<RegionFlowFeatureList>();
const Timestamp timestamp = flow_stream->Value().Timestamp();
const CameraMotion* camera_motion = nullptr;
if (cc->Inputs().HasTag("CAMERA")) {
InputStream* camera_stream = &(cc->Inputs().Tag("CAMERA"));
if (cc->Inputs().HasTag(kCameraTag)) {
InputStream* camera_stream = &(cc->Inputs().Tag(kCameraTag));
camera_motion = &camera_stream->Get<CameraMotion>();
}
@@ -161,7 +168,7 @@ absl::Status FlowPackagerCalculator::Process(CalculatorContext* cc) {
if (frame_idx_ > 0) {
item->set_prev_timestamp_usec(prev_timestamp_.Value());
}
if (cc->Outputs().HasTag("TRACKING")) {
if (cc->Outputs().HasTag(kTrackingTag)) {
// Need to copy as output is requested.
*item->mutable_tracking_data() = *tracking_data;
} else {
@@ -172,9 +179,9 @@ absl::Status FlowPackagerCalculator::Process(CalculatorContext* cc) {
options_.caching_chunk_size_msec() * (chunk_idx_ + 1);
if (timestamp.Value() / 1000 >= next_chunk_msec) {
if (cc->Outputs().HasTag("TRACKING_CHUNK")) {
if (cc->Outputs().HasTag(kTrackingChunkTag)) {
cc->Outputs()
.Tag("TRACKING_CHUNK")
.Tag(kTrackingChunkTag)
.Add(new TrackingDataChunk(tracking_chunk_),
Timestamp(tracking_chunk_.item(0).timestamp_usec()));
}
@@ -185,9 +192,9 @@ absl::Status FlowPackagerCalculator::Process(CalculatorContext* cc) {
}
}
if (cc->Outputs().HasTag("TRACKING")) {
if (cc->Outputs().HasTag(kTrackingTag)) {
cc->Outputs()
.Tag("TRACKING")
.Tag(kTrackingTag)
.Add(tracking_data.release(), flow_stream->Value().Timestamp());
}
@@ -199,9 +206,9 @@ absl::Status FlowPackagerCalculator::Process(CalculatorContext* cc) {
absl::Status FlowPackagerCalculator::Close(CalculatorContext* cc) {
if (frame_idx_ > 0) {
tracking_chunk_.set_last_chunk(true);
if (cc->Outputs().HasTag("TRACKING_CHUNK")) {
if (cc->Outputs().HasTag(kTrackingChunkTag)) {
cc->Outputs()
.Tag("TRACKING_CHUNK")
.Tag(kTrackingChunkTag)
.Add(new TrackingDataChunk(tracking_chunk_),
Timestamp(tracking_chunk_.item(0).timestamp_usec()));
}
@@ -211,8 +218,8 @@ absl::Status FlowPackagerCalculator::Close(CalculatorContext* cc) {
}
}
if (cc->Outputs().HasTag("COMPLETE")) {
cc->Outputs().Tag("COMPLETE").Add(new bool(true), Timestamp::PreStream());
if (cc->Outputs().HasTag(kCompleteTag)) {
cc->Outputs().Tag(kCompleteTag).Add(new bool(true), Timestamp::PreStream());
}
return absl::OkStatus();
mediapipe/calculators/video/motion_analysis_calculator.cc
+93 -74
View File
@@ -38,6 +38,18 @@
namespace mediapipe {
constexpr char kDownsampleTag[] = "DOWNSAMPLE";
constexpr char kCsvFileTag[] = "CSV_FILE";
constexpr char kGrayVideoOutTag[] = "GRAY_VIDEO_OUT";
constexpr char kVideoOutTag[] = "VIDEO_OUT";
constexpr char kDenseFgTag[] = "DENSE_FG";
constexpr char kVizTag[] = "VIZ";
constexpr char kSaliencyTag[] = "SALIENCY";
constexpr char kCameraTag[] = "CAMERA";
constexpr char kFlowTag[] = "FLOW";
constexpr char kSelectionTag[] = "SELECTION";
constexpr char kVideoTag[] = "VIDEO";
using mediapipe::AffineAdapter;
using mediapipe::CameraMotion;
using mediapipe::FrameSelectionResult;
@@ -190,55 +202,56 @@ class MotionAnalysisCalculator : public CalculatorBase {
REGISTER_CALCULATOR(MotionAnalysisCalculator);
absl::Status MotionAnalysisCalculator::GetContract(CalculatorContract* cc) {
if (cc->Inputs().HasTag("VIDEO")) {
cc->Inputs().Tag("VIDEO").Set<ImageFrame>();
if (cc->Inputs().HasTag(kVideoTag)) {
cc->Inputs().Tag(kVideoTag).Set<ImageFrame>();
}
// Optional input stream from frame selection calculator.
if (cc->Inputs().HasTag("SELECTION")) {
cc->Inputs().Tag("SELECTION").Set<FrameSelectionResult>();
if (cc->Inputs().HasTag(kSelectionTag)) {
cc->Inputs().Tag(kSelectionTag).Set<FrameSelectionResult>();
}
RET_CHECK(cc->Inputs().HasTag("VIDEO") || cc->Inputs().HasTag("SELECTION"))
RET_CHECK(cc->Inputs().HasTag(kVideoTag) ||
cc->Inputs().HasTag(kSelectionTag))
<< "Either VIDEO, SELECTION must be specified.";
if (cc->Outputs().HasTag("FLOW")) {
cc->Outputs().Tag("FLOW").Set<RegionFlowFeatureList>();
if (cc->Outputs().HasTag(kFlowTag)) {
cc->Outputs().Tag(kFlowTag).Set<RegionFlowFeatureList>();
}
if (cc->Outputs().HasTag("CAMERA")) {
cc->Outputs().Tag("CAMERA").Set<CameraMotion>();
if (cc->Outputs().HasTag(kCameraTag)) {
cc->Outputs().Tag(kCameraTag).Set<CameraMotion>();
}
if (cc->Outputs().HasTag("SALIENCY")) {
cc->Outputs().Tag("SALIENCY").Set<SalientPointFrame>();
if (cc->Outputs().HasTag(kSaliencyTag)) {
cc->Outputs().Tag(kSaliencyTag).Set<SalientPointFrame>();
}
if (cc->Outputs().HasTag("VIZ")) {
cc->Outputs().Tag("VIZ").Set<ImageFrame>();
if (cc->Outputs().HasTag(kVizTag)) {
cc->Outputs().Tag(kVizTag).Set<ImageFrame>();
}
if (cc->Outputs().HasTag("DENSE_FG")) {
cc->Outputs().Tag("DENSE_FG").Set<ImageFrame>();
if (cc->Outputs().HasTag(kDenseFgTag)) {
cc->Outputs().Tag(kDenseFgTag).Set<ImageFrame>();
}
if (cc->Outputs().HasTag("VIDEO_OUT")) {
cc->Outputs().Tag("VIDEO_OUT").Set<ImageFrame>();
if (cc->Outputs().HasTag(kVideoOutTag)) {
cc->Outputs().Tag(kVideoOutTag).Set<ImageFrame>();
}
if (cc->Outputs().HasTag("GRAY_VIDEO_OUT")) {
if (cc->Outputs().HasTag(kGrayVideoOutTag)) {
// We only output grayscale video if we're actually performing full region-
// flow analysis on the video.
RET_CHECK(cc->Inputs().HasTag("VIDEO") &&
!cc->Inputs().HasTag("SELECTION"));
cc->Outputs().Tag("GRAY_VIDEO_OUT").Set<ImageFrame>();
RET_CHECK(cc->Inputs().HasTag(kVideoTag) &&
!cc->Inputs().HasTag(kSelectionTag));
cc->Outputs().Tag(kGrayVideoOutTag).Set<ImageFrame>();
}
if (cc->InputSidePackets().HasTag("CSV_FILE")) {
cc->InputSidePackets().Tag("CSV_FILE").Set<std::string>();
if (cc->InputSidePackets().HasTag(kCsvFileTag)) {
cc->InputSidePackets().Tag(kCsvFileTag).Set<std::string>();
}
if (cc->InputSidePackets().HasTag("DOWNSAMPLE")) {
cc->InputSidePackets().Tag("DOWNSAMPLE").Set<float>();
if (cc->InputSidePackets().HasTag(kDownsampleTag)) {
cc->InputSidePackets().Tag(kDownsampleTag).Set<float>();
}
if (cc->InputSidePackets().HasTag(kOptionsTag)) {
@@ -253,16 +266,16 @@ absl::Status MotionAnalysisCalculator::Open(CalculatorContext* cc) {
tool::RetrieveOptions(cc->Options<MotionAnalysisCalculatorOptions>(),
cc->InputSidePackets(), kOptionsTag);
video_input_ = cc->Inputs().HasTag("VIDEO");
selection_input_ = cc->Inputs().HasTag("SELECTION");
region_flow_feature_output_ = cc->Outputs().HasTag("FLOW");
camera_motion_output_ = cc->Outputs().HasTag("CAMERA");
saliency_output_ = cc->Outputs().HasTag("SALIENCY");
visualize_output_ = cc->Outputs().HasTag("VIZ");
dense_foreground_output_ = cc->Outputs().HasTag("DENSE_FG");
video_output_ = cc->Outputs().HasTag("VIDEO_OUT");
grayscale_output_ = cc->Outputs().HasTag("GRAY_VIDEO_OUT");
csv_file_input_ = cc->InputSidePackets().HasTag("CSV_FILE");
video_input_ = cc->Inputs().HasTag(kVideoTag);
selection_input_ = cc->Inputs().HasTag(kSelectionTag);
region_flow_feature_output_ = cc->Outputs().HasTag(kFlowTag);
camera_motion_output_ = cc->Outputs().HasTag(kCameraTag);
saliency_output_ = cc->Outputs().HasTag(kSaliencyTag);
visualize_output_ = cc->Outputs().HasTag(kVizTag);
dense_foreground_output_ = cc->Outputs().HasTag(kDenseFgTag);
video_output_ = cc->Outputs().HasTag(kVideoOutTag);
grayscale_output_ = cc->Outputs().HasTag(kGrayVideoOutTag);
csv_file_input_ = cc->InputSidePackets().HasTag(kCsvFileTag);
hybrid_meta_analysis_ = options_.meta_analysis() ==
MotionAnalysisCalculatorOptions::META_ANALYSIS_HYBRID;
@@ -310,7 +323,7 @@ absl::Status MotionAnalysisCalculator::Open(CalculatorContext* cc) {
if (csv_file_input_) {
// Read from file and parse.
const std::string filename =
cc->InputSidePackets().Tag("CSV_FILE").Get<std::string>();
cc->InputSidePackets().Tag(kCsvFileTag).Get<std::string>();
std::string file_contents;
std::ifstream input_file(filename, std::ios::in);
@@ -327,11 +340,12 @@ absl::Status MotionAnalysisCalculator::Open(CalculatorContext* cc) {
// Get video header from video or selection input if present.
const VideoHeader* video_header = nullptr;
if (video_input_ && !cc->Inputs().Tag("VIDEO").Header().IsEmpty()) {
video_header = &(cc->Inputs().Tag("VIDEO").Header().Get<VideoHeader>());
if (video_input_ && !cc->Inputs().Tag(kVideoTag).Header().IsEmpty()) {
video_header = &(cc->Inputs().Tag(kVideoTag).Header().Get<VideoHeader>());
} else if (selection_input_ &&
!cc->Inputs().Tag("SELECTION").Header().IsEmpty()) {
video_header = &(cc->Inputs().Tag("SELECTION").Header().Get<VideoHeader>());
!cc->Inputs().Tag(kSelectionTag).Header().IsEmpty()) {
video_header =
&(cc->Inputs().Tag(kSelectionTag).Header().Get<VideoHeader>());
} else {
LOG(WARNING) << "No input video header found. Downstream calculators "
"expecting video headers are likely to fail.";
@@ -339,7 +353,7 @@ absl::Status MotionAnalysisCalculator::Open(CalculatorContext* cc) {
with_saliency_ = options_.analysis_options().compute_motion_saliency();
// Force computation of saliency if requested as output.
if (cc->Outputs().HasTag("SALIENCY")) {
if (cc->Outputs().HasTag(kSaliencyTag)) {
with_saliency_ = true;
if (!options_.analysis_options().compute_motion_saliency()) {
LOG(WARNING) << "Enable saliency computation. Set "
@@ -353,11 +367,11 @@ absl::Status MotionAnalysisCalculator::Open(CalculatorContext* cc) {
cc->SetOffset(TimestampDiff(0));
}
if (cc->InputSidePackets().HasTag("DOWNSAMPLE")) {
if (cc->InputSidePackets().HasTag(kDownsampleTag)) {
options_.mutable_analysis_options()
->mutable_flow_options()
->set_downsample_factor(
cc->InputSidePackets().Tag("DOWNSAMPLE").Get<float>());
cc->InputSidePackets().Tag(kDownsampleTag).Get<float>());
}
// If no video header is provided, just return and initialize on the first
@@ -369,30 +383,33 @@ absl::Status MotionAnalysisCalculator::Open(CalculatorContext* cc) {
////////////// EARLY RETURN; ONLY HEADER OUTPUT SHOULD GO HERE ///////////////
if (visualize_output_) {
cc->Outputs().Tag("VIZ").SetHeader(Adopt(new VideoHeader(*video_header)));
cc->Outputs().Tag(kVizTag).SetHeader(Adopt(new VideoHeader(*video_header)));
}
if (video_output_) {
cc->Outputs()
.Tag("VIDEO_OUT")
.Tag(kVideoOutTag)
.SetHeader(Adopt(new VideoHeader(*video_header)));
}
if (cc->Outputs().HasTag("DENSE_FG")) {
if (cc->Outputs().HasTag(kDenseFgTag)) {
std::unique_ptr<VideoHeader> foreground_header(
new VideoHeader(*video_header));
foreground_header->format = ImageFormat::GRAY8;
cc->Outputs().Tag("DENSE_FG").SetHeader(Adopt(foreground_header.release()));
}
if (cc->Outputs().HasTag("CAMERA")) {
cc->Outputs().Tag("CAMERA").SetHeader(
Adopt(new VideoHeader(*video_header)));
}
if (cc->Outputs().HasTag("SALIENCY")) {
cc->Outputs()
.Tag("SALIENCY")
.Tag(kDenseFgTag)
.SetHeader(Adopt(foreground_header.release()));
}
if (cc->Outputs().HasTag(kCameraTag)) {
cc->Outputs()
.Tag(kCameraTag)
.SetHeader(Adopt(new VideoHeader(*video_header)));
}
if (cc->Outputs().HasTag(kSaliencyTag)) {
cc->Outputs()
.Tag(kSaliencyTag)
.SetHeader(Adopt(new VideoHeader(*video_header)));
}
@@ -405,9 +422,9 @@ absl::Status MotionAnalysisCalculator::Process(CalculatorContext* cc) {
}
InputStream* video_stream =
video_input_ ? &(cc->Inputs().Tag("VIDEO")) : nullptr;
video_input_ ? &(cc->Inputs().Tag(kVideoTag)) : nullptr;
InputStream* selection_stream =
selection_input_ ? &(cc->Inputs().Tag("SELECTION")) : nullptr;
selection_input_ ? &(cc->Inputs().Tag(kSelectionTag)) : nullptr;
// Checked on Open.
CHECK(video_stream || selection_stream);
@@ -425,8 +442,9 @@ absl::Status MotionAnalysisCalculator::Process(CalculatorContext* cc) {
CameraMotion output_motion = meta_motions_.front();
meta_motions_.pop_front();
output_motion.set_timestamp_usec(timestamp.Value());
cc->Outputs().Tag("CAMERA").Add(new CameraMotion(output_motion),
timestamp);
cc->Outputs()
.Tag(kCameraTag)
.Add(new CameraMotion(output_motion), timestamp);
}
if (region_flow_feature_output_) {
@@ -435,8 +453,8 @@ absl::Status MotionAnalysisCalculator::Process(CalculatorContext* cc) {
meta_features_.pop_front();
output_features.set_timestamp_usec(timestamp.Value());
cc->Outputs().Tag("FLOW").Add(new RegionFlowFeatureList(output_features),
timestamp);
cc->Outputs().Tag(kFlowTag).Add(
new RegionFlowFeatureList(output_features), timestamp);
}
++frame_idx_;
@@ -478,16 +496,17 @@ absl::Status MotionAnalysisCalculator::Process(CalculatorContext* cc) {
MotionAnalysisCalculatorOptions::NO_ANALYSIS_USE_SELECTION) {
// Output concatenated results, nothing to compute here.
if (camera_motion_output_) {
cc->Outputs().Tag("CAMERA").Add(
frame_selection_result->release_camera_motion(), timestamp);
cc->Outputs()
.Tag(kCameraTag)
.Add(frame_selection_result->release_camera_motion(), timestamp);
}
if (region_flow_feature_output_) {
cc->Outputs().Tag("FLOW").Add(frame_selection_result->release_features(),
timestamp);
cc->Outputs().Tag(kFlowTag).Add(
frame_selection_result->release_features(), timestamp);
}
if (video_output_) {
cc->Outputs().Tag("VIDEO_OUT").AddPacket(video_stream->Value());
cc->Outputs().Tag(kVideoOutTag).AddPacket(video_stream->Value());
}
return absl::OkStatus();
@@ -565,7 +584,7 @@ absl::Status MotionAnalysisCalculator::Process(CalculatorContext* cc) {
grayscale_mat.copyTo(image_frame_mat);
cc->Outputs()
.Tag("GRAY_VIDEO_OUT")
.Tag(kGrayVideoOutTag)
.Add(grayscale_image.release(), timestamp);
}
@@ -640,7 +659,7 @@ void MotionAnalysisCalculator::OutputMotionAnalyzedFrames(
*feature_list, *camera_motion,
with_saliency_ ? saliency[k].get() : nullptr, &visualization);
cc->Outputs().Tag("VIZ").Add(visualization_frame.release(), timestamp);
cc->Outputs().Tag(kVizTag).Add(visualization_frame.release(), timestamp);
}
// Output dense foreground mask.
@@ -650,26 +669,26 @@ void MotionAnalysisCalculator::OutputMotionAnalyzedFrames(
cv::Mat foreground = formats::MatView(foreground_frame.get());
motion_analysis_->ComputeDenseForeground(*feature_list, *camera_motion,
&foreground);
cc->Outputs().Tag("DENSE_FG").Add(foreground_frame.release(), timestamp);
cc->Outputs().Tag(kDenseFgTag).Add(foreground_frame.release(), timestamp);
}
// Output flow features if requested.
if (region_flow_feature_output_) {
cc->Outputs().Tag("FLOW").Add(feature_list.release(), timestamp);
cc->Outputs().Tag(kFlowTag).Add(feature_list.release(), timestamp);
}
// Output camera motion.
if (camera_motion_output_) {
cc->Outputs().Tag("CAMERA").Add(camera_motion.release(), timestamp);
cc->Outputs().Tag(kCameraTag).Add(camera_motion.release(), timestamp);
}
if (video_output_) {
cc->Outputs().Tag("VIDEO_OUT").AddPacket(packet_buffer_[k]);
cc->Outputs().Tag(kVideoOutTag).AddPacket(packet_buffer_[k]);
}
// Output saliency.
if (saliency_output_) {
cc->Outputs().Tag("SALIENCY").Add(saliency[k].release(), timestamp);
cc->Outputs().Tag(kSaliencyTag).Add(saliency[k].release(), timestamp);
}
}
mediapipe/calculators/video/opencv_video_decoder_calculator.cc
+20 -14
View File
@@ -27,6 +27,12 @@
namespace mediapipe {
namespace {
constexpr char kSavedAudioPathTag[] = "SAVED_AUDIO_PATH";
constexpr char kVideoPrestreamTag[] = "VIDEO_PRESTREAM";
constexpr char kVideoTag[] = "VIDEO";
constexpr char kInputFilePathTag[] = "INPUT_FILE_PATH";
// cv::VideoCapture set data type to unsigned char by default. Therefore, the
// image format is only related to the number of channles the cv::Mat has.
ImageFormat::Format GetImageFormat(int num_channels) {
@@ -87,20 +93,20 @@ ImageFormat::Format GetImageFormat(int num_channels) {
class OpenCvVideoDecoderCalculator : public CalculatorBase {
public:
static absl::Status GetContract(CalculatorContract* cc) {
cc->InputSidePackets().Tag("INPUT_FILE_PATH").Set<std::string>();
cc->Outputs().Tag("VIDEO").Set<ImageFrame>();
if (cc->Outputs().HasTag("VIDEO_PRESTREAM")) {
cc->Outputs().Tag("VIDEO_PRESTREAM").Set<VideoHeader>();
cc->InputSidePackets().Tag(kInputFilePathTag).Set<std::string>();
cc->Outputs().Tag(kVideoTag).Set<ImageFrame>();
if (cc->Outputs().HasTag(kVideoPrestreamTag)) {
cc->Outputs().Tag(kVideoPrestreamTag).Set<VideoHeader>();
}
if (cc->OutputSidePackets().HasTag("SAVED_AUDIO_PATH")) {
cc->OutputSidePackets().Tag("SAVED_AUDIO_PATH").Set<std::string>();
if (cc->OutputSidePackets().HasTag(kSavedAudioPathTag)) {
cc->OutputSidePackets().Tag(kSavedAudioPathTag).Set<std::string>();
}
return absl::OkStatus();
}
absl::Status Open(CalculatorContext* cc) override {
const std::string& input_file_path =
cc->InputSidePackets().Tag("INPUT_FILE_PATH").Get<std::string>();
cc->InputSidePackets().Tag(kInputFilePathTag).Get<std::string>();
cap_ = absl::make_unique<cv::VideoCapture>(input_file_path);
if (!cap_->isOpened()) {
return mediapipe::InvalidArgumentErrorBuilder(MEDIAPIPE_LOC)
@@ -140,16 +146,16 @@ class OpenCvVideoDecoderCalculator : public CalculatorBase {
header->frame_rate = fps;
header->duration = frame_count_ / fps;
if (cc->Outputs().HasTag("VIDEO_PRESTREAM")) {
if (cc->Outputs().HasTag(kVideoPrestreamTag)) {
cc->Outputs()
.Tag("VIDEO_PRESTREAM")
.Tag(kVideoPrestreamTag)
.Add(header.release(), Timestamp::PreStream());
cc->Outputs().Tag("VIDEO_PRESTREAM").Close();
cc->Outputs().Tag(kVideoPrestreamTag).Close();
}
// Rewind to the very first frame.
cap_->set(cv::CAP_PROP_POS_AVI_RATIO, 0);
if (cc->OutputSidePackets().HasTag("SAVED_AUDIO_PATH")) {
if (cc->OutputSidePackets().HasTag(kSavedAudioPathTag)) {
#ifdef HAVE_FFMPEG
std::string saved_audio_path = std::tmpnam(nullptr);
std::string ffmpeg_command =
@@ -159,14 +165,14 @@ class OpenCvVideoDecoderCalculator : public CalculatorBase {
int status_code = system(absl::StrCat("ls ", saved_audio_path).c_str());
if (status_code == 0) {
cc->OutputSidePackets()
.Tag("SAVED_AUDIO_PATH")
.Tag(kSavedAudioPathTag)
.Set(MakePacket<std::string>(saved_audio_path));
} else {
LOG(WARNING) << "FFmpeg can't extract audio from " << input_file_path
<< " by executing the following command: "
<< ffmpeg_command;
cc->OutputSidePackets()
.Tag("SAVED_AUDIO_PATH")
.Tag(kSavedAudioPathTag)
.Set(MakePacket<std::string>(std::string()));
}
#else
@@ -208,7 +214,7 @@ class OpenCvVideoDecoderCalculator : public CalculatorBase {
// If the timestamp of the current frame is not greater than the one of the
// previous frame, the new frame will be discarded.
if (prev_timestamp_ < timestamp) {
cc->Outputs().Tag("VIDEO").Add(image_frame.release(), timestamp);
cc->Outputs().Tag(kVideoTag).Add(image_frame.release(), timestamp);
prev_timestamp_ = timestamp;
decoded_frames_++;
}
mediapipe/calculators/video/opencv_video_decoder_calculator_test.cc
+22 -18
View File
@@ -29,6 +29,10 @@ namespace mediapipe {
namespace {
constexpr char kVideoTag[] = "VIDEO";
constexpr char kVideoPrestreamTag[] = "VIDEO_PRESTREAM";
constexpr char kInputFilePathTag[] = "INPUT_FILE_PATH";
TEST(OpenCvVideoDecoderCalculatorTest, TestMp4Avc720pVideo) {
CalculatorGraphConfig::Node node_config =
ParseTextProtoOrDie<CalculatorGraphConfig::Node>(R"pb(
@@ -37,19 +41,19 @@ TEST(OpenCvVideoDecoderCalculatorTest, TestMp4Avc720pVideo) {
output_stream: "VIDEO:video"
output_stream: "VIDEO_PRESTREAM:video_prestream")pb");
CalculatorRunner runner(node_config);
runner.MutableSidePackets()->Tag("INPUT_FILE_PATH") = MakePacket<std::string>(
runner.MutableSidePackets()->Tag(kInputFilePathTag) = MakePacket<std::string>(
file::JoinPath("./",
"/mediapipe/calculators/video/"
"testdata/format_MP4_AVC720P_AAC.video"));
MP_EXPECT_OK(runner.Run());
EXPECT_EQ(runner.Outputs().Tag("VIDEO_PRESTREAM").packets.size(), 1);
EXPECT_EQ(runner.Outputs().Tag(kVideoPrestreamTag).packets.size(), 1);
MP_EXPECT_OK(runner.Outputs()
.Tag("VIDEO_PRESTREAM")
.Tag(kVideoPrestreamTag)
.packets[0]
.ValidateAsType<VideoHeader>());
const mediapipe::VideoHeader& header =
runner.Outputs().Tag("VIDEO_PRESTREAM").packets[0].Get<VideoHeader>();
runner.Outputs().Tag(kVideoPrestreamTag).packets[0].Get<VideoHeader>();
EXPECT_EQ(ImageFormat::SRGB, header.format);
EXPECT_EQ(1280, header.width);
EXPECT_EQ(640, header.height);
@@ -58,10 +62,10 @@ TEST(OpenCvVideoDecoderCalculatorTest, TestMp4Avc720pVideo) {
// The number of the output packets should be 180.
// Some OpenCV version returns the first two frames with the same timestamp on
// macos and we might miss one frame here.
int num_of_packets = runner.Outputs().Tag("VIDEO").packets.size();
int num_of_packets = runner.Outputs().Tag(kVideoTag).packets.size();
EXPECT_GE(num_of_packets, 179);
for (int i = 0; i < num_of_packets; ++i) {
Packet image_frame_packet = runner.Outputs().Tag("VIDEO").packets[i];
Packet image_frame_packet = runner.Outputs().Tag(kVideoTag).packets[i];
cv::Mat output_mat =
formats::MatView(&(image_frame_packet.Get<ImageFrame>()));
EXPECT_EQ(1280, output_mat.size().width);
@@ -83,19 +87,19 @@ TEST(OpenCvVideoDecoderCalculatorTest, TestFlvH264Video) {
output_stream: "VIDEO:video"
output_stream: "VIDEO_PRESTREAM:video_prestream")pb");
CalculatorRunner runner(node_config);
runner.MutableSidePackets()->Tag("INPUT_FILE_PATH") = MakePacket<std::string>(
runner.MutableSidePackets()->Tag(kInputFilePathTag) = MakePacket<std::string>(
file::JoinPath("./",
"/mediapipe/calculators/video/"
"testdata/format_FLV_H264_AAC.video"));
MP_EXPECT_OK(runner.Run());
EXPECT_EQ(runner.Outputs().Tag("VIDEO_PRESTREAM").packets.size(), 1);
EXPECT_EQ(runner.Outputs().Tag(kVideoPrestreamTag).packets.size(), 1);
MP_EXPECT_OK(runner.Outputs()
.Tag("VIDEO_PRESTREAM")
.Tag(kVideoPrestreamTag)
.packets[0]
.ValidateAsType<VideoHeader>());
const mediapipe::VideoHeader& header =
runner.Outputs().Tag("VIDEO_PRESTREAM").packets[0].Get<VideoHeader>();
runner.Outputs().Tag(kVideoPrestreamTag).packets[0].Get<VideoHeader>();
EXPECT_EQ(ImageFormat::SRGB, header.format);
EXPECT_EQ(640, header.width);
EXPECT_EQ(320, header.height);
@@ -103,9 +107,9 @@ TEST(OpenCvVideoDecoderCalculatorTest, TestFlvH264Video) {
// can be either 30.30303f (with opencv2) or 30f (with opencv3 and opencv4).
// EXPECT_FLOAT_EQ(6.0f, header.duration);
// EXPECT_FLOAT_EQ(30.0f, header.frame_rate);
EXPECT_EQ(180, runner.Outputs().Tag("VIDEO").packets.size());
EXPECT_EQ(180, runner.Outputs().Tag(kVideoTag).packets.size());
for (int i = 0; i < 180; ++i) {
Packet image_frame_packet = runner.Outputs().Tag("VIDEO").packets[i];
Packet image_frame_packet = runner.Outputs().Tag(kVideoTag).packets[i];
cv::Mat output_mat =
formats::MatView(&(image_frame_packet.Get<ImageFrame>()));
EXPECT_EQ(640, output_mat.size().width);
@@ -127,19 +131,19 @@ TEST(OpenCvVideoDecoderCalculatorTest, TestMkvVp8Video) {
output_stream: "VIDEO:video"
output_stream: "VIDEO_PRESTREAM:video_prestream")pb");
CalculatorRunner runner(node_config);
runner.MutableSidePackets()->Tag("INPUT_FILE_PATH") = MakePacket<std::string>(
runner.MutableSidePackets()->Tag(kInputFilePathTag) = MakePacket<std::string>(
file::JoinPath("./",
"/mediapipe/calculators/video/"
"testdata/format_MKV_VP8_VORBIS.video"));
MP_EXPECT_OK(runner.Run());
EXPECT_EQ(runner.Outputs().Tag("VIDEO_PRESTREAM").packets.size(), 1);
EXPECT_EQ(runner.Outputs().Tag(kVideoPrestreamTag).packets.size(), 1);
MP_EXPECT_OK(runner.Outputs()
.Tag("VIDEO_PRESTREAM")
.Tag(kVideoPrestreamTag)
.packets[0]
.ValidateAsType<VideoHeader>());
const mediapipe::VideoHeader& header =
runner.Outputs().Tag("VIDEO_PRESTREAM").packets[0].Get<VideoHeader>();
runner.Outputs().Tag(kVideoPrestreamTag).packets[0].Get<VideoHeader>();
EXPECT_EQ(ImageFormat::SRGB, header.format);
EXPECT_EQ(640, header.width);
EXPECT_EQ(320, header.height);
@@ -148,10 +152,10 @@ TEST(OpenCvVideoDecoderCalculatorTest, TestMkvVp8Video) {
// The number of the output packets should be 180.
// Some OpenCV version returns the first two frames with the same timestamp on
// macos and we might miss one frame here.
int num_of_packets = runner.Outputs().Tag("VIDEO").packets.size();
int num_of_packets = runner.Outputs().Tag(kVideoTag).packets.size();
EXPECT_GE(num_of_packets, 179);
for (int i = 0; i < num_of_packets; ++i) {
Packet image_frame_packet = runner.Outputs().Tag("VIDEO").packets[i];
Packet image_frame_packet = runner.Outputs().Tag(kVideoTag).packets[i];
cv::Mat output_mat =
formats::MatView(&(image_frame_packet.Get<ImageFrame>()));
EXPECT_EQ(640, output_mat.size().width);
mediapipe/calculators/video/opencv_video_encoder_calculator.cc
+21 -16
View File
@@ -36,6 +36,11 @@
namespace mediapipe {
constexpr char kAudioFilePathTag[] = "AUDIO_FILE_PATH";
constexpr char kOutputFilePathTag[] = "OUTPUT_FILE_PATH";
constexpr char kVideoPrestreamTag[] = "VIDEO_PRESTREAM";
constexpr char kVideoTag[] = "VIDEO";
// Encodes the input video stream and produces a media file.
// The media file can be output to the output_file_path specified as a side
// packet. Currently, the calculator only supports one video stream (in
@@ -90,15 +95,15 @@ class OpenCvVideoEncoderCalculator : public CalculatorBase {
};
absl::Status OpenCvVideoEncoderCalculator::GetContract(CalculatorContract* cc) {
RET_CHECK(cc->Inputs().HasTag("VIDEO"));
cc->Inputs().Tag("VIDEO").Set<ImageFrame>();
if (cc->Inputs().HasTag("VIDEO_PRESTREAM")) {
cc->Inputs().Tag("VIDEO_PRESTREAM").Set<VideoHeader>();
RET_CHECK(cc->Inputs().HasTag(kVideoTag));
cc->Inputs().Tag(kVideoTag).Set<ImageFrame>();
if (cc->Inputs().HasTag(kVideoPrestreamTag)) {
cc->Inputs().Tag(kVideoPrestreamTag).Set<VideoHeader>();
}
RET_CHECK(cc->InputSidePackets().HasTag("OUTPUT_FILE_PATH"));
cc->InputSidePackets().Tag("OUTPUT_FILE_PATH").Set<std::string>();
if (cc->InputSidePackets().HasTag("AUDIO_FILE_PATH")) {
cc->InputSidePackets().Tag("AUDIO_FILE_PATH").Set<std::string>();
RET_CHECK(cc->InputSidePackets().HasTag(kOutputFilePathTag));
cc->InputSidePackets().Tag(kOutputFilePathTag).Set<std::string>();
if (cc->InputSidePackets().HasTag(kAudioFilePathTag)) {
cc->InputSidePackets().Tag(kAudioFilePathTag).Set<std::string>();
}
return absl::OkStatus();
}
@@ -116,7 +121,7 @@ absl::Status OpenCvVideoEncoderCalculator::Open(CalculatorContext* cc) {
<< "Video format must be specified in "
"OpenCvVideoEncoderCalculatorOptions";
output_file_path_ =
cc->InputSidePackets().Tag("OUTPUT_FILE_PATH").Get<std::string>();
cc->InputSidePackets().Tag(kOutputFilePathTag).Get<std::string>();
std::vector<std::string> splited_file_path =
absl::StrSplit(output_file_path_, '.');
RET_CHECK(splited_file_path.size() >= 2 &&
@@ -126,7 +131,7 @@ absl::Status OpenCvVideoEncoderCalculator::Open(CalculatorContext* cc) {
// If the video header will be available, the video metadata will be fetched
// from the video header directly. The calculator will receive the video
// header packet at timestamp prestream.
if (cc->Inputs().HasTag("VIDEO_PRESTREAM")) {
if (cc->Inputs().HasTag(kVideoPrestreamTag)) {
return absl::OkStatus();
}
return SetUpVideoWriter(options.fps(), options.width(), options.height());
@@ -135,13 +140,13 @@ absl::Status OpenCvVideoEncoderCalculator::Open(CalculatorContext* cc) {
absl::Status OpenCvVideoEncoderCalculator::Process(CalculatorContext* cc) {
if (cc->InputTimestamp() == Timestamp::PreStream()) {
const VideoHeader& video_header =
cc->Inputs().Tag("VIDEO_PRESTREAM").Get<VideoHeader>();
cc->Inputs().Tag(kVideoPrestreamTag).Get<VideoHeader>();
return SetUpVideoWriter(video_header.frame_rate, video_header.width,
video_header.height);
}
const ImageFrame& image_frame =
cc->Inputs().Tag("VIDEO").Value().Get<ImageFrame>();
cc->Inputs().Tag(kVideoTag).Value().Get<ImageFrame>();
ImageFormat::Format format = image_frame.Format();
cv::Mat frame;
if (format == ImageFormat::GRAY8) {
@@ -149,7 +154,7 @@ absl::Status OpenCvVideoEncoderCalculator::Process(CalculatorContext* cc) {
if (frame.empty()) {
return mediapipe::InvalidArgumentErrorBuilder(MEDIAPIPE_LOC)
<< "Receive empty frame at timestamp "
<< cc->Inputs().Tag("VIDEO").Value().Timestamp()
<< cc->Inputs().Tag(kVideoTag).Value().Timestamp()
<< " in OpenCvVideoEncoderCalculator::Process()";
}
} else {
@@ -157,7 +162,7 @@ absl::Status OpenCvVideoEncoderCalculator::Process(CalculatorContext* cc) {
if (tmp_frame.empty()) {
return mediapipe::InvalidArgumentErrorBuilder(MEDIAPIPE_LOC)
<< "Receive empty frame at timestamp "
<< cc->Inputs().Tag("VIDEO").Value().Timestamp()
<< cc->Inputs().Tag(kVideoTag).Value().Timestamp()
<< " in OpenCvVideoEncoderCalculator::Process()";
}
if (format == ImageFormat::SRGB) {
@@ -177,10 +182,10 @@ absl::Status OpenCvVideoEncoderCalculator::Close(CalculatorContext* cc) {
if (writer_ && writer_->isOpened()) {
writer_->release();
}
if (cc->InputSidePackets().HasTag("AUDIO_FILE_PATH")) {
if (cc->InputSidePackets().HasTag(kAudioFilePathTag)) {
#ifdef HAVE_FFMPEG
const std::string& audio_file_path =
cc->InputSidePackets().Tag("AUDIO_FILE_PATH").Get<std::string>();
cc->InputSidePackets().Tag(kAudioFilePathTag).Get<std::string>();
if (audio_file_path.empty()) {
LOG(WARNING) << "OpenCvVideoEncoderCalculator isn't able to attach the "
"audio tracks to the generated video because the audio "
mediapipe/calculators/video/tvl1_optical_flow_calculator.cc
+19 -14
View File
@@ -23,6 +23,11 @@
namespace mediapipe {
namespace {
constexpr char kBackwardFlowTag[] = "BACKWARD_FLOW";
constexpr char kForwardFlowTag[] = "FORWARD_FLOW";
constexpr char kSecondFrameTag[] = "SECOND_FRAME";
constexpr char kFirstFrameTag[] = "FIRST_FRAME";
// Checks that img1 and img2 have the same dimensions.
bool ImageSizesMatch(const ImageFrame& img1, const ImageFrame& img2) {
return (img1.Width() == img2.Width()) && (img1.Height() == img2.Height());
@@ -94,19 +99,19 @@ class Tvl1OpticalFlowCalculator : public CalculatorBase {
};
absl::Status Tvl1OpticalFlowCalculator::GetContract(CalculatorContract* cc) {
if (!cc->Inputs().HasTag("FIRST_FRAME") ||
!cc->Inputs().HasTag("SECOND_FRAME")) {
if (!cc->Inputs().HasTag(kFirstFrameTag) ||
!cc->Inputs().HasTag(kSecondFrameTag)) {
return absl::InvalidArgumentError(
"Missing required input streams. Both FIRST_FRAME and SECOND_FRAME "
"must be specified.");
}
cc->Inputs().Tag("FIRST_FRAME").Set<ImageFrame>();
cc->Inputs().Tag("SECOND_FRAME").Set<ImageFrame>();
if (cc->Outputs().HasTag("FORWARD_FLOW")) {
cc->Outputs().Tag("FORWARD_FLOW").Set<OpticalFlowField>();
cc->Inputs().Tag(kFirstFrameTag).Set<ImageFrame>();
cc->Inputs().Tag(kSecondFrameTag).Set<ImageFrame>();
if (cc->Outputs().HasTag(kForwardFlowTag)) {
cc->Outputs().Tag(kForwardFlowTag).Set<OpticalFlowField>();
}
if (cc->Outputs().HasTag("BACKWARD_FLOW")) {
cc->Outputs().Tag("BACKWARD_FLOW").Set<OpticalFlowField>();
if (cc->Outputs().HasTag(kBackwardFlowTag)) {
cc->Outputs().Tag(kBackwardFlowTag).Set<OpticalFlowField>();
}
return absl::OkStatus();
}
@@ -116,10 +121,10 @@ absl::Status Tvl1OpticalFlowCalculator::Open(CalculatorContext* cc) {
absl::MutexLock lock(&mutex_);
tvl1_computers_.emplace_back(cv::createOptFlow_DualTVL1());
}
if (cc->Outputs().HasTag("FORWARD_FLOW")) {
if (cc->Outputs().HasTag(kForwardFlowTag)) {
forward_requested_ = true;
}
if (cc->Outputs().HasTag("BACKWARD_FLOW")) {
if (cc->Outputs().HasTag(kBackwardFlowTag)) {
backward_requested_ = true;
}
@@ -128,15 +133,15 @@ absl::Status Tvl1OpticalFlowCalculator::Open(CalculatorContext* cc) {
absl::Status Tvl1OpticalFlowCalculator::Process(CalculatorContext* cc) {
const ImageFrame& first_frame =
cc->Inputs().Tag("FIRST_FRAME").Value().Get<ImageFrame>();
cc->Inputs().Tag(kFirstFrameTag).Value().Get<ImageFrame>();
const ImageFrame& second_frame =
cc->Inputs().Tag("SECOND_FRAME").Value().Get<ImageFrame>();
cc->Inputs().Tag(kSecondFrameTag).Value().Get<ImageFrame>();
if (forward_requested_) {
auto forward_optical_flow_field = absl::make_unique<OpticalFlowField>();
MP_RETURN_IF_ERROR(CalculateOpticalFlow(first_frame, second_frame,
forward_optical_flow_field.get()));
cc->Outputs()
.Tag("FORWARD_FLOW")
.Tag(kForwardFlowTag)
.Add(forward_optical_flow_field.release(), cc->InputTimestamp());
}
if (backward_requested_) {
@@ -144,7 +149,7 @@ absl::Status Tvl1OpticalFlowCalculator::Process(CalculatorContext* cc) {
MP_RETURN_IF_ERROR(CalculateOpticalFlow(second_frame, first_frame,
backward_optical_flow_field.get()));
cc->Outputs()
.Tag("BACKWARD_FLOW")
.Tag(kBackwardFlowTag)
.Add(backward_optical_flow_field.release(), cc->InputTimestamp());
}
return absl::OkStatus();
mediapipe/calculators/video/video_pre_stream_calculator.cc
+13 -10
View File
@@ -19,6 +19,9 @@
namespace mediapipe {
constexpr char kVideoPrestreamTag[] = "VIDEO_PRESTREAM";
constexpr char kFrameTag[] = "FRAME";
// Sets up VideoHeader based on the 1st ImageFrame and emits it with timestamp
// PreStream. Note that this calculator only fills in format, width, and height,
// i.e. frame_rate and duration will not be filled, unless:
@@ -64,8 +67,8 @@ absl::Status VideoPreStreamCalculator::GetContract(CalculatorContract* cc) {
if (!cc->Inputs().UsesTags()) {
cc->Inputs().Index(0).Set<ImageFrame>();
} else {
cc->Inputs().Tag("FRAME").Set<ImageFrame>();
cc->Inputs().Tag("VIDEO_PRESTREAM").Set<VideoHeader>();
cc->Inputs().Tag(kFrameTag).Set<ImageFrame>();
cc->Inputs().Tag(kVideoPrestreamTag).Set<VideoHeader>();
}
cc->Outputs().Index(0).Set<VideoHeader>();
return absl::OkStatus();
@@ -73,8 +76,8 @@ absl::Status VideoPreStreamCalculator::GetContract(CalculatorContract* cc) {
absl::Status VideoPreStreamCalculator::Open(CalculatorContext* cc) {
frame_rate_in_prestream_ = cc->Inputs().UsesTags() &&
cc->Inputs().HasTag("FRAME") &&
cc->Inputs().HasTag("VIDEO_PRESTREAM");
cc->Inputs().HasTag(kFrameTag) &&
cc->Inputs().HasTag(kVideoPrestreamTag);
header_ = absl::make_unique<VideoHeader>();
return absl::OkStatus();
}
@@ -82,15 +85,15 @@ absl::Status VideoPreStreamCalculator::ProcessWithFrameRateInPreStream(
CalculatorContext* cc) {
cc->GetCounter("ProcessWithFrameRateInPreStream")->Increment();
if (cc->InputTimestamp() == Timestamp::PreStream()) {
RET_CHECK(cc->Inputs().Tag("FRAME").IsEmpty());
RET_CHECK(!cc->Inputs().Tag("VIDEO_PRESTREAM").IsEmpty());
*header_ = cc->Inputs().Tag("VIDEO_PRESTREAM").Get<VideoHeader>();
RET_CHECK(cc->Inputs().Tag(kFrameTag).IsEmpty());
RET_CHECK(!cc->Inputs().Tag(kVideoPrestreamTag).IsEmpty());
*header_ = cc->Inputs().Tag(kVideoPrestreamTag).Get<VideoHeader>();
RET_CHECK_NE(header_->frame_rate, 0.0) << "frame rate should be non-zero";
} else {
RET_CHECK(cc->Inputs().Tag("VIDEO_PRESTREAM").IsEmpty())
RET_CHECK(cc->Inputs().Tag(kVideoPrestreamTag).IsEmpty())
<< "Packet on VIDEO_PRESTREAM must come in at Timestamp::PreStream().";
RET_CHECK(!cc->Inputs().Tag("FRAME").IsEmpty());
const auto& frame = cc->Inputs().Tag("FRAME").Get<ImageFrame>();
RET_CHECK(!cc->Inputs().Tag(kFrameTag).IsEmpty());
const auto& frame = cc->Inputs().Tag(kFrameTag).Get<ImageFrame>();
header_->format = frame.Format();
header_->width = frame.Width();
header_->height = frame.Height();