DepthGuide/src/image_process_v5/image_process.cpp

#include "image_process.h"

#include <glm/ext/matrix_transform.hpp>
#include <glm/gtx/matrix_transform_2d.hpp>
#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
    thread_local std::optional<cv::cuda::Stream> cv_stream;

    auto &get_cv_stream() {
        current_cuda_stream(); // initialize CUDA
        if ( !cv_stream ) [[unlikely]] { cv_stream.emplace(); }
        return *cv_stream;
    }
} // namespace

size_t normal_height_to_nv12(const size_t height) {
    assert(height % 2 == 0);
    return height / 2 * 3;
}

size_t nv12_height_to_normal(const size_t height) {
    assert(height % 3 == 0);
    return height / 3 * 2;
}

cv::Size normal_size_to_nv12(const cv::Size size) {
    return cv::Size(size.width, normal_height_to_nv12(size.height));
}

cv::Size nv12_size_to_normal(cv::Size size) {
    return cv::Size(size.width, nv12_height_to_normal(size.height));
}

sp_image nv12_luma_view(const sp_image &img) {
    assert(img.cv_type() == CV_8UC1);
    const auto luma_size = nv12_size_to_normal(img.cv_size());
    return img.sub_view(luma_size);
}

sp_image nv12_chrome_view(const sp_image &img) {
    assert(img.cv_type() == CV_8UC1);
    const auto chroma_size = cv::Size(img.width(), img.height() / 3);
    const auto img_chrome  = img.sub_view(chroma_size, cv::Size(0, nv12_height_to_normal(img.height())));
    return img_chrome.cast_view(CV_8UC2);
}

cv_stream_guard::cv_stream_guard() {
    push_cuda_stream((cudaStream_t) get_cv_stream().cudaPtr());
}

cv_stream_guard::~cv_stream_guard() {
    pop_cuda_stream();
}

cv::cuda::Stream &cv_stream_guard::cv_stream() {
    return get_cv_stream();
}

namespace {
    struct image_opencv_cuda_helper {
        const sp_image *read;
        sp_image       *write;
        using proxy_type = auto_memory_info::cuda_proxy;
        cv_stream_guard                            stream_guard;
        pair_access_helper<proxy_type, proxy_type> access_helper;

        image_opencv_cuda_helper(const sp_image &src, sp_image &dst)
            : read(&src), write(&dst), //
              access_helper(read->cuda(), write->cuda()) {
            (void) 0;
        }

        [[nodiscard]] cv::cuda::GpuMat input() const {
            return read->cv_gpu_mat(access_helper.read_ptr());
        }

        [[nodiscard]] cv::cuda::GpuMat output() const {
            return write->cv_gpu_mat(access_helper.write_ptr());
        }
    };

    struct image_opencv_host_helper {
        const sp_image *read;
        sp_image       *write;
        using proxy_type = auto_memory_info::host_proxy;
        cv_stream_guard                            stream_guard;
        pair_access_helper<proxy_type, proxy_type> access_helper;

        image_opencv_host_helper(const sp_image &src, sp_image &dst)
            : read(&src), write(&dst), //
              access_helper(read->host(), write->host()) {
            (void) 0;
        }

        [[nodiscard]] cv::Mat input() const {
            return read->cv_mat(access_helper.read_ptr());
        }

        [[nodiscard]] cv::Mat output() const {
            return write->cv_mat(access_helper.write_ptr());
        }
    };
} // namespace

sp_image image_debayer(const sp_image &img) {
    assert(img.cv_type() == CV_8UC1);
    auto       ret    = sp_image::create<uchar3>(img.cv_size());
    const auto helper = image_opencv_cuda_helper(img, ret);
    cv::cuda::cvtColor(helper.input(), helper.output(), //
        cv::COLOR_BayerRG2BGR, 3, get_cv_stream());
    ret.merge_meta(img);
    return ret;
}

void image_resize(const sp_image &src, sp_image &dst) {
    assert(src.cv_type() == dst.cv_type());
    const auto helper = image_opencv_cuda_helper(src, dst);
    cv::cuda::resize(helper.input(), helper.output(), //
        dst.cv_size(), 0, 0, cv::INTER_LINEAR, get_cv_stream());
    dst.merge_meta(src);
}

