// Copyright 2019 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.
// MediaPipe calculator to take a flow field as input, and outputs a normalized
// RGB image where the B channel is forced to zero.
// TODO: Add video stream header for visualization
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include "absl/strings/str_format.h"
#include "absl/strings/string_view.h"
#include "mediapipe/calculators/video/flow_to_image_calculator.pb.h"
#include "mediapipe/calculators/video/tool/flow_quantizer_model.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/formats/image_format.pb.h"
#include "mediapipe/framework/formats/image_frame.h"
#include "mediapipe/framework/formats/image_frame_opencv.h"
#include "mediapipe/framework/formats/motion/optical_flow_field.h"
#include "mediapipe/framework/port/opencv_core_inc.h"
#include "mediapipe/framework/port/parse_text_proto.h"
namespace mediapipe {
// Reads optical flow fields defined in
// mediapipe/framework/formats/motion/optical_flow_field.h,
// returns a VideoFrame with 2 channels (v_x and v_y), each channel is quantized
// to 0-255.
//
// Example config:
// node {
// calculator: "FlowToImageCalculator"
// input_stream: "flow_fields"
// output_stream: "frames"
// options: {
// [type.googleapis.com/mediapipe.FlowToImageCalculatorOptions]:{
// min_value: -40.0
// max_value: 40.0
// }
// }
// }
class FlowToImageCalculator : public CalculatorBase {
public:
FlowToImageCalculator() {}
~FlowToImageCalculator() override {}
static absl::Status GetContract(CalculatorContract* cc);
absl::Status Open(CalculatorContext* cc) override;
absl::Status Process(CalculatorContext* cc) override;
private:
FlowQuantizerModel model_;
};
absl::Status FlowToImageCalculator::GetContract(CalculatorContract* cc) {
cc->Inputs().Index(0).Set<OpticalFlowField>();
cc->Outputs().Index(0).Set<ImageFrame>();
// Model sanity check
const auto& options = cc->Options<FlowToImageCalculatorOptions>();
if (options.min_value() >= options.max_value()) {
return absl::InvalidArgumentError("Invalid quantizer model.");
}
return absl::OkStatus();
}
absl::Status FlowToImageCalculator::Open(CalculatorContext* cc) {
const auto& options = cc->Options<FlowToImageCalculatorOptions>();
// Fill the the model_data, ideally we want to train the model, but we omit
// the step for now, and takes the (min, max) range from protobuf.
const QuantizerModelData& model_data =
ParseTextProtoOrDie<QuantizerModelData>(
absl::StrFormat("min_value:%f min_value:%f max_value:%f max_value:%f",
options.min_value(), options.min_value(),
options.max_value(), options.max_value()));
model_.LoadFromProto(model_data);
return absl::OkStatus();
}
absl::Status FlowToImageCalculator::Process(CalculatorContext* cc) {
const auto& input = cc->Inputs().Index(0).Get<OpticalFlowField>();
// Input flow is 2-channel with x-dim flow and y-dim flow.
// Convert it to a ImageFrame in SRGB space, the 3rd channel is not used (0).
const cv::Mat_<cv::Point2f>& flow = input.flow_data();
std::unique_ptr<ImageFrame> output(
new ImageFrame(ImageFormat::SRGB, input.width(), input.height()));
cv::Mat image = ::mediapipe::formats::MatView(output.get());
for (int j = 0; j != input.height(); ++j) {
for (int i = 0; i != input.width(); ++i) {
image.at<cv::Vec3b>(j, i) =
cv::Vec3b(model_.Apply(flow.at<cv::Point2f>(j, i).x, 0),
model_.Apply(flow.at<cv::Point2f>(j, i).y, 1), 0);
}
}
cc->Outputs().Index(0).Add(output.release(), cc->InputTimestamp());
return absl::OkStatus();
}
REGISTER_CALCULATOR(FlowToImageCalculator);
} // namespace mediapipe