Project import generated by Copybara.
PiperOrigin-RevId: 253489161
This commit is contained in:
MediaPipe Team
@@ -0,0 +1,157 @@
|
||||
# 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.
|
||||
#
|
||||
|
||||
licenses(["notice"]) # Apache 2.0
|
||||
|
||||
package(default_visibility = ["//visibility:private"])
|
||||
|
||||
load("//mediapipe/framework/port:build_config.bzl", "mediapipe_cc_proto_library")
|
||||
|
||||
proto_library(
|
||||
name = "flow_to_image_calculator_proto",
|
||||
srcs = ["flow_to_image_calculator.proto"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = ["//mediapipe/framework:calculator_proto"],
|
||||
)
|
||||
|
||||
proto_library(
|
||||
name = "opencv_video_encoder_calculator_proto",
|
||||
srcs = ["opencv_video_encoder_calculator.proto"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = ["//mediapipe/framework:calculator_proto"],
|
||||
)
|
||||
|
||||
mediapipe_cc_proto_library(
|
||||
name = "flow_to_image_calculator_cc_proto",
|
||||
srcs = ["flow_to_image_calculator.proto"],
|
||||
cc_deps = ["//mediapipe/framework:calculator_cc_proto"],
|
||||
visibility = ["//mediapipe:__subpackages__"],
|
||||
deps = [":flow_to_image_calculator_proto"],
|
||||
)
|
||||
|
||||
mediapipe_cc_proto_library(
|
||||
name = "opencv_video_encoder_calculator_cc_proto",
|
||||
srcs = ["opencv_video_encoder_calculator.proto"],
|
||||
cc_deps = ["//mediapipe/framework:calculator_cc_proto"],
|
||||
visibility = ["//mediapipe:__subpackages__"],
|
||||
deps = [":opencv_video_encoder_calculator_proto"],
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "flow_to_image_calculator",
|
||||
srcs = ["flow_to_image_calculator.cc"],
|
||||
visibility = [
|
||||
"//visibility:public",
|
||||
],
|
||||
deps = [
|
||||
"//mediapipe/calculators/video:flow_to_image_calculator_cc_proto",
|
||||
"//mediapipe/calculators/video/tool:flow_quantizer_model",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/formats:image_format_cc_proto",
|
||||
"//mediapipe/framework/formats:image_frame",
|
||||
"//mediapipe/framework/formats:image_frame_opencv",
|
||||
"//mediapipe/framework/formats/motion:optical_flow_field",
|
||||
"//mediapipe/framework/port:opencv_core",
|
||||
"//mediapipe/framework/port:parse_text_proto",
|
||||
"@com_google_absl//absl/strings",
|
||||
"@com_google_absl//absl/strings:str_format",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "opencv_video_decoder_calculator",
|
||||
srcs = ["opencv_video_decoder_calculator.cc"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/formats:image_format_cc_proto",
|
||||
"//mediapipe/framework/formats:image_frame",
|
||||
"//mediapipe/framework/formats:image_frame_opencv",
|
||||
"//mediapipe/framework/formats:video_stream_header",
|
||||
"//mediapipe/framework/port:opencv_imgproc",
|
||||
"//mediapipe/framework/port:opencv_video",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
"//mediapipe/framework/port:status",
|
||||
"//mediapipe/framework/tool:status_util",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "opencv_video_encoder_calculator",
|
||||
srcs = ["opencv_video_encoder_calculator.cc"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
":opencv_video_encoder_calculator_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/formats:image_frame",
|
||||
"//mediapipe/framework/formats:image_frame_opencv",
|
||||
"//mediapipe/framework/formats:video_stream_header",
|
||||
"//mediapipe/framework/port:file_helpers",
|
||||
"//mediapipe/framework/port:opencv_highgui",
|
||||
"//mediapipe/framework/port:opencv_imgproc",
|
||||
"//mediapipe/framework/port:opencv_video",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
"//mediapipe/framework/port:source_location",
|
||||
"//mediapipe/framework/port:status",
|
||||
"//mediapipe/framework/tool:status_util",
|
||||
"@com_google_absl//absl/strings",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_test(
|
||||
name = "opencv_video_decoder_calculator_test",
|
||||
srcs = ["opencv_video_decoder_calculator_test.cc"],
|
||||
data = ["//mediapipe/calculators/video/testdata:test_videos"],
|
||||
deps = [
|
||||
":opencv_video_decoder_calculator",
|
||||
"//mediapipe/framework:calculator_runner",
|
||||
"//mediapipe/framework/deps:file_path",
|
||||
"//mediapipe/framework/formats:image_frame",
|
||||
"//mediapipe/framework/formats:image_frame_opencv",
|
||||
"//mediapipe/framework/formats:video_stream_header",
|
||||
"//mediapipe/framework/port:gtest_main",
|
||||
"//mediapipe/framework/port:integral_types",
|
||||
"//mediapipe/framework/port:logging",
|
||||
"//mediapipe/framework/port:opencv_core",
|
||||
"//mediapipe/framework/port:parse_text_proto",
|
||||
],
|
||||
)
|
||||
|
||||
cc_test(
|
||||
name = "opencv_video_encoder_calculator_test",
|
||||
srcs = ["opencv_video_encoder_calculator_test.cc"],
|
||||
data = ["//mediapipe/calculators/video/testdata:test_videos"],
|
||||
deps = [
|
||||
":opencv_video_decoder_calculator",
|
||||
":opencv_video_encoder_calculator",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:packet",
|
||||
"//mediapipe/framework/deps:file_path",
|
||||
"//mediapipe/framework/formats:deleting_file",
|
||||
"//mediapipe/framework/formats:image_frame",
|
||||
"//mediapipe/framework/formats:image_frame_opencv",
|
||||
"//mediapipe/framework/formats:video_stream_header",
|
||||
"//mediapipe/framework/port:gtest_main",
|
||||
"//mediapipe/framework/port:integral_types",
|
||||
"//mediapipe/framework/port:logging",
|
||||
"//mediapipe/framework/port:opencv_highgui",
|
||||
"//mediapipe/framework/port:opencv_imgproc",
|
||||
"//mediapipe/framework/port:opencv_video",
|
||||
"//mediapipe/framework/port:parse_text_proto",
|
||||
],
|
||||
)
|
||||
@@ -0,0 +1,114 @@
|
||||
// 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 ::mediapipe::Status GetContract(CalculatorContract* cc);
|
||||
::mediapipe::Status Open(CalculatorContext* cc) override;
|
||||
::mediapipe::Status Process(CalculatorContext* cc) override;
|
||||
|
||||
private:
|
||||
FlowQuantizerModel model_;
|
||||
};
|
||||
|
||||
::mediapipe::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 ::mediapipe::InvalidArgumentError("Invalid quantizer model.");
|
||||
}
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::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 ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::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 ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
REGISTER_CALCULATOR(FlowToImageCalculator);
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,29 @@
|
||||
// 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.
|
||||
|
||||
syntax = "proto2";
|
||||
|
||||
package mediapipe;
|
||||
|
||||
import "mediapipe/framework/calculator.proto";
|
||||
|
||||
// Specifies the maximum and minimum value to truncate when normalize optical
|
||||
// flow fields.
|
||||
message FlowToImageCalculatorOptions {
|
||||
extend CalculatorOptions {
|
||||
optional FlowToImageCalculatorOptions ext = 69508592;
|
||||
}
|
||||
optional float min_value = 1 [default = -40.0];
|
||||
optional float max_value = 2 [default = 40.0];
|
||||
}
|
||||
@@ -0,0 +1,184 @@
|
||||
// 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.
|
||||
|
||||
#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/video_stream_header.h"
|
||||
#include "mediapipe/framework/port/opencv_imgproc_inc.h"
|
||||
#include "mediapipe/framework/port/opencv_video_inc.h"
|
||||
#include "mediapipe/framework/port/status.h"
|
||||
#include "mediapipe/framework/tool/status_util.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
namespace {
|
||||
// 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) {
|
||||
ImageFormat::Format format;
|
||||
switch (num_channels) {
|
||||
case 1:
|
||||
format = ImageFormat::GRAY8;
|
||||
break;
|
||||
case 3:
|
||||
format = ImageFormat::SRGB;
|
||||
break;
|
||||
case 4:
|
||||
format = ImageFormat::SRGBA;
|
||||
break;
|
||||
default:
|
||||
format = ImageFormat::UNKNOWN;
|
||||
break;
|
||||
}
|
||||
return format;
|
||||
}
|
||||
} // namespace
|
||||
|
||||
// This Calculator takes no input streams and produces video packets.
|
||||
// All streams and input side packets are specified using tags and all of them
|
||||
// are optional.
|
||||
//
|
||||
// Output Streams:
|
||||
// VIDEO: Output video frames (ImageFrame).
|
||||
// VIDEO_PRESTREAM:
|
||||
// Optional video header information output at
|
||||
// Timestamp::PreStream() for the corresponding stream.
|
||||
// Input Side Packets:
|
||||
// INPUT_FILE_PATH: The input file path.
|
||||
//
|
||||
// Example config:
|
||||
// node {
|
||||
// calculator: "OpenCvVideoDecoderCalculator"
|
||||
// input_side_packet: "INPUT_FILE_PATH:input_file_path"
|
||||
// output_stream: "VIDEO:video_frames"
|
||||
// output_stream: "VIDEO_PRESTREAM:video_header"
|
||||
// }
|
||||
class OpenCvVideoDecoderCalculator : public CalculatorBase {
|
||||
public:
|
||||
static ::mediapipe::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>();
|
||||
}
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Open(CalculatorContext* cc) override {
|
||||
const std::string& input_file_path =
|
||||
cc->InputSidePackets().Tag("INPUT_FILE_PATH").Get<std::string>();
|
||||
cap_ = absl::make_unique<cv::VideoCapture>(input_file_path);
|
||||
if (!cap_->isOpened()) {
|
||||
return ::mediapipe::InvalidArgumentErrorBuilder(MEDIAPIPE_LOC)
|
||||
<< "Fail to open video file at " << input_file_path;
|
||||
}
|
||||
width_ = static_cast<int>(cap_->get(cv::CAP_PROP_FRAME_WIDTH));
|
||||
height_ = static_cast<int>(cap_->get(cv::CAP_PROP_FRAME_HEIGHT));
|
||||
double fps = static_cast<double>(cap_->get(cv::CAP_PROP_FPS));
|
||||
frame_count_ = static_cast<int>(cap_->get(cv::CAP_PROP_FRAME_COUNT));
|
||||
// Unfortunately, cap_->get(cv::CAP_PROP_FORMAT) always returns CV_8UC1
|
||||
// back. To get correct image format, we read the first frame from the video
|
||||
// and get the number of channels.
|
||||
cv::Mat frame;
|
||||
cap_->read(frame);
|
||||
if (frame.empty()) {
|
||||
return ::mediapipe::InvalidArgumentErrorBuilder(MEDIAPIPE_LOC)
|
||||
<< "Fail to read any frames from the video file at "
|
||||
<< input_file_path;
|
||||
}
|
||||
format_ = GetImageFormat(frame.channels());
|
||||
if (format_ == ImageFormat::UNKNOWN) {
|
||||
return ::mediapipe::InvalidArgumentErrorBuilder(MEDIAPIPE_LOC)
|
||||
<< "Unsupported video format of the video file at "
|
||||
<< input_file_path;
|
||||
}
|
||||
|
||||
if (fps <= 0 || frame_count_ <= 0 || width_ <= 0 || height_ <= 0) {
|
||||
return ::mediapipe::InvalidArgumentErrorBuilder(MEDIAPIPE_LOC)
|
||||
<< "Fail to make video header due to the incorrect metadata from "
|
||||
"the video file at "
|
||||
<< input_file_path;
|
||||
}
|
||||
auto header = absl::make_unique<VideoHeader>();
|
||||
header->format = format_;
|
||||
header->width = width_;
|
||||
header->height = height_;
|
||||
header->frame_rate = fps;
|
||||
header->duration = frame_count_ / fps;
|
||||
|
||||
if (cc->Outputs().HasTag("VIDEO_PRESTREAM")) {
|
||||
cc->Outputs()
|
||||
.Tag("VIDEO_PRESTREAM")
|
||||
.Add(header.release(), Timestamp::PreStream());
|
||||
}
|
||||
// Rewind to the very first frame.
|
||||
cap_->set(cv::CAP_PROP_POS_AVI_RATIO, 0);
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Process(CalculatorContext* cc) override {
|
||||
auto image_frame = absl::make_unique<ImageFrame>(format_, width_, height_,
|
||||
/*alignment_boundary=*/1);
|
||||
// Use microsecond as the unit of time.
|
||||
Timestamp timestamp(cap_->get(cv::CAP_PROP_POS_MSEC) * 1000);
|
||||
if (format_ == ImageFormat::GRAY8) {
|
||||
cv::Mat frame = formats::MatView(image_frame.get());
|
||||
cap_->read(frame);
|
||||
if (frame.empty()) {
|
||||
return tool::StatusStop();
|
||||
}
|
||||
} else {
|
||||
cv::Mat tmp_frame;
|
||||
cap_->read(tmp_frame);
|
||||
if (tmp_frame.empty()) {
|
||||
return tool::StatusStop();
|
||||
}
|
||||
if (format_ == ImageFormat::SRGB) {
|
||||
cv::cvtColor(tmp_frame, formats::MatView(image_frame.get()),
|
||||
cv::COLOR_BGR2RGB);
|
||||
} else if (format_ == ImageFormat::SRGBA) {
|
||||
cv::cvtColor(tmp_frame, formats::MatView(image_frame.get()),
|
||||
cv::COLOR_BGRA2RGBA);
|
||||
}
|
||||
}
|
||||
cc->Outputs().Tag("VIDEO").Add(image_frame.release(), timestamp);
|
||||
decoded_frames_++;
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Close(CalculatorContext* cc) override {
|
||||
if (cap_ && cap_->isOpened()) {
|
||||
cap_->release();
|
||||
}
|
||||
if (decoded_frames_ != frame_count_) {
|
||||
LOG(WARNING) << "Not all the frames are decoded (total frames: "
|
||||
<< frame_count_ << " vs decoded frames: " << decoded_frames_
|
||||
<< ").";
|
||||
}
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
private:
|
||||
std::unique_ptr<cv::VideoCapture> cap_;
|
||||
int width_;
|
||||
int height_;
|
||||
int frame_count_;
|
||||
int decoded_frames_ = 0;
|
||||
ImageFormat::Format format_;
|
||||
};
|
||||
|
||||
REGISTER_CALCULATOR(OpenCvVideoDecoderCalculator);
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,161 @@
|
||||
// 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.
|
||||
|
||||
#include "mediapipe/framework/calculator_runner.h"
|
||||
#include "mediapipe/framework/deps/file_path.h"
|
||||
#include "mediapipe/framework/formats/image_frame.h"
|
||||
#include "mediapipe/framework/formats/image_frame_opencv.h"
|
||||
#include "mediapipe/framework/formats/video_stream_header.h"
|
||||
#include "mediapipe/framework/port/gmock.h"
|
||||
#include "mediapipe/framework/port/gtest.h"
|
||||
#include "mediapipe/framework/port/integral_types.h"
|
||||
#include "mediapipe/framework/port/logging.h"
|
||||
#include "mediapipe/framework/port/opencv_core_inc.h"
|
||||
#include "mediapipe/framework/port/parse_text_proto.h"
|
||||
#include "mediapipe/framework/port/status_matchers.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
namespace {
|
||||
|
||||
TEST(OpenCvVideoDecoderCalculatorTest, TestMp4Avc720pVideo) {
|
||||
CalculatorGraphConfig::Node node_config =
|
||||
ParseTextProtoOrDie<CalculatorGraphConfig::Node>(R"(
|
||||
calculator: "OpenCvVideoDecoderCalculator"
|
||||
input_side_packet: "INPUT_FILE_PATH:input_file_path"
|
||||
output_stream: "VIDEO:video"
|
||||
output_stream: "VIDEO_PRESTREAM:video_prestream")");
|
||||
CalculatorRunner runner(node_config);
|
||||
runner.MutableSidePackets()->Tag("INPUT_FILE_PATH") = MakePacket<std::string>(
|
||||
file::JoinPath("./",
|
||||
"/mediapipe/calculators/video/"
|
||||
"testdata/format_MP4_AVC720P_AAC.video"));
|
||||
MEDIAPIPE_EXPECT_OK(runner.Run());
|
||||
|
||||
EXPECT_EQ(runner.Outputs().Tag("VIDEO_PRESTREAM").packets.size(), 1);
|
||||
MEDIAPIPE_EXPECT_OK(runner.Outputs()
|
||||
.Tag("VIDEO_PRESTREAM")
|
||||
.packets[0]
|
||||
.ValidateAsType<VideoHeader>());
|
||||
const mediapipe::VideoHeader& header =
|
||||
runner.Outputs().Tag("VIDEO_PRESTREAM").packets[0].Get<VideoHeader>();
|
||||
EXPECT_EQ(ImageFormat::SRGB, header.format);
|
||||
EXPECT_EQ(1280, header.width);
|
||||
EXPECT_EQ(640, header.height);
|
||||
EXPECT_FLOAT_EQ(6.0f, header.duration);
|
||||
EXPECT_FLOAT_EQ(30.0f, header.frame_rate);
|
||||
EXPECT_EQ(180, runner.Outputs().Tag("VIDEO").packets.size());
|
||||
for (int i = 0; i < 180; ++i) {
|
||||
Packet image_frame_packet = runner.Outputs().Tag("VIDEO").packets[i];
|
||||
cv::Mat output_mat =
|
||||
formats::MatView(&(image_frame_packet.Get<ImageFrame>()));
|
||||
EXPECT_EQ(1280, output_mat.size().width);
|
||||
EXPECT_EQ(640, output_mat.size().height);
|
||||
EXPECT_EQ(3, output_mat.channels());
|
||||
cv::Scalar s = cv::mean(output_mat);
|
||||
for (int i = 0; i < 3; ++i) {
|
||||
EXPECT_GT(s[i], 0);
|
||||
EXPECT_LT(s[i], 255);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
TEST(OpenCvVideoDecoderCalculatorTest, TestFlvH264Video) {
|
||||
CalculatorGraphConfig::Node node_config =
|
||||
ParseTextProtoOrDie<CalculatorGraphConfig::Node>(R"(
|
||||
calculator: "OpenCvVideoDecoderCalculator"
|
||||
input_side_packet: "INPUT_FILE_PATH:input_file_path"
|
||||
output_stream: "VIDEO:video"
|
||||
output_stream: "VIDEO_PRESTREAM:video_prestream")");
|
||||
CalculatorRunner runner(node_config);
|
||||
runner.MutableSidePackets()->Tag("INPUT_FILE_PATH") = MakePacket<std::string>(
|
||||
file::JoinPath("./",
|
||||
"/mediapipe/calculators/video/"
|
||||
"testdata/format_FLV_H264_AAC.video"));
|
||||
MEDIAPIPE_EXPECT_OK(runner.Run());
|
||||
|
||||
EXPECT_EQ(runner.Outputs().Tag("VIDEO_PRESTREAM").packets.size(), 1);
|
||||
MEDIAPIPE_EXPECT_OK(runner.Outputs()
|
||||
.Tag("VIDEO_PRESTREAM")
|
||||
.packets[0]
|
||||
.ValidateAsType<VideoHeader>());
|
||||
const mediapipe::VideoHeader& header =
|
||||
runner.Outputs().Tag("VIDEO_PRESTREAM").packets[0].Get<VideoHeader>();
|
||||
EXPECT_EQ(ImageFormat::SRGB, header.format);
|
||||
EXPECT_EQ(640, header.width);
|
||||
EXPECT_EQ(320, header.height);
|
||||
// TODO: The actual header.duration is 6.0666666f and the frame_rate
|
||||
// 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());
|
||||
for (int i = 0; i < 180; ++i) {
|
||||
Packet image_frame_packet = runner.Outputs().Tag("VIDEO").packets[i];
|
||||
cv::Mat output_mat =
|
||||
formats::MatView(&(image_frame_packet.Get<ImageFrame>()));
|
||||
EXPECT_EQ(640, output_mat.size().width);
|
||||
EXPECT_EQ(320, output_mat.size().height);
|
||||
EXPECT_EQ(3, output_mat.channels());
|
||||
cv::Scalar s = cv::mean(output_mat);
|
||||
for (int i = 0; i < 3; ++i) {
|
||||
EXPECT_GT(s[i], 0);
|
||||
EXPECT_LT(s[i], 255);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
TEST(OpenCvVideoDecoderCalculatorTest, TestMkvVp8Video) {
|
||||
CalculatorGraphConfig::Node node_config =
|
||||
ParseTextProtoOrDie<CalculatorGraphConfig::Node>(R"(
|
||||
calculator: "OpenCvVideoDecoderCalculator"
|
||||
input_side_packet: "INPUT_FILE_PATH:input_file_path"
|
||||
output_stream: "VIDEO:video"
|
||||
output_stream: "VIDEO_PRESTREAM:video_prestream")");
|
||||
CalculatorRunner runner(node_config);
|
||||
runner.MutableSidePackets()->Tag("INPUT_FILE_PATH") = MakePacket<std::string>(
|
||||
file::JoinPath("./",
|
||||
"/mediapipe/calculators/video/"
|
||||
"testdata/format_MKV_VP8_VORBIS.video"));
|
||||
MEDIAPIPE_EXPECT_OK(runner.Run());
|
||||
|
||||
EXPECT_EQ(runner.Outputs().Tag("VIDEO_PRESTREAM").packets.size(), 1);
|
||||
MEDIAPIPE_EXPECT_OK(runner.Outputs()
|
||||
.Tag("VIDEO_PRESTREAM")
|
||||
.packets[0]
|
||||
.ValidateAsType<VideoHeader>());
|
||||
const mediapipe::VideoHeader& header =
|
||||
runner.Outputs().Tag("VIDEO_PRESTREAM").packets[0].Get<VideoHeader>();
|
||||
EXPECT_EQ(ImageFormat::SRGB, header.format);
|
||||
EXPECT_EQ(640, header.width);
|
||||
EXPECT_EQ(320, header.height);
|
||||
EXPECT_FLOAT_EQ(6.0f, header.duration);
|
||||
EXPECT_FLOAT_EQ(30.0f, header.frame_rate);
|
||||
EXPECT_EQ(180, runner.Outputs().Tag("VIDEO").packets.size());
|
||||
for (int i = 0; i < 180; ++i) {
|
||||
Packet image_frame_packet = runner.Outputs().Tag("VIDEO").packets[i];
|
||||
cv::Mat output_mat =
|
||||
formats::MatView(&(image_frame_packet.Get<ImageFrame>()));
|
||||
EXPECT_EQ(640, output_mat.size().width);
|
||||
EXPECT_EQ(320, output_mat.size().height);
|
||||
EXPECT_EQ(3, output_mat.channels());
|
||||
cv::Scalar s = cv::mean(output_mat);
|
||||
for (int i = 0; i < 3; ++i) {
|
||||
EXPECT_GT(s[i], 0);
|
||||
EXPECT_LT(s[i], 255);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,176 @@
|
||||
// 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.
|
||||
|
||||
#include <memory>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
#include "absl/strings/str_split.h"
|
||||
#include "mediapipe/calculators/video/opencv_video_encoder_calculator.pb.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/formats/image_frame.h"
|
||||
#include "mediapipe/framework/formats/image_frame_opencv.h"
|
||||
#include "mediapipe/framework/formats/video_stream_header.h"
|
||||
#include "mediapipe/framework/port/file_helpers.h"
|
||||
#include "mediapipe/framework/port/opencv_highgui_inc.h"
|
||||
#include "mediapipe/framework/port/opencv_imgproc_inc.h"
|
||||
#include "mediapipe/framework/port/opencv_video_inc.h"
|
||||
#include "mediapipe/framework/port/ret_check.h"
|
||||
#include "mediapipe/framework/port/source_location.h"
|
||||
#include "mediapipe/framework/port/status.h"
|
||||
#include "mediapipe/framework/port/status_builder.h"
|
||||
#include "mediapipe/framework/tool/status_util.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
// 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
|
||||
// mediapipe::ImageFrame).
|
||||
//
|
||||
// Example config to generate the output video file:
|
||||
//
|
||||
// node {
|
||||
// calculator: "OpenCvVideoEncoderCalculator"
|
||||
// input_stream: "VIDEO:video"
|
||||
// input_stream: "VIDEO_PRESTREAM:video_header"
|
||||
// input_side_packet: "OUTPUT_FILE_PATH:output_file_path"
|
||||
// node_options {
|
||||
// [type.googleapis.com/mediapipe.OpenCvVideoEncoderCalculatorOptions]: {
|
||||
// codec: "avc1"
|
||||
// video_format: "mp4"
|
||||
// }
|
||||
// }
|
||||
// }
|
||||
class OpenCvVideoEncoderCalculator : public CalculatorBase {
|
||||
public:
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc);
|
||||
::mediapipe::Status Open(CalculatorContext* cc) override;
|
||||
::mediapipe::Status Process(CalculatorContext* cc) override;
|
||||
::mediapipe::Status Close(CalculatorContext* cc) override;
|
||||
|
||||
private:
|
||||
::mediapipe::Status SetUpVideoWriter(float frame_rate, int width, int height);
|
||||
|
||||
std::string output_file_path_;
|
||||
int four_cc_;
|
||||
std::unique_ptr<cv::VideoWriter> writer_;
|
||||
};
|
||||
|
||||
::mediapipe::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->InputSidePackets().HasTag("OUTPUT_FILE_PATH"));
|
||||
cc->InputSidePackets().Tag("OUTPUT_FILE_PATH").Set<std::string>();
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status OpenCvVideoEncoderCalculator::Open(CalculatorContext* cc) {
|
||||
OpenCvVideoEncoderCalculatorOptions options =
|
||||
cc->Options<OpenCvVideoEncoderCalculatorOptions>();
|
||||
RET_CHECK(options.has_codec() && options.codec().length() == 4)
|
||||
<< "A 4-character codec code must be specified in "
|
||||
"OpenCvVideoEncoderCalculatorOptions";
|
||||
const char* codec_array = options.codec().c_str();
|
||||
four_cc_ = mediapipe::fourcc(codec_array[0], codec_array[1], codec_array[2],
|
||||
codec_array[3]);
|
||||
RET_CHECK(!options.video_format().empty())
|
||||
<< "Video format must be specified in "
|
||||
"OpenCvVideoEncoderCalculatorOptions";
|
||||
output_file_path_ =
|
||||
cc->InputSidePackets().Tag("OUTPUT_FILE_PATH").Get<std::string>();
|
||||
std::vector<std::string> splited_file_path =
|
||||
absl::StrSplit(output_file_path_, '.');
|
||||
RET_CHECK(splited_file_path.size() >= 2 &&
|
||||
splited_file_path[splited_file_path.size() - 1] ==
|
||||
options.video_format())
|
||||
<< "The output file path is invalid.";
|
||||
// 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")) {
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
return SetUpVideoWriter(options.fps(), options.width(), options.height());
|
||||
}
|
||||
|
||||
::mediapipe::Status OpenCvVideoEncoderCalculator::Process(
|
||||
CalculatorContext* cc) {
|
||||
if (cc->InputTimestamp() == Timestamp::PreStream()) {
|
||||
const VideoHeader& video_header =
|
||||
cc->Inputs().Tag("VIDEO_PRESTREAM").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>();
|
||||
ImageFormat::Format format = image_frame.Format();
|
||||
cv::Mat frame;
|
||||
if (format == ImageFormat::GRAY8) {
|
||||
frame = formats::MatView(&image_frame);
|
||||
if (frame.empty()) {
|
||||
return ::mediapipe::InvalidArgumentErrorBuilder(MEDIAPIPE_LOC)
|
||||
<< "Receive empty frame at timestamp "
|
||||
<< cc->Inputs().Tag("VIDEO").Value().Timestamp()
|
||||
<< " in OpenCvVideoEncoderCalculator::Process()";
|
||||
}
|
||||
} else {
|
||||
cv::Mat tmp_frame = formats::MatView(&image_frame);
|
||||
if (tmp_frame.empty()) {
|
||||
return ::mediapipe::InvalidArgumentErrorBuilder(MEDIAPIPE_LOC)
|
||||
<< "Receive empty frame at timestamp "
|
||||
<< cc->Inputs().Tag("VIDEO").Value().Timestamp()
|
||||
<< " in OpenCvVideoEncoderCalculator::Process()";
|
||||
}
|
||||
if (format == ImageFormat::SRGB) {
|
||||
cv::cvtColor(tmp_frame, frame, cv::COLOR_BGR2RGB);
|
||||
} else if (format == ImageFormat::SRGBA) {
|
||||
cv::cvtColor(tmp_frame, frame, cv::COLOR_BGRA2RGBA);
|
||||
} else {
|
||||
return ::mediapipe::InvalidArgumentErrorBuilder(MEDIAPIPE_LOC)
|
||||
<< "Unsupported image format: " << format;
|
||||
}
|
||||
}
|
||||
writer_->write(frame);
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status OpenCvVideoEncoderCalculator::Close(CalculatorContext* cc) {
|
||||
if (writer_ && writer_->isOpened()) {
|
||||
writer_->release();
|
||||
}
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status OpenCvVideoEncoderCalculator::SetUpVideoWriter(
|
||||
float frame_rate, int width, int height) {
|
||||
RET_CHECK(frame_rate > 0 && width > 0 && height > 0)
|
||||
<< "Invalid video metadata: frame_rate=" << frame_rate
|
||||
<< ", width=" << width << ", height=" << height;
|
||||
writer_ = absl::make_unique<cv::VideoWriter>(
|
||||
output_file_path_, four_cc_, frame_rate, cv::Size(width, height));
|
||||
if (!writer_->isOpened()) {
|
||||
return ::mediapipe::InvalidArgumentErrorBuilder(MEDIAPIPE_LOC)
|
||||
<< "Fail to open file at " << output_file_path_;
|
||||
}
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
REGISTER_CALCULATOR(OpenCvVideoEncoderCalculator);
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,37 @@
|
||||
// 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.
|
||||
|
||||
syntax = "proto2";
|
||||
|
||||
package mediapipe;
|
||||
|
||||
import "mediapipe/framework/calculator.proto";
|
||||
|
||||
message OpenCvVideoEncoderCalculatorOptions {
|
||||
extend CalculatorOptions {
|
||||
optional OpenCvVideoEncoderCalculatorOptions ext = 207936763;
|
||||
}
|
||||
// The 4-character code of the codec to encode the video.
|
||||
optional string codec = 1;
|
||||
|
||||
// The video format of the output video file.
|
||||
optional string video_format = 2;
|
||||
|
||||
// The frame rate in Hz at which the video frames are output.
|
||||
optional double fps = 3;
|
||||
|
||||
// Dimensions of the video in pixels.
|
||||
optional int32 width = 4;
|
||||
optional int32 height = 5;
|
||||
}
|
||||
@@ -0,0 +1,218 @@
|
||||
// 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.
|
||||
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/deps/file_path.h"
|
||||
#include "mediapipe/framework/formats/deleting_file.h"
|
||||
#include "mediapipe/framework/formats/image_frame.h"
|
||||
#include "mediapipe/framework/formats/image_frame_opencv.h"
|
||||
#include "mediapipe/framework/formats/video_stream_header.h"
|
||||
#include "mediapipe/framework/packet.h"
|
||||
#include "mediapipe/framework/port/gmock.h"
|
||||
#include "mediapipe/framework/port/gtest.h"
|
||||
#include "mediapipe/framework/port/integral_types.h"
|
||||
#include "mediapipe/framework/port/logging.h"
|
||||
#include "mediapipe/framework/port/opencv_highgui_inc.h"
|
||||
#include "mediapipe/framework/port/opencv_imgproc_inc.h"
|
||||
#include "mediapipe/framework/port/opencv_video_inc.h"
|
||||
#include "mediapipe/framework/port/parse_text_proto.h"
|
||||
#include "mediapipe/framework/port/status_matchers.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
namespace {
|
||||
// Temporarily disable the test.
|
||||
// TODO: Investigate the “Could not open codec 'libx264'” error with
|
||||
// opencv2.
|
||||
TEST(OpenCvVideoEncoderCalculatorTest, DISABLED_TestMp4Avc720pVideo) {
|
||||
CalculatorGraphConfig config = ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
|
||||
node {
|
||||
calculator: "OpenCvVideoDecoderCalculator"
|
||||
input_side_packet: "INPUT_FILE_PATH:input_file_path"
|
||||
output_stream: "VIDEO:video"
|
||||
output_stream: "VIDEO_PRESTREAM:video_prestream"
|
||||
}
|
||||
node {
|
||||
calculator: "OpenCvVideoEncoderCalculator"
|
||||
input_stream: "VIDEO:video"
|
||||
input_stream: "VIDEO_PRESTREAM:video_prestream"
|
||||
input_side_packet: "OUTPUT_FILE_PATH:output_file_path"
|
||||
node_options {
|
||||
[type.googleapis.com/mediapipe.OpenCvVideoEncoderCalculatorOptions]: {
|
||||
codec: "avc1"
|
||||
video_format: "mp4"
|
||||
}
|
||||
}
|
||||
}
|
||||
)");
|
||||
std::map<std::string, Packet> input_side_packets;
|
||||
input_side_packets["input_file_path"] = MakePacket<std::string>(
|
||||
file::JoinPath("./",
|
||||
"/mediapipe/calculators/video/"
|
||||
"testdata/format_MP4_AVC720P_AAC.video"));
|
||||
const std::string output_file_path = "/tmp/tmp_video.mp4";
|
||||
DeletingFile deleting_file(output_file_path, true);
|
||||
input_side_packets["output_file_path"] =
|
||||
MakePacket<std::string>(output_file_path);
|
||||
CalculatorGraph graph;
|
||||
MEDIAPIPE_ASSERT_OK(graph.Initialize(config, input_side_packets));
|
||||
StatusOrPoller status_or_poller =
|
||||
graph.AddOutputStreamPoller("video_prestream");
|
||||
ASSERT_TRUE(status_or_poller.ok());
|
||||
OutputStreamPoller poller = std::move(status_or_poller.ValueOrDie());
|
||||
|
||||
MEDIAPIPE_ASSERT_OK(graph.StartRun({}));
|
||||
Packet packet;
|
||||
while (poller.Next(&packet)) {
|
||||
}
|
||||
MEDIAPIPE_ASSERT_OK(graph.WaitUntilDone());
|
||||
const VideoHeader& video_header = packet.Get<VideoHeader>();
|
||||
|
||||
// Checks the generated video file has the same width, height, fps, and
|
||||
// duration as the original one.
|
||||
cv::VideoCapture cap(output_file_path);
|
||||
ASSERT_TRUE(cap.isOpened());
|
||||
EXPECT_EQ(video_header.width,
|
||||
static_cast<int>(cap.get(cv::CAP_PROP_FRAME_WIDTH)));
|
||||
EXPECT_EQ(video_header.height,
|
||||
static_cast<int>(cap.get(cv::CAP_PROP_FRAME_HEIGHT)));
|
||||
EXPECT_EQ(video_header.frame_rate,
|
||||
static_cast<double>(cap.get(cv::CAP_PROP_FPS)));
|
||||
EXPECT_EQ(video_header.duration,
|
||||
static_cast<int>(cap.get(cv::CAP_PROP_FRAME_COUNT) /
|
||||
cap.get(cv::CAP_PROP_FPS)));
|
||||
}
|
||||
|
||||
TEST(OpenCvVideoEncoderCalculatorTest, TestFlvH264Video) {
|
||||
CalculatorGraphConfig config = ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
|
||||
node {
|
||||
calculator: "OpenCvVideoDecoderCalculator"
|
||||
input_side_packet: "INPUT_FILE_PATH:input_file_path"
|
||||
output_stream: "VIDEO:video"
|
||||
output_stream: "VIDEO_PRESTREAM:video_prestream"
|
||||
}
|
||||
node {
|
||||
calculator: "OpenCvVideoEncoderCalculator"
|
||||
input_stream: "VIDEO:video"
|
||||
input_stream: "VIDEO_PRESTREAM:video_prestream"
|
||||
input_side_packet: "OUTPUT_FILE_PATH:output_file_path"
|
||||
node_options {
|
||||
[type.googleapis.com/mediapipe.OpenCvVideoEncoderCalculatorOptions]: {
|
||||
codec: "MJPG"
|
||||
video_format: "avi"
|
||||
}
|
||||
}
|
||||
}
|
||||
)");
|
||||
std::map<std::string, Packet> input_side_packets;
|
||||
input_side_packets["input_file_path"] = MakePacket<std::string>(
|
||||
file::JoinPath("./",
|
||||
"/mediapipe/calculators/video/"
|
||||
"testdata/format_FLV_H264_AAC.video"));
|
||||
const std::string output_file_path = "/tmp/tmp_video.avi";
|
||||
DeletingFile deleting_file(output_file_path, true);
|
||||
input_side_packets["output_file_path"] =
|
||||
MakePacket<std::string>(output_file_path);
|
||||
CalculatorGraph graph;
|
||||
MEDIAPIPE_ASSERT_OK(graph.Initialize(config, input_side_packets));
|
||||
StatusOrPoller status_or_poller =
|
||||
graph.AddOutputStreamPoller("video_prestream");
|
||||
ASSERT_TRUE(status_or_poller.ok());
|
||||
OutputStreamPoller poller = std::move(status_or_poller.ValueOrDie());
|
||||
|
||||
MEDIAPIPE_ASSERT_OK(graph.StartRun({}));
|
||||
Packet packet;
|
||||
while (poller.Next(&packet)) {
|
||||
}
|
||||
MEDIAPIPE_ASSERT_OK(graph.WaitUntilDone());
|
||||
const VideoHeader& video_header = packet.Get<VideoHeader>();
|
||||
|
||||
// Checks the generated video file has the same width, height, fps, and
|
||||
// duration as the original one.
|
||||
cv::VideoCapture cap(output_file_path);
|
||||
ASSERT_TRUE(cap.isOpened());
|
||||
EXPECT_EQ(video_header.width,
|
||||
static_cast<int>(cap.get(cv::CAP_PROP_FRAME_WIDTH)));
|
||||
EXPECT_EQ(video_header.height,
|
||||
static_cast<int>(cap.get(cv::CAP_PROP_FRAME_HEIGHT)));
|
||||
// TODO: The actual header.duration is 6.0666666f and the frame_rate
|
||||
// can be either 30.30303f (with opencv2) or 30f (with opencv3 and opencv4).
|
||||
// EXPECT_EQ(video_header.frame_rate,
|
||||
// static_cast<double>(cap.get(cv::CAP_PROP_FPS)));
|
||||
// EXPECT_EQ(video_header.duration,
|
||||
// static_cast<int>(cap.get(cv::CAP_PROP_FRAME_COUNT) /
|
||||
// cap.get(cv::CAP_PROP_FPS)));
|
||||
}
|
||||
|
||||
TEST(OpenCvVideoEncoderCalculatorTest, TestMkvVp8Video) {
|
||||
CalculatorGraphConfig config = ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
|
||||
node {
|
||||
calculator: "OpenCvVideoDecoderCalculator"
|
||||
input_side_packet: "INPUT_FILE_PATH:input_file_path"
|
||||
output_stream: "VIDEO:video"
|
||||
output_stream: "VIDEO_PRESTREAM:video_prestream"
|
||||
}
|
||||
node {
|
||||
calculator: "OpenCvVideoEncoderCalculator"
|
||||
input_stream: "VIDEO:video"
|
||||
input_stream: "VIDEO_PRESTREAM:video_prestream"
|
||||
input_side_packet: "OUTPUT_FILE_PATH:output_file_path"
|
||||
node_options {
|
||||
[type.googleapis.com/mediapipe.OpenCvVideoEncoderCalculatorOptions]: {
|
||||
codec: "PIM1"
|
||||
video_format: "mkv"
|
||||
}
|
||||
}
|
||||
}
|
||||
)");
|
||||
std::map<std::string, Packet> input_side_packets;
|
||||
input_side_packets["input_file_path"] = MakePacket<std::string>(
|
||||
file::JoinPath("./",
|
||||
"/mediapipe/calculators/video/"
|
||||
"testdata/format_MKV_VP8_VORBIS.video"));
|
||||
const std::string output_file_path = "/tmp/tmp_video.mkv";
|
||||
DeletingFile deleting_file(output_file_path, true);
|
||||
input_side_packets["output_file_path"] =
|
||||
MakePacket<std::string>(output_file_path);
|
||||
CalculatorGraph graph;
|
||||
MEDIAPIPE_ASSERT_OK(graph.Initialize(config, input_side_packets));
|
||||
StatusOrPoller status_or_poller =
|
||||
graph.AddOutputStreamPoller("video_prestream");
|
||||
ASSERT_TRUE(status_or_poller.ok());
|
||||
OutputStreamPoller poller = std::move(status_or_poller.ValueOrDie());
|
||||
|
||||
MEDIAPIPE_ASSERT_OK(graph.StartRun({}));
|
||||
Packet packet;
|
||||
while (poller.Next(&packet)) {
|
||||
}
|
||||
MEDIAPIPE_ASSERT_OK(graph.WaitUntilDone());
|
||||
const VideoHeader& video_header = packet.Get<VideoHeader>();
|
||||
|
||||
// Checks the generated video file has the same width, height, fps, and
|
||||
// duration as the original one.
|
||||
cv::VideoCapture cap(output_file_path);
|
||||
ASSERT_TRUE(cap.isOpened());
|
||||
EXPECT_EQ(video_header.width,
|
||||
static_cast<int>(cap.get(cv::CAP_PROP_FRAME_WIDTH)));
|
||||
EXPECT_EQ(video_header.height,
|
||||
static_cast<int>(cap.get(cv::CAP_PROP_FRAME_HEIGHT)));
|
||||
EXPECT_EQ(video_header.frame_rate,
|
||||
static_cast<double>(cap.get(cv::CAP_PROP_FPS)));
|
||||
EXPECT_EQ(video_header.duration,
|
||||
static_cast<int>(cap.get(cv::CAP_PROP_FRAME_COUNT) /
|
||||
cap.get(cv::CAP_PROP_FPS)));
|
||||
}
|
||||
|
||||
} // namespace
|
||||
} // namespace mediapipe
|
||||
+26
@@ -0,0 +1,26 @@
|
||||
# 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.
|
||||
#
|
||||
|
||||
licenses(["notice"]) # Apache 2.0
|
||||
|
||||
filegroup(
|
||||
name = "test_videos",
|
||||
srcs = [
|
||||
"format_FLV_H264_AAC.video",
|
||||
"format_MKV_VP8_VORBIS.video",
|
||||
"format_MP4_AVC720P_AAC.video",
|
||||
],
|
||||
visibility = ["//visibility:public"],
|
||||
)
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,50 @@
|
||||
#
|
||||
# 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.
|
||||
|
||||
licenses(["notice"]) # Apache 2.0
|
||||
|
||||
package(default_visibility = ["//mediapipe/calculators/video:__subpackages__"])
|
||||
|
||||
exports_files(["LICENSE"])
|
||||
|
||||
load("//mediapipe/framework/port:build_config.bzl", "mediapipe_cc_proto_library")
|
||||
|
||||
proto_library(
|
||||
name = "flow_quantizer_model_proto",
|
||||
srcs = ["flow_quantizer_model.proto"],
|
||||
)
|
||||
|
||||
mediapipe_cc_proto_library(
|
||||
name = "flow_quantizer_model_cc_proto",
|
||||
srcs = ["flow_quantizer_model.proto"],
|
||||
visibility = ["//mediapipe:__subpackages__"],
|
||||
deps = [":flow_quantizer_model_proto"],
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "flow_quantizer_model",
|
||||
srcs = ["flow_quantizer_model.cc"],
|
||||
hdrs = ["flow_quantizer_model.h"],
|
||||
deps = [
|
||||
"//mediapipe/calculators/video/tool:flow_quantizer_model_cc_proto",
|
||||
"//mediapipe/framework:type_map",
|
||||
"//mediapipe/framework/formats:image_frame",
|
||||
"//mediapipe/framework/formats/motion:optical_flow_field",
|
||||
"//mediapipe/framework/port:integral_types",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
"//mediapipe/framework/tool:status_util",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
@@ -0,0 +1,78 @@
|
||||
// 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.
|
||||
|
||||
#include "mediapipe/calculators/video/tool/flow_quantizer_model.h"
|
||||
|
||||
#include "mediapipe/framework/port/ret_check.h"
|
||||
#include "mediapipe/framework/type_map.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
// Uniform normalization to 0-255.
|
||||
uint8 FlowQuantizerModel::Apply(const float val, const int channel) const {
|
||||
CHECK_LT(channel, model_.min_value_size());
|
||||
const auto& min_value = model_.min_value(channel);
|
||||
const auto& max_value = model_.max_value(channel);
|
||||
QCHECK_GT(max_value, min_value);
|
||||
float res = (val - min_value) / (max_value - min_value);
|
||||
if (res < 0.0) {
|
||||
res = 0.0;
|
||||
} else if (res > 1.0) {
|
||||
res = 1.0;
|
||||
}
|
||||
return static_cast<uint8>(res * 255);
|
||||
}
|
||||
|
||||
void FlowQuantizerModel::LoadFromProto(const QuantizerModelData& data) {
|
||||
QCHECK_GT(data.max_value(0), data.min_value(0));
|
||||
QCHECK_GT(data.max_value(1), data.min_value(1));
|
||||
|
||||
model_ = data;
|
||||
}
|
||||
|
||||
const QuantizerModelData& FlowQuantizerModel::GetModelData() const {
|
||||
return model_;
|
||||
}
|
||||
|
||||
// Used for training, update the (min, max) range. We want to estimate the range
|
||||
// of optical flow fields (Theorectically it is (-num_pixels_along_diag,
|
||||
// num_pixels_along_diag).
|
||||
// TODO: Taking the min and max over all training flow fields might be
|
||||
// sensitive to noise. We should use more robust statistics.
|
||||
void FlowQuantizerModel::AddSampleFlowField(const OpticalFlowField& flow) {
|
||||
CHECK_EQ(model_.min_value_size(), 2);
|
||||
const cv::Mat_<cv::Point2f>& flow_mat = flow.flow_data();
|
||||
for (int i = 0; i != flow.width(); ++i) {
|
||||
for (int j = 0; j != flow.height(); ++j) {
|
||||
const auto& x = flow_mat.at<cv::Point2f>(i, j).x;
|
||||
const auto& y = flow_mat.at<cv::Point2f>(i, j).y;
|
||||
// Always use the minimum and maximum value occurred in training flow
|
||||
// fields.
|
||||
model_.set_min_value(0, std::min<float>(x, model_.min_value(0)));
|
||||
model_.set_min_value(1, std::min<float>(y, model_.min_value(1)));
|
||||
model_.set_max_value(0, std::max<float>(x, model_.max_value(0)));
|
||||
model_.set_max_value(1, std::max<float>(y, model_.max_value(1)));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void FlowQuantizerModel::Init() {
|
||||
model_.Clear();
|
||||
// Initialize the values.
|
||||
for (int i = 0; i != 2; ++i) {
|
||||
model_.add_min_value(std::numeric_limits<float>::max());
|
||||
model_.add_max_value(-std::numeric_limits<float>::max());
|
||||
}
|
||||
}
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,47 @@
|
||||
// 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.
|
||||
|
||||
// Quantization model to convert a real value float number (flow field) to a
|
||||
// 8-bit discrete number.
|
||||
#ifndef MEDIAPIPE_CALCULATORS_VIDEO_TOOL_FLOW_QUANTIZER_MODEL_H_
|
||||
#define MEDIAPIPE_CALCULATORS_VIDEO_TOOL_FLOW_QUANTIZER_MODEL_H_
|
||||
|
||||
#include "mediapipe/calculators/video/tool/flow_quantizer_model.pb.h"
|
||||
#include "mediapipe/framework/formats/motion/optical_flow_field.h"
|
||||
#include "mediapipe/framework/port/integral_types.h"
|
||||
#include "mediapipe/framework/tool/status_util.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
class FlowQuantizerModel {
|
||||
public:
|
||||
// Initializes the model proto.
|
||||
void Init();
|
||||
// Quantizes flow field with the model.
|
||||
uint8 Apply(const float val, const int channel) const;
|
||||
// Loads model from proto.
|
||||
void LoadFromProto(const QuantizerModelData& data);
|
||||
// Gets proto from model.
|
||||
const QuantizerModelData& GetModelData() const;
|
||||
// Used in training. Updates the model proto by reading the flow fields.
|
||||
// TODO: This model is currently manually set. Need to find a way to
|
||||
// learn from flow fields directly.
|
||||
void AddSampleFlowField(const OpticalFlowField& flow);
|
||||
|
||||
private:
|
||||
QuantizerModelData model_;
|
||||
};
|
||||
} // namespace mediapipe
|
||||
|
||||
#endif // MEDIAPIPE_CALCULATORS_VIDEO_TOOL_FLOW_QUANTIZER_MODEL_H_
|
||||
@@ -0,0 +1,24 @@
|
||||
// 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.
|
||||
|
||||
syntax = "proto2";
|
||||
|
||||
package mediapipe;
|
||||
|
||||
// Message storing min value and max value for normalization in all channels.
|
||||
message QuantizerModelData {
|
||||
// For all channels.
|
||||
repeated float min_value = 1;
|
||||
repeated float max_value = 2;
|
||||
}
|
||||
Reference in New Issue
Block a user