sp_image image_resize(const sp_image &img, const cv::Size size) {
    auto ret = sp_image::create(img.cv_type(), size);
    image_resize(img, ret);
    return ret;
}

sp_image image_flip_y(const sp_image &img) {
    auto       ret    = sp_image::create(img.cv_type(), img.cv_size());
    const auto helper = image_opencv_cuda_helper(img, ret);
    cv::cuda::flip(helper.input(), helper.output(), 1, get_cv_stream()); // flip vertically
    ret.merge_meta(img);
    return ret;
}

sp_image image_warp_affine(const sp_image &img, const glm::mat3 &matrix) {
    auto cv_matrix = cv::Mat(2, 3, CV_32FC1);
    for ( auto i = 0; i < 3; ++i )
        for ( auto j = 0; j < 2; ++j ) { cv_matrix.at<float>(j, i) = matrix[i][j]; }
    auto       ret    = sp_image::create_like(img);
    const auto helper = image_opencv_cuda_helper(img, ret);
    cv::cuda::warpAffine(helper.input(), helper.output(), cv_matrix, img.cv_size(), //
        cv::INTER_LINEAR, cv::BORDER_CONSTANT, {}, get_cv_stream());
    ret.merge_meta(img);
    return ret;
}

namespace {
    float pixel_center(const float size) {
        return 0.5f * size - 0.5f;
    }
} // namespace

namespace {
    glm::mat3 cv_mat23_to_mat3(const cv::Mat &mat) {
        assert(mat.rows == 2 && mat.cols == 3);
        assert(mat.type() == CV_64FC1);
        auto ret = glm::mat3(0.f);
        for ( auto i = 0; i < 3; ++i )
            for ( auto j = 0; j < 2; ++j ) { //
                ret[i][j] = mat.at<double>(j, i);
            }
        return ret;
    }
} // namespace

sp_image image_rotate(const sp_image &img, const float angle, std::optional<glm::vec2> center) {
    if ( !center ) {
        center = glm::vec2(            //
            pixel_center(img.width()), //
            pixel_center(img.height()));
    }
    const auto center_cv = cv::Point2f(center->x, center->y);
    const auto matrix_cv = cv::getRotationMatrix2D(center_cv, angle, 1);
    return image_warp_affine(img, cv_mat23_to_mat3(matrix_cv));
}

sp_image image_translate(const sp_image &img, const glm::vec2 offset) {
    const auto matrix = glm::translate(glm::identity<glm::mat3>(), offset);
    return image_warp_affine(img, matrix);
}

image_stereo_pair image_stereo_split(const sp_image &img) {
    assert(img.width() % 2 == 0);
    const auto mono_size = cv::Size(img.width() / 2, img.height());
    auto       ret_left  = sp_image::create(img.cv_type(), mono_size);
    const auto img_left  = img.sub_view(mono_size);
    copy_sp_image(img_left, ret_left);
    auto       ret_right = sp_image::create(img.cv_type(), mono_size);
    const auto img_right = img.sub_view(mono_size, cv::Size(img.width() / 2, 0));
    copy_sp_image(img_right, ret_right);
    return std::make_tuple(ret_left, ret_right);
}

image_stereo_pair image_stereo_split_view(const sp_image &img) {
    assert(img.width() % 2 == 0);
    const auto mono_size = cv::Size(img.width() / 2, img.height());
    const auto img_left  = img.sub_view(mono_size);
    const auto img_right = img.sub_view(mono_size, cv::Size(img.width() / 2, 0));
    return std::make_tuple(img_left, img_right);
}

sp_image image_stereo_combine(const sp_image &left, const sp_image &right) {
    assert(left.cv_type() == right.cv_type());
    assert(left.shape_array() == right.shape_array());
    const auto stereo_shape      = cv::Size(left.width() * 2, left.height());
    auto       ret_img           = sp_image::create(left.cv_type(), stereo_shape);
    auto [left_view, right_view] = image_stereo_split_view(ret_img);
    copy_sp_image(left, left_view);
    copy_sp_image(right, right_view);
    ret_img.merge_meta(left);
    ret_img.merge_meta(right);
    return ret_img;
}

