[Fix] Add an option to flip webcam inputs for pose tracker demo (#1725)

(cherry picked from commit 5de0ecf)
open-mmlab · Feb 7, 2023 · 2731b61 · 2731b61
1 parent 682cb79
commit 2731b61
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 6 deletions.
diff --git 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
@@ -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