Browse Source

Implemented versatile saver.

jcsyshc 1 year ago
parent
commit
9df9843230

+ 1 - 0
CMakeLists.txt

@@ -76,6 +76,7 @@ add_executable(${PROJECT_NAME} src/main.cpp
 add_subdirectory(src/core_v2)
 add_subdirectory(src/image_process_v5)
 add_subdirectory(src/render_osg)
+add_subdirectory(src/module_v5)
 
 target_include_directories(${PROJECT_NAME} PRIVATE src)
 

+ 12 - 0
src/core_v2/utility.hpp

@@ -4,6 +4,8 @@
 #include <cassert>
 #include <cstdint>
 
+#include <BS_thread_pool.hpp>
+
 template<size_t Align = 1>
 size_t alignment_round(size_t size, const size_t align = Align) {
     assert(std::popcount(align) == 1);
@@ -13,4 +15,14 @@ size_t alignment_round(size_t size, const size_t align = Align) {
     return size;
 }
 
+extern BS::thread_pool *g_thread_pool;
+
+template<typename Func>
+void thread_pool_detach_task(Func &&func) {
+    g_thread_pool->detach_task(std::forward<Func>(func));
+}
+
+#define TP_DETACH(func) thread_pool_detach_task(func)
+#define TP_SYNC g_thread_pool->wait()
+
 #endif //UTILITY_HPP

+ 21 - 0
src/image_process/cuda_impl/pixel_convert.cu

@@ -27,6 +27,27 @@ void call_cvt_rgb_bgra_u8(image_type_v2<uchar3> in,
     func_type(in, out, stream);
 }
 
+template<typename PixType>
+struct cvt_rgb_bgr {
+    __device__ static constexpr PixType Op(PixType in) {
+        auto out = PixType();
+        out.z = in.x;
+        out.y = in.y;
+        out.x = in.z;
+        return out;
+    }
+};
+
+using cvt_rgb_bgr_u8 = cvt_rgb_bgr<uchar3>;
+
+void call_cvt_rgb_bgr_u8(image_type_v2<uchar3> in,
+                         image_type_v2<uchar3> out,
+                         cudaStream_t stream) {
+    auto func_type = call_image_element_wise_unary<
+        uchar3, uchar3, cvt_rgb_bgr_u8>;
+    func_type(in, out, stream);
+}
+
 __device__ constexpr glm::vec3 to_vec3(uchar3 vec) {
     return glm::vec3(vec.x, vec.y, vec.z) / 255.f;
 }

+ 4 - 0
src/image_process/cuda_impl/pixel_convert.cuh

@@ -7,6 +7,10 @@ void call_cvt_rgb_bgra_u8(image_type_v2<uchar3> in,
                           image_type_v2<uchar4> out,
                           cudaStream_t stream);
 
+void call_cvt_rgb_bgr_u8(image_type_v2<uchar3> in,
+                         image_type_v2<uchar3> out,
+                         cudaStream_t stream);
+
 void call_yuv_to_rgb(image_type_v2<uchar3> in,
                      image_type_v2<uchar3> out,
                      cudaStream_t stream);

+ 29 - 0
src/image_process_v5/image_process.cpp

@@ -3,6 +3,7 @@
 #include <opencv2/cudaarithm.hpp>
 #include <opencv2/cudaimgproc.hpp>
 #include <opencv2/cudawarping.hpp>
+#include <opencv2/imgcodecs.hpp>
 
 namespace {
     // TODO: hack OpenCV code to make it support construction from cudaStream_t
@@ -115,6 +116,15 @@ namespace {
     };
 }
 
+sp_image image_rgb_to_bgr(const sp_image &img) {
+    assert(img.cv_type() == CV_8UC3);
+    auto ret = sp_image::create(CV_8UC3, img.cv_size());
+    auto helper = image_cuda_v2_helper<uchar3, uchar3>(img, ret);
+    call_cvt_rgb_bgr_u8(helper.input(), helper.output(), current_cuda_stream());
+    ret.merge_meta(img);
+    return ret;
+}
+
 sp_image image_rgb_to_bgra(const sp_image &img) {
     assert(img.cv_type() == CV_8UC3);
     auto ret = sp_image::create(CV_8UC4, img.cv_size());
@@ -151,6 +161,25 @@ sp_image image_yuyv_to_rgb(const sp_image &img) {
     return ret;
 }
 
+namespace {
+    void image_save_opencv(sp_image img, const std::string &filename) {
+        if (CV_MAT_CN(img.cv_type()) == 3) {
+            img = image_rgb_to_bgr(img);
+        }
+        const auto helper = read_access_helper(img.host());
+        const auto img_mat = img.cv_mat(helper.ptr());
+        cv::imwrite(filename, img_mat);
+    }
+}
+
+void image_save_jpg(const sp_image &img, const std::string &filename) {
+    image_save_opencv(img, fmt::format("{}.jpg", filename));
+}
+
+void image_save_png(const sp_image &img, const std::string &filename) {
+    image_save_opencv(img, fmt::format("{}.png", filename));
+}
+
 #include "render/render_utility.h"
 
 struct image_output_helper::impl {

+ 4 - 0
src/image_process_v5/image_process.h

@@ -15,11 +15,15 @@ sp_image image_debayer(const sp_image &img); // TODO: add an option for bayer ty
 void image_resize(const sp_image &src, sp_image &dst);
 sp_image image_resize(const sp_image &img, cv::Size size);
 sp_image image_flip_y(const sp_image &img);
+sp_image image_rgb_to_bgr(const sp_image &img);
 sp_image image_rgb_to_bgra(const sp_image &img);
 sp_image image_rgb_to_nv12(const sp_image &img);
 sp_image image_nv12_to_rgb(const sp_image &img);
 sp_image image_yuyv_to_rgb(const sp_image &img);
 
+void image_save_jpg(const sp_image &img, const std::string &filename); // filename without extension
+void image_save_png(const sp_image &img, const std::string &filename);
+
 #include <core_v2/object_manager.h>
 
 class image_output_helper {

+ 13 - 1
src/impl/apps/remote_ar/remote_ar_v2.cpp

@@ -42,9 +42,9 @@ app_remote_ar_v2::app_remote_ar_v2(create_config _conf)
 
     if (true) {
         auto sub_conf = image_viewer_v2::create_config();
+        sub_conf.items.emplace_back(uvc_img_id, "Endoscope", true);
         sub_conf.items.emplace_back(left_img_id, "Left", true);
         sub_conf.items.emplace_back(right_img_id, "Right", true);
-        sub_conf.items.emplace_back(uvc_img_id, "Endoscope", true);
         bg_viewer.emplace(sub_conf);
     }
 
@@ -54,6 +54,14 @@ app_remote_ar_v2::app_remote_ar_v2(create_config _conf)
         sub_conf.asio_ctx = main_conf.asio_ctx;
         streamer.emplace(sub_conf);
     }
+
+    if (true) {
+        auto sub_conf = versatile_saver::create_config();
+        // sub_conf.items.emplace_back(left_img_id, "Left", true);
+        // sub_conf.items.emplace_back(right_img_id, "Right", true);
+        sub_conf.items.emplace_back(uvc_img_id, "Endoscope", true);
+        saver.emplace(sub_conf);
+    }
 }
 
 app_remote_ar_v2::~app_remote_ar_v2() = default;
@@ -86,6 +94,10 @@ void app_remote_ar_v2::show_ui() {
                 bg_viewer->show_ui();
                 ImGui::TreePop();
             }
+            if (ImGui::TreeNode("Saver")) {
+                saver->show_ui();
+                ImGui::TreePop();
+            }
         }
 
         ImGui::PopItemWidth();

+ 2 - 0
src/impl/apps/remote_ar/remote_ar_v2.h

@@ -7,6 +7,7 @@
 #include "image_process_v5/image_viewer.h"
 #include "image_process_v5/image_process.h"
 #include "module/image_streamer.h"
+#include "module_v5/versatile_saver.h"
 
 class app_remote_ar_v2 final : public app_base {
 public:
@@ -35,6 +36,7 @@ private:
     std::optional<stereo_output_helper> output_helper;
     std::optional<image_viewer_v2> bg_viewer;
     std::optional<image_streamer> streamer;
+    std::optional<versatile_saver> saver;
 };
 
 

+ 13 - 0
src/impl/main_impl.cpp

@@ -2,6 +2,7 @@
 #include "core/object_manager.h"
 #include "apps/app_selector/app_selector.h"
 #include "core_v2/memory_manager.h"
+#include "core_v2/utility.hpp"
 
 #include <boost/asio/io_context.hpp>
 #include <boost/asio/post.hpp>
@@ -25,6 +26,7 @@ using boost::system::error_code;
 GLFWwindow *window = nullptr;
 smart_cuda_stream *default_cuda_stream = nullptr;
 io_context *main_ctx;
+BS::thread_pool *g_thread_pool;
 
 using cleanup_list_type =
 std::vector<cleanup_func_type>;
@@ -140,6 +142,9 @@ void init_all() {
     main_ob = new object_manager_v2({.ctx = main_ctx});
     g_memory_manager = new memory_manager();
 
+    constexpr auto background_thread_count = 6; // TODO: load this in config file
+    g_thread_pool = new BS::thread_pool(background_thread_count);
+
     auto app_conf = app_selector::create_config();
     app_conf.asio_ctx = main_ctx;
     app_conf.cuda_ctx = &cuda_ctx;
@@ -241,6 +246,11 @@ void show_memory_usage() {
                 status.cuda_allocated * kb_to_mb, status.cuda_cached * kb_to_mb);
 }
 
+void show_thread_pool_usage() {
+    ImGui::Text("Running: %ld", g_thread_pool->get_tasks_running());
+    ImGui::Text("Queued: %ld", g_thread_pool->get_tasks_queued());
+}
+
 void show_debug_ui() {
     if (ImGui::Begin("Debug")) {
         ImGui::SeparatorText("Display Config");
@@ -249,6 +259,9 @@ void show_debug_ui() {
         ImGui::SeparatorText("Memory Usage");
         show_memory_usage();
 
+        ImGui::SeparatorText("Thread Pool Usage");
+        show_thread_pool_usage();
+
         ImGui::SeparatorText("Miscellaneous");
         ImGui::Checkbox("Show Demo", &show_demo);
     }

+ 1 - 1
src/module_v5/CMakeLists.txt

@@ -1,2 +1,2 @@
 target_sources(${PROJECT_NAME} PRIVATE
-    )
+    versatile_saver.cpp)

+ 182 - 0
src/module_v5/versatile_saver.cpp

@@ -0,0 +1,182 @@
+#include "versatile_saver.h"
+#include "core/imgui_utility.hpp"
+#include "image_process_v5/sp_image.h"
+#include "image_process_v5/image_process.h"
+
+#include <ImGuiFileDialog.h>
+
+#include <boost/asio/post.hpp>
+
+using obj_type = object_manager_v2::obj_query_config;
+using boost::asio::post;
+
+extern boost::asio::io_context *main_ctx;
+
+namespace {
+    bool can_save_jpg(const obj_type &obj) {
+        return true; // TODO
+    }
+
+    bool can_save_png(const obj_type &obj) {
+        return true; // TODO
+    }
+}
+
+struct versatile_saver::impl {
+    enum save_type_enum: int {
+        JPG, PNG, // for sp_image
+        UNKNOWN
+    };
+
+    create_config conf;
+
+    struct item_info_type : create_config::item_info {
+        save_type_enum save_type = UNKNOWN;
+        obj_conn_type conn;
+        timestamp_type last_update_ts = 0;
+        bool is_selected = true;
+    };
+
+    using item_list_type = std::vector<item_info_type>;
+    item_list_type items;
+
+    size_t save_cnt = 0;
+    bool enable_autosave = false;
+    timestamp_type last_save_ts = 0;
+
+    constexpr static auto filepath_max_length = 256;
+    char filepath_buf[filepath_max_length] = {};
+    int min_interval_ms = 0;
+
+    void save_jpg(item_info_type &item) {
+        image_save_jpg(OBJ_QUERY(sp_image, item.name), get_save_filename(item));
+    }
+
+    void save_png(item_info_type &item) {
+        image_save_png(OBJ_QUERY(sp_image, item.name), get_save_filename(item));
+    }
+
+    void save_item(item_info_type &item) {
+        switch (item.save_type) {
+            //@formatter:off
+            case JPG: { save_jpg(item); break; }
+            case PNG: { save_png(item); break; }
+            default: { assert(false); }
+            //@formatter:on
+        }
+    }
+
+    void save_all() {
+        for (auto &item: items) {
+            if (!item.is_selected) continue;
+            TP_DETACH([&] { save_item(item); });
+        }
+        last_save_ts = current_timestamp();
+        ++save_cnt;
+    }
+
+    item_info_type &query_item(obj_name_type name) {
+        const auto iter = std::ranges::find_if(
+            items, [=](const auto &item) { return item.name == name; });
+        assert(iter != items.end());
+        return *iter;
+    }
+
+    bool should_auto_save() {
+        assert(enable_autosave);
+        return std::ranges::all_of(items, [this](const auto &item) {
+            if (!item.is_selected) return true;
+            return item.last_update_ts - last_save_ts > min_interval_ms * 1e3;
+        });
+    }
+
+    void object_callback(const obj_name_type name) {
+        auto &item = query_item(name);
+        item.last_update_ts = OBJ_TS(name);
+        if (enable_autosave && should_auto_save()) {
+            save_all();
+        }
+    }
+
+    std::string get_save_filename(item_info_type &item) {
+        const auto filename = fmt::format("{}_{}", item.save_name, save_cnt);
+        const auto save_folder = std::filesystem::path(filepath_buf);
+        if (!exists(save_folder)) {
+            create_directories(save_folder);
+        }
+        return save_folder / filename;
+    }
+
+    void show_ui() {
+        ImGui::Checkbox("Autosave", &enable_autosave);
+        ImGui::SameLine();
+        if (!enable_autosave) {
+            if (ImGui::Button("Save")) {
+                post(*main_ctx, [this] { save_all(); });
+            }
+        }
+
+        if (true) {
+            auto guard = imgui_disable_guard(enable_autosave);
+            ImGui::InputText("Save Folder", filepath_buf, filepath_max_length);
+            ImGui::InputInt("Autosave Interval (ms)", &min_interval_ms);
+        }
+
+        for (auto &item: items) {
+            if (ImGui::TreeNode(item.save_name.c_str())) {
+                ImGui::Checkbox("##select", &item.is_selected);
+
+                // query object info
+                auto item_info = obj_type();
+                item_info.name = item.name;
+                main_ob->query_all(item_info);
+
+                // save type
+                ImGui::SameLine();
+                if (can_save_jpg(item_info)) {
+                    ImGui::SameLine();
+                    ImGui::RadioButton("JPG", (int *) &item.save_type, JPG);
+                }
+                if (can_save_png(item_info)) {
+                    ImGui::SameLine();
+                    ImGui::RadioButton("PNG", (int *) &item.save_type, PNG);
+                }
+                ImGui::TreePop();
+            }
+        }
+    }
+
+    explicit impl(create_config _conf)
+        : conf(std::move(_conf)) {
+        std::ranges::transform(
+            conf.items, std::back_inserter(items),
+            [this](const create_config::item_info &item) {
+                auto ret = item_info_type();
+                *(create_config::item_info *) &ret = item;
+                ret.conn = OBJ_SIG(ret.name)->connect(
+                    [this](auto name) { object_callback(name); });
+                ret.save_type = JPG; // TODO: check for item type
+                return ret;
+            });
+
+        strcpy(filepath_buf, "./capture");
+        min_interval_ms = conf.min_interval_ms;
+    }
+
+    ~impl() {
+        for (auto &item: items) {
+            item.conn.disconnect();
+        }
+        TP_SYNC;
+    }
+};
+
+versatile_saver::versatile_saver(const create_config &conf)
+    : pimpl(std::make_unique<impl>(conf)) {
+}
+
+versatile_saver::~versatile_saver() = default;
+
+void versatile_saver::show_ui() const {
+    pimpl->show_ui();
+}

+ 34 - 0
src/module_v5/versatile_saver.h

@@ -0,0 +1,34 @@
+#ifndef VERSTAILE_SAVER_H
+#define VERSTAILE_SAVER_H
+
+#include "core_v2/object_manager.h"
+
+#include <filesystem>
+
+class versatile_saver {
+public:
+    struct create_config {
+        struct item_info {
+            obj_name_type name = {};
+            std::string save_name;
+            bool image_flip = false; // flip if first row comes first
+        };
+
+        std::vector<item_info> items;
+        uint64_t min_interval_ms = 0;
+    };
+
+    explicit versatile_saver(const create_config &conf);
+
+    ~versatile_saver();
+
+    void show_ui() const;
+
+private:
+    struct impl;
+    std::unique_ptr<impl> pimpl;
+};
+
+#undef image_viewer
+
+#endif //VERSTAILE_SAVER_H