Skip to content

Commit

Permalink
tmp
Browse files Browse the repository at this point in the history
  • Loading branch information
@xmlyqing00
xmlyqing00 committed Dec 2, 2018
1 parent 6c4e9f7 commit 46b7abd
Show file tree
Hide file tree
Showing 5 changed files with 63 additions and 28 deletions.
0 data/models/eng.traineddata → data/tesseract_model/eng.traineddata
View file
File renamed without changes.
22 changes: 15 additions & 7 deletions include/stripes_solver.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,8 +24,9 @@ class StripePair {
int stripe_idx0;
int stripe_idx1;
double m_score;
double ac_prob;

StripePair(int _stripe_idx0, int _stripe_idx1, double _m_score);
StripePair(int _stripe_idx0, int _stripe_idx1, double _m_score, double _ac_prob=1);
bool operator < (const StripePair & _sp) const {
return m_score > _sp.m_score;
}
Expand Down Expand Up @@ -54,21 +55,18 @@ class StripesSolver {
vector<int> comp_idx;
cv::Mat comp_img;

const string tesseract_model_path {"data/tesseract_model/"};
OcrExtractor ocr_ectractor;

StripesSolver();
~StripesSolver();

void push(const cv::Mat & stripe_img);

void m_metric();

bool reassemble(Metric _metric_mode, Composition _composition_mode);

double m_metric_pixel(const cv::Mat & piece0, const cv::Mat & piece1);

double m_metric_word(const cv::Mat & piece0, const cv::Mat & piece1);

double m_metric_comp_eva(const cv::Mat & piece0, const cv::Mat & piece1);

void save_result(const string & case_name);

private:
Expand All @@ -86,7 +84,17 @@ class StripesSolver {

Metric metric_mode;
Composition composition_mode;

vector<StripePair> stripe_pairs;

double m_metric_pixel(const cv::Mat & piece0, const cv::Mat & piece1);
double m_metric_word(const cv::Mat & piece0, const cv::Mat & piece1);
double m_metric_comp_eva(const cv::Mat & piece0, const cv::Mat & piece1);

bool reassemble_greedy();
bool reassemble_greedy_probability();



};

Expand Down
1 change: 0 additions & 1 deletion src/ocr_extractor.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -101,7 +101,6 @@ void OcrExtractor::extract_img(const cv::Mat & piece) {

for (int y = y0; y < y1; y++) {
for (int x = width_half - 2; width_half + 2; width_half++) {
if ()
}
}
}
Expand Down
17 changes: 9 additions & 8 deletions src/solver/solve_puzzle.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,10 +6,6 @@ void solve_stripes( const string & stripes_folder,
StripesSolver::Metric metric_mode,
StripesSolver::Composition composition_mode) {

cout << "Metric mode: \t" << metric_mode << endl;
cout << "Composition mode: \t" << composition_mode << endl;
cout << endl;

StripesSolver stripes_solver;

cout << "[INFO] Import stripes." << endl;
Expand All @@ -20,7 +16,7 @@ void solve_stripes( const string & stripes_folder,
cerr << "[ERR] Stripe img does not exist." << endl;
exit(-1);
}
stripes_solver.push(stripe_img);
stripes_solver.push(move(stripe_img));
}

stripes_solver.reassemble(metric_mode, composition_mode);
Expand Down Expand Up @@ -78,7 +74,7 @@ int main(int argc, char ** argv) {
string case_name = "doc0";
PuzzleType puzzle_type = PuzzleType::STRIPES;
int vertical_n = 4;
StripesSolver::Composition composition_mode = StripesSolver::GREEDY_PROBABILITY;
StripesSolver::Composition composition_mode = StripesSolver::GREEDY;
StripesSolver::Metric metric_mode = StripesSolver::PIXEL;

// Parse command line parameters
Expand Down Expand Up @@ -107,11 +103,16 @@ int main(int argc, char ** argv) {
opt = getopt(argc, argv, opt_str.c_str());
}

const string metric_mode_str =
metric_mode == StripesSolver::Metric::PIXEL ? "Pixel" :
metric_mode == StripesSolver::Metric::WORD ? "Word" :
"Compability Evaluator";

cout << "Test case name: \t" << case_name << endl;
cout << "Vertical cut num: \t" << vertical_n << endl;
cout << "Puzzle type: \t" << (puzzle_type == PuzzleType::SQUARES ? "Squares": "Stripes") << endl;
cout << "Metric mode: \t" << (static_cast<int>(metric_mode)) << endl;
cout << "Composition mode: \t" << (static_cast<int>(composition_mode)) << endl;
cout << "Metric mode: \t" << metric_mode_str << endl;
cout << "Composition mode: \t" << (composition_mode == StripesSolver::Composition::GREEDY ? "Greedy" : "Greedy Probability") << endl;

// Import stripes
if (puzzle_type == PuzzleType::STRIPES) {
Expand Down
51 changes: 39 additions & 12 deletions src/solver/stripes_solver.cpp
View file
Original file line number Diff line number Diff line change
@@ -1,17 +1,18 @@
#include <stripes_solver.h>

StripePair::StripePair(int _stripe_idx0, int _stripe_idx1, double _m_score) {
StripePair::StripePair(int _stripe_idx0, int _stripe_idx1, double _m_score, double _ac_prob) {
stripe_idx0 = _stripe_idx0;
stripe_idx1 = _stripe_idx1;
m_score = _m_score;
ac_prob = _ac_prob;
}

StripesSolver::StripesSolver() {

stripes_n = 0;

ocr = new tesseract::TessBaseAPI();
if (ocr->Init(nullptr, "eng", tesseract::OEM_LSTM_ONLY)) {
if (ocr->Init(tesseract_model_path.c_str(), "eng", tesseract::OEM_LSTM_ONLY)) {
cerr << "Could not initialize tesseract." << endl;
exit(-1);
}
Expand Down Expand Up @@ -53,10 +54,13 @@ bool StripesSolver::reassemble(Metric _metric_mode, Composition _composition_mod
metric_mode = _metric_mode;
composition_mode = _composition_mode;

m_metric();

switch (composition_mode) {
case Composition::GREEDY:
cout << "Reassemble mode: \t" << "GREEDY" << endl;
return reassemble_greedy();
case Composition::GREEDY_PROBABILITY:
return reassemble_greedy_probability();
default:
return false;
}
Expand Down Expand Up @@ -166,13 +170,11 @@ double StripesSolver::m_metric_comp_eva(const cv::Mat & piece0, const cv::Mat &

}

bool StripesSolver::reassemble_greedy() {
void StripesSolver::m_metric() {

// Compute matching score for each pair
cout << "[INFO] Compute matching score for each pair." << endl;

vector<StripePair> stripe_pairs;

if (metric_mode == COMP_EVA) {

int filters_n = int(stripes_n * filter_rate);
Expand Down Expand Up @@ -208,6 +210,7 @@ bool StripesSolver::reassemble_greedy() {

vector<StripePair> candidates;
for (int j = 0; j < stripes_n; j++) {
if (i == j) continue;
double m_score = m_metric_pixel(stripes[i], stripes[j]);
candidates.push_back(StripePair(i, j, m_score));
}
Expand All @@ -227,6 +230,8 @@ bool StripesSolver::reassemble_greedy() {
} else {

for (int i = 0; i < stripes_n; i++) {

vector<StripePair> candidates;
for (int j = 0; j < stripes_n; j++) {

if (i == j) continue;
Expand All @@ -239,21 +244,39 @@ bool StripesSolver::reassemble_greedy() {
case WORD:
m_score = m_metric_word(stripes[i], stripes[j]);
break;
case COMP_EVA:
m_score = m_metric_comp_eva(stripes[i], stripes[j]);
break;
}

stripe_pairs.push_back(StripePair(i, j, m_score));
candidates.push_back(StripePair(i, j, m_score));

}

if (composition_mode == Composition::GREEDY_PROBABILITY) {

sort(candidates.begin(), candidates.end());
for (int j = 0; j < candidates.size() - 1; j++) {
// m_score are negatives.
double alpha = candidates[j + 1].m_score / candidates[j].m_score;
candidates[j].ac_prob = alpha / (1 + alpha);
}

stripe_pairs.insert(
stripe_pairs.end(),
make_move_iterator(candidates.begin()),
make_move_iterator(candidates.end()));

}



}
}



sort(stripe_pairs.begin(), stripe_pairs.end());

}

bool StripesSolver::reassemble_greedy() {

// Greedy algorithm
cout << "[INFO] Reassemble the document by greedy algorithm." << endl;

Expand Down Expand Up @@ -305,3 +328,7 @@ bool StripesSolver::reassemble_greedy() {
return true;

}

bool StripesSolver::reassemble_greedy_probability() {

}

0 comments on commit 46b7abd

Please sign in to comment.