|
|
@@ -0,0 +1,367 @@
|
|
|
+#include "video_stabilization.h"
|
|
|
+#include "third_party/scope_guard.hpp"
|
|
|
+
|
|
|
+#include <opencv2/calib3d.hpp>
|
|
|
+#include <opencv2/cudaoptflow.hpp>
|
|
|
+#include <opencv2/cudaimgproc.hpp>
|
|
|
+
|
|
|
+#include <Eigen/Dense>
|
|
|
+
|
|
|
+#include <cmath>
|
|
|
+#include <deque>
|
|
|
+#include <ranges>
|
|
|
+#include <vector>
|
|
|
+
|
|
|
+namespace {
|
|
|
+ template<typename T>
|
|
|
+ std::vector<T> download(const cv::cuda::GpuMat &d_mat, int type) {
|
|
|
+ auto ret = std::vector<T>(d_mat.cols);
|
|
|
+ auto mat = cv::Mat(1, d_mat.cols, type, ret.data());
|
|
|
+ d_mat.download(mat);
|
|
|
+ return ret;
|
|
|
+ }
|
|
|
+
|
|
|
+ struct point_tracker {
|
|
|
+ cv::Ptr<cv::cuda::CornersDetector> detector;
|
|
|
+ cv::Ptr<cv::cuda::SparsePyrLKOpticalFlow> tracker;
|
|
|
+ image_ptr prev;
|
|
|
+
|
|
|
+ using points_type = std::vector<cv::Point2f>;
|
|
|
+
|
|
|
+ struct result_type {
|
|
|
+ points_type prev_pts;
|
|
|
+ points_type curr_pts;
|
|
|
+ };
|
|
|
+
|
|
|
+ using create_config = video_stabilization::point_tracker_config;
|
|
|
+ create_config conf;
|
|
|
+
|
|
|
+ explicit point_tracker(const create_config _conf)
|
|
|
+ : conf(_conf) {
|
|
|
+ detector = cv::cuda::createGoodFeaturesToTrackDetector(
|
|
|
+ CV_8UC1, conf.max_feature_points,
|
|
|
+ 0.3, 7, 7);
|
|
|
+ tracker = cv::cuda::SparsePyrLKOpticalFlow::create();
|
|
|
+ }
|
|
|
+
|
|
|
+ // Input RGB image
|
|
|
+ std::optional<result_type> process(const image_ptr &img) {
|
|
|
+ auto closer = sg::make_scope_guard([&] { prev = img; });
|
|
|
+ if (prev == nullptr) return {};
|
|
|
+
|
|
|
+ cv::cuda::GpuMat d_prev_pts;
|
|
|
+ if (true) {
|
|
|
+ const auto prev_gray = image_rgb_to_gray(prev);
|
|
|
+ const auto helper = read_access_helper(prev_gray.cuda());
|
|
|
+ detector->detect(prev_gray.cv_gpu_mat(helper.ptr()), d_prev_pts,
|
|
|
+ cv::noArray(), stream_guard.cv_stream());
|
|
|
+ }
|
|
|
+
|
|
|
+ cv::cuda::GpuMat d_curr_pts, d_status;
|
|
|
+ if (true) {
|
|
|
+ const auto prev_helper = read_access_helper(prev.cuda());
|
|
|
+ const auto curr_helper = read_access_helper(img.cuda());
|
|
|
+ tracker->calc(prev.cv_gpu_mat(prev_helper.ptr()),
|
|
|
+ img.cv_gpu_mat(curr_helper.ptr()),
|
|
|
+ d_prev_pts, d_curr_pts, d_status);
|
|
|
+ }
|
|
|
+
|
|
|
+ auto status = download<uchar>(d_status, CV_8UC1);
|
|
|
+ auto valid_indices = std::views::iota(0ull, status.size())
|
|
|
+ | std::views::filter([&](auto i) { return status[i]; });
|
|
|
+ auto prev_pts = download<cv::Point2f>(d_prev_pts, CV_32FC2);
|
|
|
+ auto curr_pts = download<cv::Point2f>(d_curr_pts, CV_32FC2);
|
|
|
+ auto valid_prev_pts = valid_indices
|
|
|
+ | std::views::transform([&](auto i) { return prev_pts[i]; });
|
|
|
+ auto valid_curr_pts = valid_indices
|
|
|
+ | std::views::transform([&](auto i) { return curr_pts[i]; });
|
|
|
+ auto ret = result_type();
|
|
|
+ std::ranges::copy(valid_prev_pts, std::back_inserter(ret.prev_pts));
|
|
|
+ std::ranges::copy(valid_curr_pts, std::back_inserter(ret.curr_pts));
|
|
|
+ return ret;
|
|
|
+ }
|
|
|
+ };
|
|
|
+
|
|
|
+ struct transform_stabilization {
|
|
|
+ using matrix_type = Eigen::Matrix<float, 2, 3>;
|
|
|
+ using matrix_param_type = Eigen::Vector4f;
|
|
|
+ using record_type = Eigen::VectorXf;
|
|
|
+ using trajectory_type = Eigen::Matrix4Xf;
|
|
|
+
|
|
|
+ static record_type optimize_simple(const record_type &traj) {
|
|
|
+ const auto num = traj.rows();
|
|
|
+ auto ret = record_type(num);
|
|
|
+ ret.setConstant(traj.mean());
|
|
|
+ return ret;
|
|
|
+ }
|
|
|
+
|
|
|
+ static trajectory_type optimize_trajectory(const trajectory_type &traj) {
|
|
|
+ const auto num_dims = traj.rows();
|
|
|
+ auto ret = trajectory_type(num_dims, traj.cols());
|
|
|
+ for (auto i = 0; i < num_dims; ++i) {
|
|
|
+ ret.row(i) = optimize_simple(traj.row(i));
|
|
|
+ }
|
|
|
+ return ret;
|
|
|
+ }
|
|
|
+
|
|
|
+ static matrix_param_type decompose_transform(const matrix_type &mat) {
|
|
|
+ const auto scale = sqrtf(mat.block<2, 2>(0, 0).determinant());
|
|
|
+ const auto angle = atan2f(mat(1, 0), mat(0, 0));
|
|
|
+ const auto tx = mat(0, 2);
|
|
|
+ const auto ty = mat(1, 2);
|
|
|
+ return {logf(scale), angle, tx, ty};
|
|
|
+ }
|
|
|
+
|
|
|
+ static matrix_type compose_transform(const matrix_param_type ¶m) {
|
|
|
+ const auto scale = expf(param(0)), angle = param(1);
|
|
|
+ const auto tx = param(2), ty = param(3);
|
|
|
+ auto ret = matrix_type();
|
|
|
+ ret << cosf(angle), -sinf(angle), tx,
|
|
|
+ sinf(angle), cosf(angle), ty;
|
|
|
+ ret.block<2, 2>(0, 0) *= scale;
|
|
|
+ return ret;
|
|
|
+ }
|
|
|
+
|
|
|
+ using create_config = video_stabilization::transform_stabilization_config;
|
|
|
+ create_config conf;
|
|
|
+
|
|
|
+ size_t buffer_length = 0;
|
|
|
+ std::deque<matrix_param_type> transform_buf;
|
|
|
+
|
|
|
+ explicit transform_stabilization(const create_config &_conf)
|
|
|
+ : conf(_conf) {
|
|
|
+ buffer_length = 2 * conf.future_frames + 1;
|
|
|
+ }
|
|
|
+
|
|
|
+ std::optional<matrix_type> process(const matrix_type &mat) {
|
|
|
+ const auto param = decompose_transform(mat);
|
|
|
+ transform_buf.push_back(param);
|
|
|
+ if (transform_buf.size() <= conf.future_frames) return {};
|
|
|
+ if (transform_buf.size() > buffer_length) {
|
|
|
+ transform_buf.pop_front();
|
|
|
+ }
|
|
|
+
|
|
|
+ auto curr_traj = trajectory_type(4, transform_buf.size());
|
|
|
+ for (int i = 0; i < transform_buf.size(); i++) {
|
|
|
+ curr_traj.col(i) = transform_buf[i];
|
|
|
+ if (i > 0) {
|
|
|
+ curr_traj.col(i) += curr_traj.col(i - 1);
|
|
|
+ }
|
|
|
+ }
|
|
|
+
|
|
|
+ const auto traj_opt = optimize_trajectory(curr_traj);
|
|
|
+ const auto curr_pos = conf.future_frames;
|
|
|
+ const auto diff = traj_opt.col(curr_pos) - curr_traj.col(curr_pos);
|
|
|
+ const auto curr = transform_buf[curr_pos] + diff;
|
|
|
+ return compose_transform(curr);
|
|
|
+ }
|
|
|
+ };
|
|
|
+}
|
|
|
+
|
|
|
+struct video_stabilization::impl {
|
|
|
+ using matrix_type = transform_stabilization::matrix_type;
|
|
|
+
|
|
|
+ static matrix_type combine_transform(const matrix_type &mat1, const matrix_type &mat2) {
|
|
|
+ auto ret = matrix_type();
|
|
|
+ const auto r1_mat = mat1.leftCols<2>();
|
|
|
+ ret.leftCols<2>() = r1_mat * mat2.leftCols<2>();
|
|
|
+ ret.rightCols<1>() = r1_mat * mat2.rightCols<1>() + mat1.rightCols<1>();
|
|
|
+ return ret;
|
|
|
+ }
|
|
|
+
|
|
|
+ static matrix_type identity_transform() {
|
|
|
+ auto ret = matrix_type();
|
|
|
+ ret.setZero();
|
|
|
+ ret.leftCols<2>().setIdentity();
|
|
|
+ return ret;
|
|
|
+ }
|
|
|
+
|
|
|
+ static glm::mat3 to_mat3(const matrix_type &mat) {
|
|
|
+ auto ret = glm::mat3();
|
|
|
+ for (auto i = 0; i < 2; ++i)
|
|
|
+ for (auto j = 0; j < 3; ++j) {
|
|
|
+ ret[j][i] = mat(i, j);
|
|
|
+ }
|
|
|
+ return ret;
|
|
|
+ }
|
|
|
+
|
|
|
+ static matrix_type to_transform(const cv::Mat &mat) {
|
|
|
+ assert(mat.rows == 2 && mat.cols == 3);
|
|
|
+ assert(mat.type() == CV_64FC1);
|
|
|
+ auto ret = matrix_type();
|
|
|
+ for (auto i = 0; i < 2; ++i)
|
|
|
+ for (auto j = 0; j < 3; ++j) {
|
|
|
+ ret(i, j) = mat.at<double>(i, j);
|
|
|
+ }
|
|
|
+ return ret;
|
|
|
+ }
|
|
|
+
|
|
|
+ static constexpr auto min_tracked_points = 10;
|
|
|
+
|
|
|
+ create_config conf;
|
|
|
+ std::optional<point_tracker> track_left, track_right;
|
|
|
+ std::optional<transform_stabilization> transform_stab;
|
|
|
+ std::deque<sp_image> img_buf;
|
|
|
+
|
|
|
+ [[nodiscard]] sp_image transform_image(const sp_image &img, matrix_type mat) const {
|
|
|
+ const auto center = cv::Point2f(img.width() / 2.f, img.height() / 2.f);
|
|
|
+ const auto fix_mat = cv::getRotationMatrix2D(center, 0, 1.f + conf.enlarge_factor);
|
|
|
+ mat = combine_transform(to_transform(fix_mat), mat);
|
|
|
+ return image_warp_affine(img, to_mat3(mat));
|
|
|
+ }
|
|
|
+
|
|
|
+ std::optional<sp_image> process_mono(const sp_image &img) {
|
|
|
+ assert(!conf.stereo_mode);
|
|
|
+ const auto pts = track_left->process(img);
|
|
|
+ if (!pts) return {};
|
|
|
+ matrix_type mat;
|
|
|
+ if (pts->prev_pts.size() < min_tracked_points) [[unlikely]] {
|
|
|
+ mat = identity_transform();
|
|
|
+ } else {
|
|
|
+ const auto cv_mat = cv::estimateAffinePartial2D(
|
|
|
+ pts->prev_pts, pts->curr_pts);
|
|
|
+ mat = to_transform(cv_mat);
|
|
|
+ }
|
|
|
+
|
|
|
+ const auto cur_mat = transform_stab->process(mat);
|
|
|
+ img_buf.push_back(img);
|
|
|
+ if (!cur_mat) return {};
|
|
|
+
|
|
|
+ const auto cur_img = img_buf.front();
|
|
|
+ img_buf.pop_front();
|
|
|
+ return transform_image(cur_img, *cur_mat);
|
|
|
+ }
|
|
|
+
|
|
|
+ std::optional<sp_image> process_stereo(const sp_image &img) {
|
|
|
+ assert(conf.stereo_mode);
|
|
|
+ auto [img_left, img_right] = image_stereo_split_view(img);
|
|
|
+ auto pts_left = track_left->process(img_left);
|
|
|
+ auto pts_right = track_right->process(img_right);
|
|
|
+ if (!pts_left || !pts_right) return {};
|
|
|
+
|
|
|
+ auto num_points = pts_left->prev_pts.size() + pts_right->prev_pts.size();
|
|
|
+ matrix_type mat;
|
|
|
+ if (num_points < min_tracked_points) [[unlikely]] {
|
|
|
+ mat = identity_transform();
|
|
|
+ } else {
|
|
|
+ auto prev_pts = std::move(pts_left->prev_pts);
|
|
|
+ prev_pts.reserve(num_points);
|
|
|
+ std::ranges::copy(pts_right->prev_pts, std::back_inserter(prev_pts));
|
|
|
+ auto curr_pts = std::move(pts_left->curr_pts);
|
|
|
+ curr_pts.reserve(num_points);
|
|
|
+ std::ranges::copy(pts_right->curr_pts, std::back_inserter(curr_pts));
|
|
|
+ const auto cv_mat = cv::estimateAffinePartial2D(prev_pts, curr_pts);
|
|
|
+ mat = to_transform(cv_mat);
|
|
|
+ }
|
|
|
+
|
|
|
+ const auto cur_mat = transform_stab->process(mat);
|
|
|
+ img_buf.push_back(img);
|
|
|
+ if (!cur_mat) return {};
|
|
|
+
|
|
|
+ const auto cur_img = img_buf.front();
|
|
|
+ img_buf.pop_front();
|
|
|
+ auto [cur_left, cur_right] = image_stereo_split_view(cur_img);
|
|
|
+ cur_left = transform_image(cur_left, *cur_mat);
|
|
|
+ cur_right = transform_image(cur_right, *cur_mat);
|
|
|
+ return image_stereo_combine(cur_left, cur_right);
|
|
|
+ }
|
|
|
+
|
|
|
+ auto process(const sp_image &img) {
|
|
|
+ if (conf.stereo_mode) {
|
|
|
+ return process_stereo(img);
|
|
|
+ } else {
|
|
|
+ return process_mono(img);
|
|
|
+ }
|
|
|
+ }
|
|
|
+
|
|
|
+ explicit impl(const create_config &_conf) : conf(_conf) {
|
|
|
+ track_left.emplace(conf.pt_conf);
|
|
|
+ if (conf.stereo_mode) {
|
|
|
+ track_right.emplace(conf.pt_conf);
|
|
|
+ }
|
|
|
+ transform_stab.emplace(conf.st_conf);
|
|
|
+ }
|
|
|
+};
|
|
|
+
|
|
|
+video_stabilization::video_stabilization(create_config conf)
|
|
|
+ : pimpl(std::make_unique<impl>(conf)) { (void) 0; }
|
|
|
+
|
|
|
+video_stabilization::~video_stabilization() = default;
|
|
|
+
|
|
|
+std::optional<sp_image> video_stabilization::process(const sp_image &img) const {
|
|
|
+ return pimpl->process(img);
|
|
|
+}
|
|
|
+
|
|
|
+#include "core/imgui_utility.hpp"
|
|
|
+
|
|
|
+struct video_stabilization_ui::impl {
|
|
|
+ create_config conf;
|
|
|
+ obj_conn_type conn;
|
|
|
+
|
|
|
+ bool passthrough = false;
|
|
|
+ std::optional<BS::thread_pool> work_tp;
|
|
|
+ std::optional<video_stabilization> video_stab;
|
|
|
+
|
|
|
+ void image_callback_impl() const {
|
|
|
+ const auto img = OBJ_QUERY(sp_image, conf.in_name);
|
|
|
+ if (const auto ret = video_stab->process(img)) {
|
|
|
+ OBJ_SAVE(conf.out_name, *ret);
|
|
|
+ }
|
|
|
+ }
|
|
|
+
|
|
|
+ void image_callback() {
|
|
|
+ if (passthrough) {
|
|
|
+ OBJ_SAVE(conf.out_name, OBJ_QUERY(sp_image, conf.in_name));
|
|
|
+ return;
|
|
|
+ }
|
|
|
+
|
|
|
+ if (!video_stab) [[unlikely]] {
|
|
|
+ video_stab.emplace(conf.opts);
|
|
|
+ }
|
|
|
+ auto task = [this] {
|
|
|
+ try {
|
|
|
+ image_callback_impl();
|
|
|
+ } catch (...) { (void) 0; }
|
|
|
+ };
|
|
|
+ if (work_tp->get_tasks_queued() <= 1) {
|
|
|
+ work_tp->detach_task(task);
|
|
|
+ } else {
|
|
|
+ SPDLOG_WARN("Too many frames for stabilization, frame dropped.");
|
|
|
+ }
|
|
|
+ }
|
|
|
+
|
|
|
+ void show_ui() {
|
|
|
+ if (ImGui::Checkbox("Passthrough", &passthrough)) {
|
|
|
+ if (!passthrough) {
|
|
|
+ video_stab.reset();
|
|
|
+ }
|
|
|
+ }
|
|
|
+ if (!passthrough) {
|
|
|
+ ImGui::SameLine();
|
|
|
+ if (ImGui::Button("Reset")) {
|
|
|
+ video_stab.reset();
|
|
|
+ }
|
|
|
+ ImGui::DragInt("Frame Delay",
|
|
|
+ &conf.opts.st_conf.future_frames, 1, 1);
|
|
|
+ }
|
|
|
+ }
|
|
|
+
|
|
|
+ explicit impl(const create_config &_conf) : conf(_conf) {
|
|
|
+ work_tp.emplace(1);
|
|
|
+ conn = OBJ_SIG(conf.in_name)->connect(
|
|
|
+ [this](auto _) { image_callback(); });
|
|
|
+ }
|
|
|
+
|
|
|
+ ~impl() {
|
|
|
+ conn.disconnect();
|
|
|
+ }
|
|
|
+};
|
|
|
+
|
|
|
+video_stabilization_ui::video_stabilization_ui(create_config conf)
|
|
|
+ : pimpl(std::make_unique<impl>(conf)) { (void) 0; }
|
|
|
+
|
|
|
+video_stabilization_ui::~video_stabilization_ui() = default;
|
|
|
+
|
|
|
+void video_stabilization_ui::show_ui() const {
|
|
|
+ pimpl->show_ui();
|
|
|
+}
|
