demo.cc

// 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.
//
// An example of sending OpenCV webcam frames into a MediaPipe graph.
#include <cstdlib>
#include <fstream>

#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/formats/image_frame.h"
#include "mediapipe/framework/formats/image_frame_opencv.h"
#include "mediapipe/framework/port/commandlineflags.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/parse_text_proto.h"
#include "mediapipe/framework/port/status.h"
#include "mediapipe/framework/formats/rect.pb.h"
#include "mediapipe/framework/formats/landmark.pb.h"
#include "json.hpp"
#include "pipeline.h"

using json = nlohmann::json;

constexpr char kInputStream[] = "input_video";
constexpr char kOutputStream[] = "output_video";
constexpr char kOutputStream_face[] = "face_rect";
constexpr char kOutputStream_eye_contour_landmarks[] = "left_eye_contour_landmarks";
constexpr char kWindowName[] = "MediaPipe";

DEFINE_string(
    calculator_graph_config_file, "",
    "Name of file containing text format CalculatorGraphConfig proto.");
DEFINE_string(input_video_path, "",
              "Full path of video to load. "
              "If not provided, attempt to use a webcam.");
DEFINE_string(output_video_path, "",
              "Full path of where to save result (.mp4 only). "
              "If not provided, show result in a window.");

int pipe_cpp_to_py[2];
int pipe_py_to_cpp[2];

bool set_pipeline()
{
  if (::pipe(pipe_cpp_to_py) || ::pipe(pipe_py_to_cpp))
  {
    std::cout << "Couldn't open pipes" << std::endl;
    return false;
  }

  // ::fcntl(pipe_py_to_cpp[1], F_SETFL, O_NONBLOCK);
  // ::fcntl(pipe_cpp_to_py[1], F_SETFL, O_NONBLOCK);
  ::fcntl(pipe_cpp_to_py[1], F_SETPIPE_SZ, 640 * 480 * 3);

  pid_t pid = fork();

  if (pid == 0)
  {
    // ::close(pipe_py_to_cpp[0]);
    ::close(pipe_cpp_to_py[1]);
    std::ostringstream oss;

    oss << "export PY_READ_FD=" << pipe_cpp_to_py[0] << " && "
        << "export PY_WRITE_FD=" << pipe_py_to_cpp[1] << " && "
        << "export PYTHONUNBUFFERED=true && " // Force stdin, stdout and stderr to be totally unbuffered.
        << "python `python path here`/py_libs/main.py";

    ::system(oss.str().c_str());
    // ::close(pipe_py_to_cpp[1]);
    // ::close(pipe_cpp_to_py[0]);

    return false;
  }
  else if (pid < 0)
  {
    std::cout << "Fork failed." << std::endl;
    return false;
  }
  else
  {
    std::cout << "parent process" << std::endl;
    ::close(pipe_py_to_cpp[1]);
    // ::close(pipe_cpp_to_py[0]);
    return true;
    //   ::close(pipe_py_to_cpp[0]);
    //   ::close(pipe_cpp_to_py[1]);
  }
}

void end_pipeline()
{
  transmit_msg(pipe_py_to_cpp[0], pipe_cpp_to_py[1], "close");

  ::close(pipe_cpp_to_py[0]);
  ::close(pipe_cpp_to_py[1]);

  ::close(pipe_py_to_cpp[0]);
}

void processFrame(cv::Mat &frame)
{
  auto msg_pipe = check_pipe_status(pipe_py_to_cpp[0]);

  if (!msg_pipe.empty())
  {
    // transmit raw frame
    const bool load_video = !FLAGS_input_video_path.empty();

    cv::Mat flipMat;
    if (load_video)
    {
      flipMat = frame;
    }
    else
    {
      cv::flip(frame, flipMat, /*flipcode=HORIZONTAL*/ 1);
    }

    int mat_size = frame.rows * frame.cols * frame.elemSize();
    transmit_buf(pipe_py_to_cpp[0], pipe_cpp_to_py[1], flipMat.data, mat_size);
  }
}

