// Copyright 2020 Google LLC
//
// 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.
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/port/ret_check.h"
#include "mediapipe/framework/port/status.h"
#include "mediapipe/graphs/instant_motion_tracking/calculators/transformations.h"
#include "mediapipe/util/tracking/box_tracker.pb.h"
namespace mediapipe {
constexpr char kSentinelTag[] = "SENTINEL";
constexpr char kAnchorsTag[] = "ANCHORS";
constexpr char kBoxesInputTag[] = "BOXES";
constexpr char kBoxesOutputTag[] = "START_POS";
constexpr char kCancelTag[] = "CANCEL_ID";
// TODO: Find optimal Height/Width (0.1-0.3)
constexpr float kBoxEdgeSize =
0.2f; // Used to establish tracking box dimensions
constexpr float kUsToMs =
1000.0f; // Used to convert from microseconds to millis
// This calculator manages the regions being tracked for each individual sticker
// and adjusts the regions being tracked if a change is detected in a sticker's
// initial anchor placement. Regions being tracked that have no associated
// sticker will be automatically removed upon the next iteration of the graph to
// optimize performance and remove all sticker artifacts
//
// Input:
// SENTINEL - ID of sticker which has an anchor that must be reset (-1 when no
// anchor must be reset) [REQUIRED]
// ANCHORS - Initial anchor data (tracks changes and where to re/position)
// [REQUIRED] BOXES - Used in cycle, boxes being tracked meant to update
// positions [OPTIONAL
// - provided by subgraph]
// Output:
// START_POS - Positions of boxes being tracked (can be overwritten with ID)
// [REQUIRED] CANCEL_ID - Single integer ID of tracking box to remove from
// tracker subgraph [OPTIONAL] ANCHORS - Updated set of anchors with tracked
// and normalized X,Y,Z [REQUIRED]
//
// Example config:
// node {
// calculator: "TrackedAnchorManagerCalculator"
// input_stream: "SENTINEL:sticker_sentinel"
// input_stream: "ANCHORS:initial_anchor_data"
// input_stream: "BOXES:boxes"
// input_stream_info: {
// tag_index: 'BOXES'
// back_edge: true
// }
// output_stream: "START_POS:start_pos"
// output_stream: "CANCEL_ID:cancel_object_id"
// output_stream: "ANCHORS:tracked_scaled_anchor_data"
// }
class TrackedAnchorManagerCalculator : public CalculatorBase {
private:
// Previous graph iteration anchor data
std::vector<Anchor> previous_anchor_data_;
public:
static absl::Status GetContract(CalculatorContract* cc) {
RET_CHECK(cc->Inputs().HasTag(kAnchorsTag) &&
cc->Inputs().HasTag(kSentinelTag));
RET_CHECK(cc->Outputs().HasTag(kAnchorsTag) &&
cc->Outputs().HasTag(kBoxesOutputTag));
cc->Inputs().Tag(kAnchorsTag).Set<std::vector<Anchor>>();
cc->Inputs().Tag(kSentinelTag).Set<int>();
if (cc->Inputs().HasTag(kBoxesInputTag)) {
cc->Inputs().Tag(kBoxesInputTag).Set<mediapipe::TimedBoxProtoList>();
}
cc->Outputs().Tag(kAnchorsTag).Set<std::vector<Anchor>>();
cc->Outputs().Tag(kBoxesOutputTag).Set<mediapipe::TimedBoxProtoList>();
if (cc->Outputs().HasTag(kCancelTag)) {
cc->Outputs().Tag(kCancelTag).Set<int>();
}
return absl::OkStatus();
}
absl::Status Open(CalculatorContext* cc) override { return absl::OkStatus(); }
absl::Status Process(CalculatorContext* cc) override;
};
REGISTER_CALCULATOR(TrackedAnchorManagerCalculator);
absl::Status TrackedAnchorManagerCalculator::Process(CalculatorContext* cc) {
mediapipe::Timestamp timestamp = cc->InputTimestamp();
const int sticker_sentinel = cc->Inputs().Tag(kSentinelTag).Get<int>();
std::vector<Anchor> current_anchor_data =
cc->Inputs().Tag(kAnchorsTag).Get<std::vector<Anchor>>();
auto pos_boxes = absl::make_unique<mediapipe::TimedBoxProtoList>();
std::vector<Anchor> tracked_scaled_anchor_data;
// Delete any boxes being tracked without an associated anchor
for (const mediapipe::TimedBoxProto& box :
cc->Inputs()
.Tag(kBoxesInputTag)
.Get<mediapipe::TimedBoxProtoList>()
.box()) {
bool anchor_exists = false;
for (Anchor anchor : current_anchor_data) {
if (box.id() == anchor.sticker_id) {
anchor_exists = true;
break;
}
}
if (!anchor_exists) {
cc->Outputs()
.Tag(kCancelTag)
.AddPacket(MakePacket<int>(box.id()).At(timestamp++));
}
}
// Perform tracking or updating for each anchor position
for (const Anchor& anchor : current_anchor_data) {
Anchor output_anchor = anchor;
// Check if anchor position is being reset by user in this graph iteration
if (sticker_sentinel == anchor.sticker_id) {
// Delete associated tracking box
// TODO: BoxTrackingSubgraph should accept vector to avoid breaking
// timestamp rules
cc->Outputs()
.Tag(kCancelTag)
.AddPacket(MakePacket<int>(anchor.sticker_id).At(timestamp++));
// Add a tracking box
mediapipe::TimedBoxProto* box = pos_boxes->add_box();
box->set_left(anchor.x - kBoxEdgeSize * 0.5f);
box->set_right(anchor.x + kBoxEdgeSize * 0.5f);
box->set_top(anchor.y - kBoxEdgeSize * 0.5f);
box->set_bottom(anchor.y + kBoxEdgeSize * 0.5f);
box->set_id(anchor.sticker_id);
box->set_time_msec((timestamp++).Microseconds() / kUsToMs);
// Default value for normalized z (scale factor)
output_anchor.z = 1.0f;
} else {
// Anchor position was not reset by user
// Attempt to update anchor position from tracking subgraph
// (TimedBoxProto)
bool updated_from_tracker = false;
const mediapipe::TimedBoxProtoList box_list =
cc->Inputs().Tag(kBoxesInputTag).Get<mediapipe::TimedBoxProtoList>();
for (const auto& box : box_list.box()) {
if (box.id() == anchor.sticker_id) {
// Get center x normalized coordinate [0.0-1.0]
output_anchor.x = (box.left() + box.right()) * 0.5f;
// Get center y normalized coordinate [0.0-1.0]
output_anchor.y = (box.top() + box.bottom()) * 0.5f;
// Get center z coordinate [z starts at normalized 1.0 and scales
// inversely with box-width]
// TODO: Look into issues with uniform scaling on x-axis and y-axis
output_anchor.z = kBoxEdgeSize / (box.right() - box.left());
updated_from_tracker = true;
break;
}
}
// If anchor position was not updated from tracker, create new tracking
// box at last recorded anchor coordinates. This will allow all current
// stickers to be tracked at approximately last location even if
// re-acquisitioning in the BoxTrackingSubgraph encounters errors
if (!updated_from_tracker) {
for (const Anchor& prev_anchor : previous_anchor_data_) {
if (anchor.sticker_id == prev_anchor.sticker_id) {
mediapipe::TimedBoxProto* box = pos_boxes->add_box();
box->set_left(prev_anchor.x - kBoxEdgeSize * 0.5f);
box->set_right(prev_anchor.x + kBoxEdgeSize * 0.5f);
box->set_top(prev_anchor.y - kBoxEdgeSize * 0.5f);
box->set_bottom(prev_anchor.y + kBoxEdgeSize * 0.5f);
box->set_id(prev_anchor.sticker_id);
box->set_time_msec(cc->InputTimestamp().Microseconds() / kUsToMs);
output_anchor = prev_anchor;
// Default value for normalized z (scale factor)
output_anchor.z = 1.0f;
break;
}
}
}
}
tracked_scaled_anchor_data.emplace_back(output_anchor);
}
// Set anchor data for next iteration
previous_anchor_data_ = tracked_scaled_anchor_data;
cc->Outputs()
.Tag(kAnchorsTag)
.AddPacket(MakePacket<std::vector<Anchor>>(tracked_scaled_anchor_data)
.At(cc->InputTimestamp()));
cc->Outputs()
.Tag(kBoxesOutputTag)
.Add(pos_boxes.release(), cc->InputTimestamp());
return absl::OkStatus();
}
} // namespace mediapipe
Codebase Browser
chromium