From 5de0ecfcaf716bc919f5e4ee8661ed8b6b34a709 Mon Sep 17 00:00:00 2001
From: Li Zhang <lzhang329@gmail.com>
Date: Tue, 7 Feb 2023 20:27:43 +0800
Subject: [PATCH] [Fix] Add an option to flip webcam inputs for pose tracker
 demo (#1725)

---
 demo/csrc/cpp/pose_tracker.cxx |  5 +++--
 demo/csrc/cpp/utils/mediaio.h  | 13 +++++++++----
 2 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/demo/csrc/cpp/pose_tracker.cxx b/demo/csrc/cpp/pose_tracker.cxx
index c57cc3e4a4..c50bdfbe14 100644
--- a/demo/csrc/cpp/pose_tracker.cxx
+++ b/demo/csrc/cpp/pose_tracker.cxx
@@ -14,6 +14,7 @@ DEFINE_string(device, "cpu", "Device name, e.g. \"cpu\", \"cuda\"");
 DEFINE_string(output, "", "Output video path or format string");
 
 DEFINE_int32(output_size, 0, "Long-edge of output frames");
+DEFINE_int32(flip, 0, "Set to 1 for flipping the input horizontally");
 DEFINE_int32(show, 1, "Delay passed to `cv::waitKey` when using `cv::imshow`; -1: disable");
 
 DEFINE_string(skeleton, "coco", R"(Path to skeleton data or name of predefined skeletons: "coco")");
@@ -38,7 +39,7 @@ int main(int argc, char* argv[]) {
   // create a tracker state for each video
   mmdeploy::PoseTracker::State state = tracker.CreateState(params);
 
-  utils::mediaio::Input input(ARGS_input);
+  utils::mediaio::Input input(ARGS_input, FLAGS_flip);
   utils::mediaio::Output output(FLAGS_output, FLAGS_show);
 
   utils::Visualize v(FLAGS_output_size);
@@ -59,7 +60,7 @@ int main(int argc, char* argv[]) {
 
     // write to output stream
     if (!output.write(sess.get())) {
-      // user requested exit by pressing 'q'
+      // user requested exit by pressing ESC
       break;
     }
   }
diff --git a/demo/csrc/cpp/utils/mediaio.h b/demo/csrc/cpp/utils/mediaio.h
index 3426dfb3e0..65018602c2 100644
--- a/demo/csrc/cpp/utils/mediaio.h
+++ b/demo/csrc/cpp/utils/mediaio.h
@@ -154,8 +154,8 @@ class BatchInputIterator {
 
 class Input {
  public:
-  explicit Input(const std::string& path, MediaType type = MediaType::kUnknown)
-      : path_(path), type_(type) {
+  explicit Input(const std::string& path, bool flip = false, MediaType type = MediaType::kUnknown)
+      : path_(path), flip_(flip), type_(type) {
     if (type_ == MediaType::kUnknown) {
       auto ext = detail::get_extension(path);
       if (detail::is_image(ext)) {
@@ -212,6 +212,9 @@ class Input {
         }
       }
     }
+    if (flip_ && !img.empty()) {
+      cv::flip(img, img, 1);
+    }
     return img;
   }
 
@@ -270,8 +273,9 @@ class Input {
   }
 
  private:
-  MediaType type_{MediaType::kUnknown};
   std::string path_;
+  bool flip_{};
+  MediaType type_{MediaType::kUnknown};
   std::vector<std::string> items_;
   cv::VideoCapture cap_;
   size_t index_{};
@@ -350,7 +354,7 @@ class Output {
     }
     if (show_ >= 0) {
       cv::imshow("", frame);
-      exit = cv::waitKey(show_) == 'q';
+      exit = cv::waitKey(show_) == 27;  // ESC
     }
     ++frame_id_;
     return !exit;
@@ -385,4 +389,5 @@ OutputIterator& OutputIterator::operator=(const cv::Mat& frame) {
 
 }  // namespace mediaio
 }  // namespace utils
+
 #endif  // MMDEPLOY_MEDIAIO_H