add KM

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)
 
 

include/KM.h

@@ -0,0 +1,22 @@
+#ifndef KM_H
+#define KM_H
+
+#include <limits>
+#include <vector>
+
+using namespace std;
+
+class KM {
+
+public:
+    KM(const vector< vector<int> > & _edges);
+    vector<int> solve();
+
+private:
+
+    int nodes_n;
+    const vector< vector<int> > edges;
+    vector<int> lx, ly, match_x, slack;
+
+};
+#endif

include/stripes_solver.h

@@ -17,6 +17,7 @@
 #include <fragment.h>
 #include <path_manager.h>
 #include <utils.h>
+#include <KM.h>
 #include <compatibility_net.h>
 #include <stripe_pair.h>
 
@@ -94,6 +95,7 @@ class StripesSolver {
     cv::Mat word_detection( const cv::Mat & img, 
                             const vector<int> & sol,
                             int sol_cnt=1);
+    void merge_single_sol(vector< vector<int> > & composition_orders);
     void finetune_sols(const vector< vector<int> > & composition_orders);
 
 };

src/solver/KM.cpp

@@ -0,0 +1,34 @@
+#include <KM.h>
+
+KM::KM(const vector< vector<int> > & _edges) :
+    edges(_edges) {
+
+    nodes_n = edges.size();
+    lx = vector<int>(nodes_n, 0);
+    ly = vector<int>(nodes_n, 0);
+    match_x = vector<int>(nodes_n, -1);
+    slack = vector<int>(nodes_n, numeric_limits<int>::max());
+
+    for (int i = 0; i < nodes_n; i++) {
+        lx[i] = numeric_limits<int>::min();
+        for (int j = 0; j < nodes_n; j++) {
+            lx[i] = max(lx[i], edges[i][j]);
+        }
+    }
+
+}
+
+vector<int> KM::solve() {
+
+    for (int i = 0; i < nodes_n; i++) {
+
+        for (int j = 0; j < nodes_n; j++) {
+            slack[j] = numeric_limits<int>::max();
+        }
+
+        while (true) {
+
+        }
+
+    }
+}

src/solver/stripes_solver.cpp

@@ -495,6 +495,7 @@ void StripesSolver::reassemble_greedy_probability() {
     stripe_pairs = path_manager.build_stripe_pairs();
     vector< vector<int> > && 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<int> > & 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<double>(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<int> > & 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<int> > & composition_orde
     cout << endl;
 
     vector<int> 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<StripePair> end_pairs;
+    int nodes_n = in_nodes.size();
+    vector< vector<int> > edges(nodes_n, vector<int>(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<int> && matches = KM_solver.solve();
 
     composition_order = composition_orders[0];