#include "image_process/cuda_impl/pixel_convert.cuh"

namespace {
    template<typename Input, typename Output>
    struct image_cuda_v2_helper {
        const sp_image *read;
        sp_image       *write;
        using proxy_type = auto_memory_info::cuda_proxy;
        pair_access_helper<proxy_type, proxy_type> access_helper;

        image_cuda_v2_helper(const sp_image &src, sp_image &dst)
            : read(&src), write(&dst), access_helper(read->cuda(), write->cuda()) {
            (void) 0;
        }

        template<typename T = Input>
        image_type_v2<T> input() {
            return to_cuda_v2(read->as_ndarray<T>(access_helper.read_ptr()));
        }

        template<typename T = Output>
        image_type_v2<T> output() {
            return to_cuda_v2(write->as_ndarray<T>(access_helper.write_ptr()));
        }
    };
} // 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());
    auto helper = image_cuda_v2_helper<uchar3, uchar4>(img, ret);
    call_cvt_rgb_bgra_u8(helper.input(), helper.output(), current_cuda_stream());
    ret.merge_meta(img);
    return ret;
}

sp_image image_rgb_to_gray(const sp_image &img) {
    assert(img.cv_type() == CV_8UC3);
    auto       ret    = sp_image::create<uchar1>(img.cv_size());
    const auto helper = image_opencv_cuda_helper(img, ret);
    cv::cuda::cvtColor(helper.input(), helper.output(), cv::COLOR_RGB2GRAY, 1, get_cv_stream());
    ret.merge_meta(img);
    return ret;
}

sp_image image_gray_to_binary(const sp_image &img, float threshold) {
    assert(img.cv_type() == CV_8UC1);
    auto ret    = sp_image::create<uchar1>(img.cv_size());
    auto status = img.mem->status();
    if ( status.cuda_available ) {
        const auto helper = image_opencv_cuda_helper(img, ret);
        cv::cuda::threshold(helper.input(), helper.output(), //
            threshold, 255, cv::THRESH_BINARY);
    } else {
        assert(status.host_available);
        const auto helper = image_opencv_host_helper(img, ret);
        cv::threshold(helper.input(), helper.output(), //
            threshold, 255, cv::THRESH_BINARY);
    }
    ret.merge_meta(img);
    return ret;
}

sp_image image_rgb_to_nv12(const sp_image &img) {
    assert(img.cv_type() == CV_8UC3);
    auto ret    = sp_image::create(CV_8UC1, normal_size_to_nv12(img.cv_size()));
    auto helper = image_cuda_v2_helper<uchar3, uchar1>(img, ret);
    call_rgb_to_nv12(helper.input(), helper.output(), current_cuda_stream());
    ret.merge_meta(img);
    return ret;
}

sp_image image_nv12_to_rgb(const sp_image &img) {
    assert(img.cv_type() == CV_8UC1);
    auto ret    = sp_image::create(CV_8UC3, nv12_size_to_normal(img.cv_size()));
    auto helper = image_cuda_v2_helper<uchar1, uchar3>(img, ret);
    call_nv12_to_rgb(helper.input(), helper.output(), current_cuda_stream());
    ret.merge_meta(img);
    return ret;
}

sp_image image_yuyv_to_rgb(const sp_image &img) {
    assert(img.cv_type() == CV_8UC2);
    auto ret    = sp_image::create(CV_8UC3, img.cv_size());
    auto helper = image_cuda_v2_helper<uchar2, uchar3>(img, ret);
    call_yuyv_to_rgb(helper.input(), helper.output(), current_cuda_stream());
    ret.merge_meta(img);
    return ret;
}

