|
|
|
|
|
|
|
|
|
|
|
|
|
|
#include "quad_rectify/quad_rectify.h" |
|
|
#include "non_maximal_suppression/non_maximal_suppression.h" |
|
|
#include "geometry_api/geometry_api.h" |
|
|
#include "beam_decode/beam_decode.h" |
|
|
#include "better_grid_sample/grid_sample.h" |
|
|
#include "sparse_select/sparse_select.h" |
|
|
#include "text_region_grouping/text_region_grouping.h" |
|
|
#include "local_ips/local_ips.h" |
|
|
|
|
|
#include <torch/extension.h> |
|
|
#include <pybind11/stl.h> |
|
|
|
|
|
|
|
|
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { |
|
|
m.def("quad_rectify_calc_quad_width", &quad_rectify_calc_quad_width, |
|
|
"Quad Rectify Calc Quad Width C++", |
|
|
py::arg("quads"), |
|
|
py::arg("output_height"), |
|
|
py::arg("round_factor") = 16, |
|
|
py::arg("max_width") = 0 |
|
|
); |
|
|
m.def("quad_rectify_forward", &quad_rectify_forward, "Quad Rectify Forward C++", |
|
|
py::arg("quads"), |
|
|
py::arg("image_height"), py::arg("image_width"), |
|
|
py::arg("output_height"), py::arg("output_width"), |
|
|
py::arg("isotropic") = true |
|
|
); |
|
|
m.def("quad_rectify_backward", &quad_rectify_backward, "Quad Rectify Backward C++", |
|
|
py::arg("quads"), py::arg("grad_output"), |
|
|
py::arg("image_height"), py::arg("image_width"), |
|
|
py::arg("isotropic") = true |
|
|
); |
|
|
m.def("quad_non_maximal_suppression", &quad_non_maximal_suppression, "Quad Non-Maximal Suppression C++", |
|
|
py::arg("quads"), py::arg("probs"), |
|
|
py::arg("prob_threshold"), py::arg("iou_threshold"), |
|
|
py::arg("kernel_height"), py::arg("kernel_width"), |
|
|
py::arg("max_regions"), |
|
|
py::arg("verbose") = false |
|
|
); |
|
|
|
|
|
py::class_<LanguageModel>(m, "LanguageModel"); |
|
|
|
|
|
m.def("beam_decode", &beam_decode, "beam_decode c++", |
|
|
py::arg("probs"), |
|
|
py::arg("beam_size") = 100, |
|
|
py::arg("blank") = 0, |
|
|
py::arg("min_prob") = 0.001, |
|
|
py::arg("lang_model") = static_cast<LanguageModel*>(nullptr), |
|
|
py::arg("lm_weight") = 1, |
|
|
py::arg("combine_duplicates") = true |
|
|
); |
|
|
|
|
|
py::class_<TokenMappingWrapper, TokenMappingWrapper::Ptr>(m, "TokenMapping"); |
|
|
|
|
|
m.def("create_token_mapping", &create_token_mapping, "create token mapping c++", |
|
|
py::arg("token_mapping") |
|
|
); |
|
|
|
|
|
m.def("decode_sequences", &decode_sequences, "decode_sequences c++", |
|
|
py::arg("tokens"), py::arg("language_model"), |
|
|
py::arg("probs") = nullptr |
|
|
); |
|
|
|
|
|
m.def("create_sbo_lm", &create_sbo_lm, "create_sbo_lm c++", |
|
|
py::arg("data_file_path"), |
|
|
py::arg("token_mapping"), |
|
|
py::arg("backoff") = 0.4 |
|
|
); |
|
|
|
|
|
m.def("indirect_grid_sample_forward", &indirect_grid_sample_forward, "indirect_grid_sample::forward c++", |
|
|
py::arg("input"), py::arg("grid"), py::arg("input_indices"), py::arg("method") |
|
|
); |
|
|
m.def("indirect_grad_sample_backward", &indirect_grad_sample_backward, "indirect_grid_sample::backward c++", |
|
|
py::arg("grad_output"), py::arg("input"), py::arg("grid"), py::arg("input_indices"), py::arg("method") |
|
|
); |
|
|
m.def("region_counts_to_indices", ®ion_counts_to_indices, "region counts to indices", |
|
|
py::arg("region_counts"), py::arg("num_outputs") |
|
|
); |
|
|
|
|
|
m.def("rrect_to_quads", &rrect_to_quads, "convert rotated rectangle to quadrangles", |
|
|
py::arg("rrects"), py::arg("cell_size") |
|
|
); |
|
|
m.def("rrect_to_quads_backward", &rrect_to_quads_backward, "gradient of rrect_to_quads", |
|
|
py::arg("rrects"), py::arg("grad_output") |
|
|
); |
|
|
|
|
|
m.def("sparse_select", &sparse_select, "Select sparse tensor(s) given a set of indices", |
|
|
py::arg("sparse_counts"), py::arg("sparse_tensors"), py::arg("select_indices") |
|
|
); |
|
|
|
|
|
m.def("text_region_grouping", &text_region_grouping, "Clusters all of the text into lines and phrases", |
|
|
py::arg("quads"), py::arg("counts"), |
|
|
py::arg("horizontal_tolerance") = 2.0f, |
|
|
py::arg("vertical_tolerance") = 0.5f, |
|
|
py::arg("verbose") = false |
|
|
); |
|
|
|
|
|
m.def("dense_relations_to_graph", &dense_relations_to_graph, "Converts a dense relational tensor to a graph", |
|
|
py::arg("relations") |
|
|
); |
|
|
|
|
|
m.def("ragged_quad_all_2_all_distance_v2", &ragged_quad_all_2_all_distance_v2, "get the all-to-all distances in ragged-batch quad mode", |
|
|
py::arg("embed_quads"), py::arg("region_counts"), |
|
|
py::arg("x_factor") = 1.0f, |
|
|
py::arg("y_factor") = 1.0f, |
|
|
py::arg("allow_self_distance") = true |
|
|
); |
|
|
|
|
|
m.def("calc_poly_min_rrect", &calc_poly_min_rrect, "calculate a reasonable bounding rectangle for a given text polygon", |
|
|
py::arg("vertices") |
|
|
); |
|
|
|
|
|
m.def("get_rel_continuation_cos", &get_rel_continuation_cos, "c++ get relation cosine between 2 regions", |
|
|
py::arg("rrect_a"), py::arg("rrect_b") |
|
|
); |
|
|
|
|
|
m.def("get_poly_bounds_quad", &get_poly_bounds_quad, "c++ get polygon bounds", |
|
|
py::arg("poly") |
|
|
); |
|
|
} |
|
|
|
