Explorar o código

Implemented raw image saving.

jcsyshc %!s(int64=2) %!d(string=hai) anos
pai
achega
8ecfc68df4
Modificáronse 9 ficheiros con 341 adicións e 18 borrados
  1. 1 0
      CMakeLists.txt
  2. 3 0
      src/config.h
  3. 73 10
      src/main.cpp
  4. 46 7
      src/mvs_camera.cpp
  5. 7 0
      src/mvs_camera.h
  6. 155 0
      src/raw_file_saver.cpp
  7. 42 0
      src/raw_file_saver.h
  8. 13 0
      src/stereo_camera.hpp
  9. 1 1
      src/third_party/rs.c

+ 1 - 0
CMakeLists.txt

@@ -7,6 +7,7 @@ add_executable(RemoteAR2 src/main.cpp
         src/augment_renderer.cpp
         src/frame_sender.cpp
         src/sophiar_connect.cpp
+        src/raw_file_saver.cpp
         src/third_party/rs.c)
 
 # OpenGL config

+ 3 - 0
src/config.h

@@ -35,6 +35,9 @@ static constexpr auto output_frame_height = 1080;
 
 static constexpr auto ar_width_normal = 1.0 * output_frame_height * image_width / output_frame_width / image_height;
 
+static constexpr std::string_view raw_save_prefix = "raw_save";
+static constexpr auto default_raw_save_interval_s = 10; // 10s
+
 #define RET_ERROR \
     assert(false); \
     return false; \

+ 73 - 10
src/main.cpp

@@ -2,6 +2,7 @@
 #include "config.h"
 #include "frame_buffer_helper.hpp"
 #include "frame_sender.h"
+#include "raw_file_saver.h"
 #include "scene_manager.hpp"
 #include "sophiar_connect.h"
 #include "stereo_camera.hpp"
