diff --git a/CMakeLists.txt b/CMakeLists.txt index 48a6a0a..adc0b2d 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -73,9 +73,10 @@ set(SOLVER_PATH ${SRC_ROOT_PATH}/solver) aux_source_directory(${SOLVER_PATH} SOLVER_SRC) add_library(path_manager ${SOLVER_PATH}/path_manager.cpp) +add_library(KM ${SOLVER_PATH}/KM.cpp) add_executable(solver ${SOLVER_SRC}) target_compile_definitions(solver PUBLIC ${DEBUG_FLAG}) -target_link_libraries(solver ${OpenCV_LIBRARIES} ${TORCH_LIBRARIES} tesseract utils ocr-extractor compatibility-net path_manager) +target_link_libraries(solver ${OpenCV_LIBRARIES} ${TORCH_LIBRARIES} tesseract utils ocr-extractor compatibility-net path_manager KM) diff --git a/include/KM.h b/include/KM.h new file mode 100644 index 0000000..4a06144 --- /dev/null +++ b/include/KM.h @@ -0,0 +1,22 @@ +#ifndef KM_H +#define KM_H + +#include +#include + +using namespace std; + +class KM { + +public: + KM(const vector< vector > & _edges); + vector solve(); + +private: + + int nodes_n; + const vector< vector > edges; + vector lx, ly, match_x, slack; + +}; +#endif \ No newline at end of file diff --git a/include/stripes_solver.h b/include/stripes_solver.h index 246d374..ebc96fd 100644 --- a/include/stripes_solver.h +++ b/include/stripes_solver.h @@ -17,6 +17,7 @@ #include #include #include +#include #include #include @@ -94,6 +95,7 @@ class StripesSolver { cv::Mat word_detection( const cv::Mat & img, const vector & sol, int sol_cnt=1); + void merge_single_sol(vector< vector > & composition_orders); void finetune_sols(const vector< vector > & composition_orders); }; diff --git a/src/solver/KM.cpp b/src/solver/KM.cpp new file mode 100644 index 0000000..fe6acb5 --- /dev/null +++ b/src/solver/KM.cpp @@ -0,0 +1,34 @@ +#include + +KM::KM(const vector< vector > & _edges) : + edges(_edges) { + + nodes_n = edges.size(); + lx = vector(nodes_n, 0); + ly = vector(nodes_n, 0); + match_x = vector(nodes_n, -1); + slack = vector(nodes_n, numeric_limits::max()); + + for (int i = 0; i < nodes_n; i++) { + lx[i] = numeric_limits::min(); + for (int j = 0; j < nodes_n; j++) { + lx[i] = max(lx[i], edges[i][j]); + } + } + +} + +vector KM::solve() { + + for (int i = 0; i < nodes_n; i++) { + + for (int j = 0; j < nodes_n; j++) { + slack[j] = numeric_limits::max(); + } + + while (true) { + + } + + } +} \ No newline at end of file diff --git a/src/solver/stripes_solver.cpp b/src/solver/stripes_solver.cpp index 34ae51a..80aad29 100644 --- a/src/solver/stripes_solver.cpp +++ b/src/solver/stripes_solver.cpp @@ -495,6 +495,7 @@ void StripesSolver::reassemble_greedy_probability() { stripe_pairs = path_manager.build_stripe_pairs(); vector< vector > && composition_orders = reassemble_greedy(false); + merge_single_sol(composition_orders); finetune_sols(composition_orders); } @@ -559,11 +560,59 @@ cv::Mat StripesSolver::word_detection( const cv::Mat & img, } +void StripesSolver::merge_single_sol(vector< vector > & composition_orders) { + + cout << "[INFO] Merge single composition order." << endl; + + cout << "Fragments:" << endl; + for (int i = 0; i < composition_orders.size(); i++) { + cout << "[" << i << "]\t"; + for (const int & x: composition_orders[i]) cout << x << " "; + cout << endl; + } + cout << endl; + + for (int i = 0; i < composition_orders.size(); i++) { + + if (composition_orders[i].size() > 1) continue; + + double min_score = static_cast(INFINITY); + int end_id = 0; + int single_node = composition_orders[i][0]; + + for (int j = 0; j < composition_orders.size(); j++) { + if (i == j) continue; + if (min_score > pixel_graph[single_node][composition_orders[j].front()]) { + min_score = pixel_graph[single_node][composition_orders[j].front()]; + end_id = -j; + } + if (min_score > pixel_graph[single_node][composition_orders[j].back()]) { + min_score = pixel_graph[single_node][composition_orders[j].back()]; + end_id = j; + } + } + + cout << single_node << " " << end_id << endl; + int composition_idx = abs(end_id); + if (end_id > 0) { + composition_orders[composition_idx].push_back(single_node); + } else { + auto p_st = composition_orders[composition_idx].begin(); + composition_orders[composition_idx].insert(p_st, single_node); + } + + composition_orders.erase(composition_orders.begin() + i); + i--; + + } + +} + void StripesSolver::finetune_sols(const vector< vector > & composition_orders) { cout << "[INFO] Finetune composition orders." << endl; - cout << endl << "Fragments:" << endl; + cout << "Fragments:" << endl; for (int i = 0; i < composition_orders.size(); i++) { cout << "[" << i << "]\t"; for (const int & x: composition_orders[i]) cout << x << " "; @@ -572,29 +621,24 @@ void StripesSolver::finetune_sols(const vector< vector > & composition_orde cout << endl; vector in_nodes, out_nodes; - for (const auto & sol: composition_orders) { - // for (int i = 1; i < sol.size(); i++) { - // cout << sol[i - 1] << "\t->\t" << sol[i] << "\tpixel-level score: " << pixel_graph[sol[i-1]][sol[i]] << endl; - // } in_nodes.push_back(sol.front()); out_nodes.push_back(sol.back()); } - vector end_pairs; + int nodes_n = in_nodes.size(); + vector< vector > edges(nodes_n, vector(nodes_n)); + for (int i = 0; i < out_nodes.size(); i++) { for (int j = 0; j < in_nodes.size(); j++) { if (i == j) continue; - StripePair sp(out_nodes[i], in_nodes[j], pixel_graph[out_nodes[i]][in_nodes[j]], 1, true); - end_pairs.push_back(sp); + if (pixel_graph[out_nodes[i]][in_nodes[j]] < -eps) continue; + edges[i][j] = pixel_graph[out_nodes[i]][in_nodes[j]]; } } - sort(end_pairs.begin(), end_pairs.end()); - - for (const auto & sp: end_pairs) { - cout << sp << endl; - } + KM KM_solver(edges); + vector && matches = KM_solver.solve(); composition_order = composition_orders[0];