sp_image image_remap_np_to_tex(const sp_image &img, const float fov, const float aspect) {
    assert(img.cv_type() == CV_32FC2);
    auto ret    = sp_image::create(CV_32FC2, img.cv_size());
    auto helper = image_cuda_v2_helper<float2, float2>(img, ret);
    call_np_to_tex(helper.input(), helper.output(), fov, aspect, current_cuda_stream());
    ret.merge_meta(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);
    }
} // namespace

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 {
    create_config conf;
    obj_conn_type conn;

    void image_callback_impl() {
        const auto img      = OBJ_QUERY(sp_image, conf.in_name);
        auto       ret_rect = simple_rect(0, 0, conf.size.width, conf.size.height);
        ret_rect            = ret_rect.fit_aspect(img.cv_size().aspectRatio());
        auto ret_img        = sp_image::create(img.cv_type(), conf.size);
        ret_img.initialize_meta();
        auto ret_view = ret_img.sub_view(cv::Size(ret_rect.width, ret_rect.height), //
            cv::Size(ret_rect.x, ret_rect.y));
        image_resize(img, ret_view);
        if ( conf.flip_y ) ret_img = image_flip_y(ret_img);
        OBJ_SAVE(conf.out_name, ret_img);
    }

    void image_callback(const obj_name_type _name) {
        assert(conf.in_name == _name);
        try {
            image_callback_impl();
        } catch ( ... ) { (void) 0; }
    }

    explicit impl(const create_config _conf) : conf(_conf) {
        conn = OBJ_SIG(conf.in_name)->connect([this](auto name) { image_callback(name); });
    }

    ~impl() {
        conn.disconnect();
    }
};

image_output_helper::image_output_helper(create_config conf) : pimpl(std::make_unique<impl>(conf)) { }

image_output_helper::~image_output_helper() = default;

struct stereo_output_helper::impl {
    create_config conf;
    obj_conn_type left_conn, right_conn;

    bool left_updated  = false;
    bool right_updated = false;

    void image_callback_impl() {
        const auto left_img  = OBJ_QUERY(sp_image, conf.left_name);
        const auto right_img = OBJ_QUERY(sp_image, conf.right_name);
        assert(left_img.cv_type() == right_img.cv_type());
        assert(left_img.cv_size() == right_img.cv_size());
        auto ret_size = conf.size;
        if ( ret_size.empty() ) {
            if ( conf.halve_width ) {
                ret_size = left_img.cv_size();
            } else {
                ret_size = cv::Size(left_img.width() * 2, left_img.height());
            }
        }
        assert(ret_size.width % 2 == 0);
        auto ret_rect = simple_rect(0, 0, //
            conf.halve_width ? ret_size.width : (ret_size.width / 2), ret_size.height);
        ret_rect      = ret_rect.fit_aspect(left_img.cv_size().aspectRatio());
        if ( conf.halve_width ) {
            ret_rect.x /= 2;
            ret_rect.width /= 2;
        }
        auto ret_img = sp_image::create(left_img.cv_type(), ret_size);
        ret_img.initialize_meta();
        auto left_view = ret_img.sub_view(cv::Size(ret_rect.width, ret_rect.height), //
            cv::Size(ret_rect.x, ret_rect.y));
        image_resize(left_img, left_view);
        auto right_view = ret_img.sub_view(cv::Size(ret_rect.width, ret_rect.height), //
            cv::Size(ret_rect.x + ret_size.width / 2, ret_rect.y));
        image_resize(right_img, right_view);
        if ( conf.flip_y ) ret_img = image_flip_y(ret_img);
        OBJ_SAVE(conf.out_name, ret_img);
    }

    void image_callback(const obj_name_type name) {
        if ( name == conf.left_name ) left_updated = true;
        if ( name == conf.right_name ) right_updated = true;
        if ( !left_updated || !right_updated ) return;
        try {
            image_callback_impl();
        } catch ( ... ) { (void) 0; }
        left_updated  = false;
        right_updated = false;
    }

    explicit impl(const create_config &_conf) : conf(_conf) {
        left_conn = OBJ_SIG(conf.left_name) //
                        ->connect([this](auto name) { image_callback(name); });
        right_conn = OBJ_SIG(conf.right_name) //
                         ->connect([this](auto name) { image_callback(name); });
    }

    ~impl() {
        left_conn.disconnect();
        right_conn.disconnect();
    }
};

stereo_output_helper::stereo_output_helper(create_config conf) : pimpl(std::make_unique<impl>(conf)) { }

stereo_output_helper::~stereo_output_helper() = default;