Project import generated by Copybara.

GitOrigin-RevId: e3566e5029af25b0fc4b1071a49e49ae20aa5df6
2019-12-02 17:54:10 -08:00
parent 446d7cf6b6
commit 137867d088
74 changed files with 909 additions and 504 deletions
@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.

+#include <stdlib.h>
+
 #include "mediapipe/framework/calculator_framework.h"
 #include "mediapipe/framework/formats/image_format.pb.h"
 #include "mediapipe/framework/formats/image_frame.h"
@@ -66,6 +68,20 @@ ImageFormat::Format GetImageFormat(int num_channels) {
 //   output_stream: "VIDEO:video_frames"
 //   output_stream: "VIDEO_PRESTREAM:video_header"
 // }
+//
+// OpenCV's VideoCapture doesn't decode audio tracks. If the audio tracks need
+// to be saved, specify an output side packet with tag "SAVED_AUDIO_PATH".
+// The calculator will call FFmpeg binary to save audio tracks as an aac file.
+//
+// Example config:
+// node {
+//   calculator: "OpenCvVideoDecoderCalculator"
+//   input_side_packet: "INPUT_FILE_PATH:input_file_path"
+//   output_side_packet: "SAVED_AUDIO_PATH:audio_path"
+//   output_stream: "VIDEO:video_frames"
+//   output_stream: "VIDEO_PRESTREAM:video_header"
+// }
+//
 class OpenCvVideoDecoderCalculator : public CalculatorBase {
 public:
  static ::mediapipe::Status GetContract(CalculatorContract* cc) {
@@ -74,6 +90,9 @@ class OpenCvVideoDecoderCalculator : public CalculatorBase {
    if (cc->Outputs().HasTag("VIDEO_PRESTREAM")) {
      cc->Outputs().Tag("VIDEO_PRESTREAM").Set<VideoHeader>();
    }
+    if (cc->OutputSidePackets().HasTag("SAVED_AUDIO_PATH")) {
+      cc->OutputSidePackets().Tag("SAVED_AUDIO_PATH").Set<std::string>();
+    }
    return ::mediapipe::OkStatus();
  }

@@ -127,6 +146,25 @@ class OpenCvVideoDecoderCalculator : public CalculatorBase {
    }
    // Rewind to the very first frame.
    cap_->set(cv::CAP_PROP_POS_AVI_RATIO, 0);
+
+    if (cc->OutputSidePackets().HasTag("SAVED_AUDIO_PATH")) {
+#ifdef HAVE_FFMPEG
+      std::string saved_audio_path = std::tmpnam(nullptr);
+      system(absl::StrCat("ffmpeg -nostats -loglevel 0 -i ", input_file_path,
+                          " -vn -f adts ", saved_audio_path)
+                 .c_str());
+      cc->OutputSidePackets()
+          .Tag("SAVED_AUDIO_PATH")
+          .Set(MakePacket<std::string>(saved_audio_path));
+
+#else
+      return ::mediapipe::InvalidArgumentErrorBuilder(MEDIAPIPE_LOC)
+             << "OpenCVVideoDecoderCalculator can't save the audio file "
+                "because FFmpeg is not installed. Please remove "
+                "output_side_packet: \"SAVED_AUDIO_PATH\" from the node "
+                "config.";
+#endif
+    }
    return ::mediapipe::OkStatus();
  }

@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.

+#include <stdlib.h>
+
 #include <memory>
 #include <string>
 #include <vector>
@@ -39,8 +41,7 @@ namespace mediapipe {
 // packet. Currently, the calculator only supports one video stream (in
 // mediapipe::ImageFrame).
 //
-// Example config to generate the output video file:
-//
+// Example config:
 // node {
 //   calculator: "OpenCvVideoEncoderCalculator"
 //   input_stream: "VIDEO:video"
@@ -53,6 +54,26 @@ namespace mediapipe {
 //     }
 //   }
 // }
+//
+// OpenCV's VideoWriter doesn't encode audio. If an input side packet with tag
+// "AUDIO_FILE_PATH" is specified, the calculator will call FFmpeg binary to
+// attach the audio file to the video as the last step in Close().
+//
+// Example config:
+// node {
+//   calculator: "OpenCvVideoEncoderCalculator"
+//   input_stream: "VIDEO:video"
+//   input_stream: "VIDEO_PRESTREAM:video_header"
+//   input_side_packet: "OUTPUT_FILE_PATH:output_file_path"
+//   input_side_packet: "AUDIO_FILE_PATH:audio_path"
+//   node_options {
+//     [type.googleapis.com/mediapipe.OpenCvVideoEncoderCalculatorOptions]: {
+//        codec: "avc1"
+//        video_format: "mp4"
+//     }
+//   }
+// }
+//
 class OpenCvVideoEncoderCalculator : public CalculatorBase {
 public:
  static ::mediapipe::Status GetContract(CalculatorContract* cc);
@@ -77,6 +98,9 @@ class OpenCvVideoEncoderCalculator : public CalculatorBase {
  }
  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>();
+  }
  return ::mediapipe::OkStatus();
 }

@@ -155,6 +179,27 @@ class OpenCvVideoEncoderCalculator : public CalculatorBase {
  if (writer_ && writer_->isOpened()) {
    writer_->release();
  }
+  if (cc->InputSidePackets().HasTag("AUDIO_FILE_PATH")) {
+#ifdef HAVE_FFMPEG
+    const std::string& audio_file_path =
+        cc->InputSidePackets().Tag("AUDIO_FILE_PATH").Get<std::string>();
+    // A temp output file is needed because FFmpeg can't do in-place editing.
+    const std::string temp_file_path = std::tmpnam(nullptr);
+    system(absl::StrCat("mv ", output_file_path_, " ", temp_file_path,
+                        "&& ffmpeg -nostats -loglevel 0 -i ", temp_file_path,
+                        " -i ", audio_file_path,
+                        "  -c copy -map 0:v:0 -map 1:a:0 ", output_file_path_,
+                        "&& rm ", temp_file_path)
+               .c_str());
+
+#else
+    return ::mediapipe::InvalidArgumentErrorBuilder(MEDIAPIPE_LOC)
+           << "OpenCVVideoEncoderCalculator can't attach the audio tracks to "
+              "the video because FFmpeg is not installed. Please remove "
+              "input_side_packet: \"AUDIO_FILE_PATH\" from the node "
+              "config.";
+#endif
+  }
  return ::mediapipe::OkStatus();
 }