::mediapipe::Status RunMPPGraph()
{
  std::string calculator_graph_config_contents;
  MP_RETURN_IF_ERROR(mediapipe::file::GetContents(
      FLAGS_calculator_graph_config_file, &calculator_graph_config_contents));
  LOG(INFO) << "Get calculator graph config contents: "
            << calculator_graph_config_contents;
  mediapipe::CalculatorGraphConfig config =
      mediapipe::ParseTextProtoOrDie<mediapipe::CalculatorGraphConfig>(
          calculator_graph_config_contents);

  LOG(INFO) << "Initialize the calculator graph.";
  mediapipe::CalculatorGraph graph;
  MP_RETURN_IF_ERROR(graph.Initialize(config));

  LOG(INFO) << "Initialize the camera or load the video.";
  cv::VideoCapture capture;
  const bool load_video = !FLAGS_input_video_path.empty();
  if (load_video)
  {
    capture.open(FLAGS_input_video_path);
  }
  else
  {
    capture.open(0);
  }
  RET_CHECK(capture.isOpened());

  cv::VideoWriter writer;
  const bool save_video = !FLAGS_output_video_path.empty();
  if (!save_video)
  {
    cv::namedWindow(kWindowName, /*flags=WINDOW_AUTOSIZE*/ 1);
#if (CV_MAJOR_VERSION >= 3) && (CV_MINOR_VERSION >= 2)
    capture.set(cv::CAP_PROP_FRAME_WIDTH, 640);
    capture.set(cv::CAP_PROP_FRAME_HEIGHT, 480);
    capture.set(cv::CAP_PROP_FPS, 30);
#endif
  }

  LOG(INFO) << "Start running the calculator graph.";
  ASSIGN_OR_RETURN(mediapipe::OutputStreamPoller poller,
                   graph.AddOutputStreamPoller(kOutputStream));
  ASSIGN_OR_RETURN(mediapipe::OutputStreamPoller face_count_poller,
                   graph.AddOutputStreamPoller("face_count"));
  ASSIGN_OR_RETURN(mediapipe::OutputStreamPoller poller_face,
                   graph.AddOutputStreamPoller(kOutputStream_face));
  ASSIGN_OR_RETURN(mediapipe::OutputStreamPoller poller_eye_contour_landmarks,
                   graph.AddOutputStreamPoller(kOutputStream_eye_contour_landmarks));

  MP_RETURN_IF_ERROR(graph.StartRun({}));

  LOG(INFO) << "Start grabbing and processing frames.";
  bool grab_frames = true;
  while (grab_frames)
  {
    // Capture opencv camera or video frame.
    cv::Mat camera_frame_raw;
    capture >> camera_frame_raw;
    if (camera_frame_raw.empty())
      break; // End of video.

    // update current frame pos
    _current_frame_pos = capture.get(cv::CAP_PROP_POS_FRAMES);

    if (camera_frame_raw.cols > 1500 || camera_frame_raw.rows > 1000)
    {
      std::cout << "Too high resolution of frame to process." << std::endl;
      break;
    }

    cv::Mat camera_frame;
    cv::cvtColor(camera_frame_raw, camera_frame, cv::COLOR_BGR2RGB);
    if (!load_video)
    {
      cv::flip(camera_frame, camera_frame, /*flipcode=HORIZONTAL*/ 1);
    }

    // Wrap Mat into an ImageFrame.
    auto input_frame = absl::make_unique<mediapipe::ImageFrame>(
        mediapipe::ImageFormat::SRGB, camera_frame.cols, camera_frame.rows,
        mediapipe::ImageFrame::kDefaultAlignmentBoundary);
    cv::Mat input_frame_mat = mediapipe::formats::MatView(input_frame.get());
    camera_frame.copyTo(input_frame_mat);

    // Send image packet into the graph.
    size_t frame_timestamp_us =
        (double)cv::getTickCount() / (double)cv::getTickFrequency() * 1e6;
    MP_RETURN_IF_ERROR(graph.AddPacketToInputStream(
        kInputStream, mediapipe::Adopt(input_frame.release())
                          .At(mediapipe::Timestamp(frame_timestamp_us))));

    // Get the graph result packet, or stop if that fails.
    mediapipe::Packet packet;
    if (!poller.Next(&packet))
      break;

    mediapipe::Packet face_count_packet;
    if (!face_count_poller.Next(&face_count_packet))
      break;

    auto face_count = face_count_packet.Get<int>();

    if (face_count > 0)
    {
      // your code here ...
      processFrame(camera_frame_raw);
    }

    auto &output_frame = packet.Get<mediapipe::ImageFrame>();

    // Convert back to opencv for display or saving.
    cv::Mat output_frame_mat = mediapipe::formats::MatView(&output_frame);
    cv::cvtColor(output_frame_mat, output_frame_mat, cv::COLOR_RGB2BGR);

    if (save_video)
    {
      if (!writer.isOpened())
      {
        LOG(INFO) << "Prepare video writer.";
        writer.open(FLAGS_output_video_path,
                    mediapipe::fourcc('a', 'v', 'c', '1'), // .mp4
                    capture.get(cv::CAP_PROP_FPS), output_frame_mat.size());
        RET_CHECK(writer.isOpened());
      }
      writer.write(output_frame_mat);
    }
    else
    {
      cv::imshow(kWindowName, output_frame_mat);
      // Press any key to exit.
      const int pressed_key = cv::waitKey(5);
      if (pressed_key >= 0 && pressed_key != 255)
        grab_frames = false;
    }
  }

  LOG(INFO) << "Shutting down.";

  // send sutting down message through pipe
  end_pipeline();

  if (writer.isOpened())
    writer.release();
  MP_RETURN_IF_ERROR(graph.CloseInputStream(kInputStream));
  return graph.WaitUntilDone();
}

int main(int argc, char **argv)
{
  google::InitGoogleLogging(argv[0]);
  gflags::ParseCommandLineFlags(&argc, &argv, true);

  if (!set_pipeline())
  {
    return EXIT_SUCCESS;
  }

  ::mediapipe::Status run_status = RunMPPGraph();
  if (!run_status.ok())
  {
    LOG(ERROR) << "Failed to run the graph: " << run_status.message();
    return EXIT_FAILURE;
  }
  else
  {
    LOG(INFO) << "Success!";
  }

  return EXIT_SUCCESS;
}