@@ -95,12 +96,16 @@ int main() {
     frame_sender sender;
     std::atomic_flag idr_flag;
 
-    frame_buffer_helper output_fbo;
-    output_fbo.initialize(output_frame_width, output_frame_height);
+    std::unique_ptr<frame_buffer_helper> output_fbo;
+//    output_fbo.initialize(output_frame_width, output_frame_height);
 
     video_encoder encoder;
     encoder.initialize();
 
+    raw_file_saver raw_saver;
+    raw_saver.start(raw_save_prefix);
+    camera.set_raw_saver(&raw_saver);
+
     // config sophiar
     using namespace sophiar;
     variable_io var_io;
@@ -154,7 +159,12 @@ int main() {
 
     int preview_cam = 0; // left camera
     bool save_video = true; // whether encoded bitstream should be saved
-    bool save_frame_length = true; // save frame length for tiny player to replay
+    bool save_frame_length = false; // save frame length for tiny player to replay
+    bool full_resolution = false;
+
+    bool auto_save_raw = false;
+    int auto_save_raw_interval = default_raw_save_interval_s;
+    std::chrono::system_clock::time_point last_save_raw_time;
 
     // main loop
     while (!glfwWindowShouldClose(main_window)) {
@@ -199,6 +209,12 @@ int main() {
                         if (ImGui::Button("Stop")) {
                             camera.stop_capture();
                         }
+                        if (!auto_save_raw) {
+                            ImGui::SameLine();
+                            if (ImGui::Button("Capture")) {
+                                camera.request_save_raw();
+                            }
+                        }
                     }
                 }
 
@@ -225,6 +241,40 @@ int main() {
                         ImGui::RadioButton("Left", &preview_cam, 0);
                         ImGui::SameLine();
                         ImGui::RadioButton("Right", &preview_cam, 1);
+
+                        // auto save raw config
+                        ImGui::SeparatorText("Auto Shoot");
+                        ImGui::PushID("Auto Shoot");
+
+                        if (!auto_save_raw) {
+                            if (ImGui::Button("Start")) {
+                                auto_save_raw = true;
+                            }
+                        } else {
+                            if (ImGui::Button("Stop")) {
+                                auto_save_raw = false;
+                            }
+                        }
+
+                        if (auto_save_raw_interval < 1) {
+                            auto_save_raw_interval = 1;
+                        }
+
+                        if (auto_save_raw) {
+                            ImGui::BeginDisabled();
+                            auto now_time = std::chrono::system_clock::now();
+                            if (now_time - last_save_raw_time >
+                                std::chrono::seconds{auto_save_raw_interval}) {
+                                camera.request_save_raw();
+                                last_save_raw_time = now_time;
+                            }
+                        }
+                        ImGui::InputInt("Shoot Interval (s)", &auto_save_raw_interval, 1, 10);
+                        if (auto_save_raw) {
+                            ImGui::EndDisabled();
+                        }
+
+                        ImGui::PopID();
                     }
                 }
 
@@ -238,7 +288,13 @@ int main() {
                 ImGui::SeparatorText("Actions");
                 if (!encoder.is_encoding()) {
                     if (ImGui::Button("Start")) {
-                        encoder.start_encode(output_fbo.tex_width, output_fbo.tex_height,
+                        output_fbo = std::make_unique<frame_buffer_helper>();
+                        if (full_resolution) {
+                            output_fbo->initialize(2 * image_width, image_height);
+                        } else {
+                            output_fbo->initialize(output_frame_width, output_frame_height);
+                        }
+                        encoder.start_encode(output_fbo->tex_width, output_fbo->tex_height,
                                              camera_fps, (int) (output_bitrate_mbps * 1e6),
                                              save_video, save_frame_length);
                     }
@@ -254,6 +310,8 @@ int main() {
                 }
 
                 ImGui::DragFloat("Bitrate (Mbps)", &output_bitrate_mbps, 0.1, 1, 20, "%.01f");
+                ImGui::Checkbox("Full Resolution", &full_resolution);
+                ImGui::SameLine();
                 ImGui::Checkbox("Save Video", &save_video);
                 if (save_video) {
                     ImGui::SameLine();
@@ -334,10 +392,15 @@ int main() {
 
         if (encoder.is_encoding()) {
             // draw frame for streaming
-            glBindFramebuffer(GL_DRAW_FRAMEBUFFER, output_fbo.fbo);
-            glViewport(0, 0, output_fbo.tex_width, output_fbo.tex_height);
-            left_ar.render({-0.5 - ar_width_normal / 2, -1, ar_width_normal, 2});
-            right_ar.render({0.5 - ar_width_normal / 2, -1, ar_width_normal, 2});
+            glBindFramebuffer(GL_DRAW_FRAMEBUFFER, output_fbo->fbo);
+            glViewport(0, 0, output_fbo->tex_width, output_fbo->tex_height);
+            if (full_resolution) {
+                left_ar.render({-1, -1, 1, 2});
+                right_ar.render({0, -1, 1, 2});
+            } else {
+                left_ar.render({-0.5 - ar_width_normal / 2, -1, ar_width_normal, 2});
+                right_ar.render({0.5 - ar_width_normal / 2, -1, ar_width_normal, 2});
+            }
 
             // send IDR frame occasionally to prevent broken video frame
             static int last_p_count = 0;
@@ -346,7 +409,7 @@ int main() {
             }
 
             // encode frame
-            output_fbo.download_pixels();
+            output_fbo->download_pixels();
             frame_sender::frame_info info;
             info.is_idr_frame = idr_flag.test();
             if (info.is_idr_frame) {
@@ -355,7 +418,7 @@ int main() {
                 idr_flag.clear();
                 last_p_count = 0;
             }
-            encoder.encode_frame(output_fbo.pbo_res, (void **) &info.data, &info.length);
+            encoder.encode_frame(output_fbo->pbo_res, (void **) &info.data, &info.length);
             SPDLOG_TRACE("Time used: {}ms, length = {}", std::chrono::duration_cast<std::chrono::milliseconds>(
                     std::chrono::high_resolution_clock::now() - start_time).count(), info.length);
 

+ 46 - 7
src/mvs_camera.cpp

@@ -9,6 +9,21 @@
 
 #include <spdlog/spdlog.h>
 
+#include <boost/predef.h>
+
+#ifdef BOOST_OS_WINDOWS_AVAILABLE
+
+#include <format>
+
+#define fmt std
+
+#else
+
+#include <fmt/chrono.h>
+#include <fmt/format.h>
+
+#endif
+
 #include <atomic>
 
 bool check_mvs_api_call(int api_ret, unsigned int line_number,
@@ -32,6 +47,9 @@ struct mvs_camera::impl {
     cv::cuda::GpuMat *inner_img = nullptr;
     std::atomic<cv::cuda::GpuMat *> next_img;
 
+    raw_file_saver *raw_saver;
+    std::atomic_flag save_raw_flag;
+
     static void on_error(unsigned int msg_type, void *user_data) {
         auto pimpl = (impl *) user_data;
         SPDLOG_ERROR("MVS camera {} exception 0x{:x}.", pimpl->cam_name, msg_type);
@@ -39,20 +57,33 @@ struct mvs_camera::impl {
         assert(false);
     }
 
-    static void on_image(unsigned char *data, MV_FRAME_OUT_INFO_EX *frame_info, void *user_data) {
+    void on_image_impl(unsigned char *data, MV_FRAME_OUT_INFO_EX *frame_info) {
         assert(frame_info->nFrameLen == raw_image_size);
-        auto pimpl = (impl *) user_data;
         auto host_img = cv::Mat{image_height, image_width, CV_8UC1, data};
 
         // upload image to gpu
-        if (pimpl->inner_img == nullptr) [[unlikely]] {
-            pimpl->inner_img = new cv::cuda::GpuMat{};
+        if (inner_img == nullptr) [[unlikely]] {
+            inner_img = new cv::cuda::GpuMat{};
         }
-        pimpl->inner_img->upload(host_img);
+        inner_img->upload(host_img);
 
         // commit new image
-        pimpl->inner_img = pimpl->next_img.exchange(pimpl->inner_img);
-        pimpl->next_img.notify_all();
+        inner_img = next_img.exchange(inner_img);
+        next_img.notify_all();
+
+        if (raw_saver != nullptr && save_raw_flag.test()) {
+            auto file_name = fmt::format("{}_{}.bmp", cam_name, frame_info->nFrameNum);
+            auto data_len = frame_info->nFrameLen;
+            auto data_bak = malloc(data_len);
+            memcpy(data_bak, data, data_len);
+//            raw_saver->save_file({data_bak, data_len, file_name});
+            raw_saver->save_image({data_bak, image_height, image_width, CV_8UC1, file_name});
+            save_raw_flag.clear();
+        }
+    }
+
+    static void on_image(unsigned char *data, MV_FRAME_OUT_INFO_EX *frame_info, void *user_data) {
+        ((impl *) user_data)->on_image_impl(data, frame_info);
     }
 };
 
@@ -149,4 +180,12 @@ bool mvs_camera::is_opened() const {
 
 bool mvs_camera::is_capturing() const {
     return pimpl->is_capturing;
+}
+
+void mvs_camera::set_raw_saver(raw_file_saver *saver) {
+    pimpl->raw_saver = saver;
+}
+
+std::atomic_flag *mvs_camera::get_save_raw_flag() const {
+    return &pimpl->save_raw_flag;
 }

+ 7 - 0
src/mvs_camera.h

@@ -1,8 +1,11 @@
 #ifndef REMOTEAR2_MVS_CAMERA_H
 #define REMOTEAR2_MVS_CAMERA_H
 
+#include "raw_file_saver.h"
+
 #include <opencv2/core/cuda.hpp>
 
+#include <atomic>
 #include <memory>
 #include <string_view>
 
@@ -34,6 +37,10 @@ public:
 
     bool is_capturing() const;
 
+    void set_raw_saver(raw_file_saver *saver);
+
+    std::atomic_flag *get_save_raw_flag() const;
+
 private:
     struct impl;
     std::unique_ptr<impl> pimpl;

+ 155 - 0
src/raw_file_saver.cpp

@@ -0,0 +1,155 @@
+#include "raw_file_saver.h"
+
+#include <boost/asio/awaitable.hpp>
+#include <boost/asio/co_spawn.hpp>
+#include <boost/asio/detached.hpp>
+#include <boost/asio/experimental/concurrent_channel.hpp>
+#include <boost/asio/io_context.hpp>
+#include <boost/asio/use_awaitable.hpp>
+#include <boost/smart_ptr.hpp>
+
+#include <opencv2/imgcodecs.hpp>
+
+#include <spdlog/spdlog.h>
+
+#include <filesystem>
+#include <thread>
+
+using boost::asio::awaitable;
+using boost::asio::co_spawn;
+using boost::asio::detached;
+using boost::asio::experimental::concurrent_channel;
+using boost::asio::io_context;
+using boost::asio::use_awaitable;
+using boost::system::error_code;
+
+struct raw_file_saver::impl {
+
+    static constexpr auto channel_capacity = 16;
+
+    boost::scoped_ptr<io_context> context;
+
+    using file_chan_type = concurrent_channel<void(error_code, file_info)>;
+    using image_chan_type = concurrent_channel<void(error_code, image_info)>;
+    boost::scoped_ptr<file_chan_type> file_chan;
+    boost::scoped_ptr<image_chan_type> image_chan;
+
+    std::unique_ptr<std::thread> work_thread;
+
+    std::filesystem::path prefix_path;
+
+    impl() {
+        context.reset(new io_context{});
+        file_chan.reset(new file_chan_type{*context, channel_capacity});
+        image_chan.reset(new image_chan_type{*context, channel_capacity});
+
+        // create main loop
+        auto error_handler = [](std::exception_ptr ep) {
+            if (!ep) {
+                SPDLOG_ERROR("Infinite loop exited with no error.");
+                return;
+            }
+            try {
+                std::rethrow_exception(ep);
+            } catch (std::exception &e) {
+                SPDLOG_ERROR("Infinite loop exited with error: {}", e.what());
+            }
+        };
+        co_spawn(*context, main_loop_file(), error_handler);
+        co_spawn(*context, main_loop_image(), error_handler);
+    }
+
+    ~impl() {
+        stop();
+    }
+
+    void handle_file(const file_info &info) {
+        static constexpr auto block_size = 512 * 1024; // 512KiB
+        auto save_file_path = prefix_path / info.file_name;
+        auto file = fopen(save_file_path.c_str(), "wb");
+        assert(file != nullptr);
+        auto block_cnt = info.length / block_size;
+        auto write_cnt = fwrite(info.data, block_size, block_cnt, file);
+        assert(write_cnt == block_cnt);
+        auto write_len = write_cnt * block_size;
+        if (info.length > write_len) {
+            fwrite((char *) info.data + write_len, info.length - write_len, 1, file);
+        }
+        fclose(file);
+        free(info.data);
+    }
+
+    void handle_image(const image_info &info) {
+        auto img = cv::Mat{info.height, info.width, info.cv_type, info.data};
+        auto save_file_path = prefix_path / info.file_name;
+        cv::imwrite(save_file_path, img);
+        free(info.data);
+    }
+
+    awaitable<void> main_loop_file() {
+        for (;;) {
+            auto info = co_await file_chan->async_receive(use_awaitable);
+            handle_file(info);
+        }
+    }
+
+    awaitable<void> main_loop_image() {
+        for (;;) {
+            auto info = co_await image_chan->async_receive(use_awaitable);
+            handle_image(info);
+        }
+    }
+
+    void start() {
+        // clean channel
+        if (file_chan != nullptr) {
+            while (file_chan->try_receive([](error_code e, file_info &&info) {
+                free(info.data);
+            }));
+        }
+
+        // restart context if needed
+        if (context->stopped()) {
+            context->restart();
+        }
+
+        // create working thread
+        assert(work_thread == nullptr);
+        work_thread = std::make_unique<std::thread>([this]() {
+            context->run();
+        });
+    }
+
+    void stop() {
+        if (work_thread == nullptr) return;
+        context->stop();
+        work_thread->join();
+        work_thread = nullptr;
+    }
+
+};
+
+raw_file_saver::raw_file_saver()
+        : pimpl(std::make_unique<impl>()) {
+}
+
+raw_file_saver::~raw_file_saver() = default;
+
+void raw_file_saver::start(std::string_view prefix) {
+    pimpl->prefix_path = prefix;
+    pimpl->start();
+}
+
+void raw_file_saver::stop() {
+    pimpl->stop();
+}
+
+void raw_file_saver::save_file(const file_info &info) {
+    auto ok = pimpl->file_chan->try_send(error_code{}, info);
+    assert(ok);
+}
+
+void raw_file_saver::save_image(const raw_file_saver::image_info &info) {
+    auto ok = pimpl->image_chan->try_send(error_code{}, info);
+    assert(ok);
+}

+ 42 - 0
src/raw_file_saver.h

@@ -0,0 +1,42 @@
+#ifndef REMOTEAR2_RAW_FILE_SAVER_H
+#define REMOTEAR2_RAW_FILE_SAVER_H
+
+#include <memory>
+#include <string>
+#include <string_view>
+
+class raw_file_saver {
+public:
+
+    raw_file_saver();
+
+    ~raw_file_saver();
+
+    void start(std::string_view prefix);
+
+    void stop();
+
+    struct file_info {
+        void *data;
+        size_t length;
+        std::string file_name;
+    };
+
+    void save_file(const file_info &info);
+
+    struct image_info {
+        void *data;
+        int height, width, cv_type;
+        std::string file_name;
+    };
+
+    void save_image(const image_info &info);
+
+private:
+    struct impl;
+    std::unique_ptr<impl> pimpl;
+
+};
+
+
+#endif //REMOTEAR2_RAW_FILE_SAVER_H

+ 13 - 0
src/stereo_camera.hpp

@@ -14,10 +14,13 @@ struct stereo_camera {
     mvs_camera left_camera, right_camera;
     cv::cuda::GpuMat *left_raw_image = nullptr, *right_raw_image = nullptr;
     cv::cuda::GpuMat left_rgb_image, right_rgb_image;
+    std::atomic_flag *left_save_raw_flag, *right_save_raw_flag;
 
     stereo_camera() {
         left_rgb_image = cv::cuda::GpuMat{image_height, image_width, CV_8UC3};
         right_rgb_image = cv::cuda::GpuMat{image_height, image_width, CV_8UC3};
+        left_save_raw_flag = left_camera.get_save_raw_flag();
+        right_save_raw_flag = right_camera.get_save_raw_flag();
     }
 
     ~stereo_camera() {
@@ -101,6 +104,16 @@ struct stereo_camera {
         cv::cuda::cvtColor(*right_raw_image, right_rgb_image, cv::COLOR_BayerRG2RGB);
     }
 
+    void set_raw_saver(raw_file_saver *saver) {
+        left_camera.set_raw_saver(saver);
+        right_camera.set_raw_saver(saver);
+    }
+
+    void request_save_raw() {
+        left_save_raw_flag->test_and_set();
+        right_save_raw_flag->test_and_set();
+    }
+
 private:
 
     std::thread *trigger_thread = nullptr;

+ 1 - 1
src/third_party/rs.c

@@ -169,7 +169,7 @@ modnn(int x)
  * A value related to the multiplication is held in a local variable
  * declared with USE_GF_MULC . See usage in addmul1().
  */
-__declspec(align(16)) static gf gf_mul_table[(GF_SIZE + 1)*(GF_SIZE + 1)];
+_Alignas(16) static gf gf_mul_table[(GF_SIZE + 1)*(GF_SIZE + 1)];
 
 #define gf_mul(x,y) gf_mul_table[(x<<8)+y]