File: text_recognition_cnn.cpp

package info (click to toggle)
opencv 4.5.1%2Bdfsg-5
  • links: PTS, VCS
  • area: main
  • in suites: bullseye
  • size: 268,248 kB
  • sloc: cpp: 969,170; xml: 682,525; python: 36,732; lisp: 30,170; java: 25,155; ansic: 7,927; javascript: 5,643; objc: 2,041; sh: 935; cs: 601; perl: 494; makefile: 145
file content (122 lines) | stat: -rw-r--r-- 4,717 bytes parent folder | download | duplicates (3)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
#include <opencv2/text.hpp>
#include <opencv2/highgui.hpp>
#include <opencv2/imgproc.hpp>
#include <opencv2/dnn.hpp>

#include  <iostream>
#include  <fstream>

using namespace cv;
using namespace std;

namespace
{
void printHelpStr(const string& progFname)
{
    cout << "   Demo of text recognition CNN for text detection." << endl
         << "   Max Jaderberg et al.: Reading Text in the Wild with Convolutional Neural Networks, IJCV 2015"<<endl<<endl
         << "   Usage: " << progFname << " <output_file> <input_image>" << endl
         << "   Caffe Model files (textbox.prototxt, TextBoxes_icdar13.caffemodel)"<<endl
         << "     must be in the current directory. See the documentation of text::TextDetectorCNN class to get download links." << endl
         << "   Obtaining text recognition Caffe Model files in linux shell:" << endl
         << "   wget http://nicolaou.homouniversalis.org/assets/vgg_text/dictnet_vgg.caffemodel" << endl
         << "   wget http://nicolaou.homouniversalis.org/assets/vgg_text/dictnet_vgg_deploy.prototxt" << endl
         << "   wget http://nicolaou.homouniversalis.org/assets/vgg_text/dictnet_vgg_labels.txt" <<endl << endl;
}

bool fileExists (const string& filename)
{
    ifstream f(filename.c_str());
    return f.good();
}

void textbox_draw(Mat src, std::vector<Rect>& groups, std::vector<float>& probs, std::vector<int>& indexes)
{
    for (size_t i = 0; i < indexes.size(); i++)
    {
        if (src.type() == CV_8UC3)
        {
            Rect currrentBox = groups[indexes[i]];
            rectangle(src, currrentBox, Scalar( 0, 255, 255 ), 2, LINE_AA);
            String label = format("%.2f", probs[indexes[i]]);
            std::cout << "text box: " << currrentBox << " confidence: " << probs[indexes[i]] << "\n";

            int baseLine = 0;
            Size labelSize = getTextSize(label, FONT_HERSHEY_PLAIN, 1, 1, &baseLine);
            int yLeftBottom = std::max(currrentBox.y, labelSize.height);
            rectangle(src, Point(currrentBox.x, yLeftBottom - labelSize.height),
                      Point(currrentBox.x + labelSize.width, yLeftBottom + baseLine), Scalar( 255, 255, 255 ), FILLED);

            putText(src, label, Point(currrentBox.x, yLeftBottom), FONT_HERSHEY_PLAIN, 1, Scalar( 0,0,0 ), 1, LINE_AA);
        }
        else
            rectangle(src, groups[i], Scalar( 255 ), 3, 8 );
    }
}

}

int main(int argc, const char * argv[])
{
    if (argc < 2)
    {
        printHelpStr(argv[0]);
        cout << "Insufiecient parameters. Aborting!" << endl;
        exit(1);
    }

    const string modelArch = "textbox.prototxt";
    const string moddelWeights = "TextBoxes_icdar13.caffemodel";

    if (!fileExists(modelArch) || !fileExists(moddelWeights))
    {
        printHelpStr(argv[0]);
        cout << "Model files not found in the current directory. Aborting!" << endl;
        exit(1);
    }

    Mat image = imread(String(argv[1]), IMREAD_COLOR);

    cout << "Starting Text Box Demo" << endl;
    Ptr<text::TextDetectorCNN> textSpotter =
            text::TextDetectorCNN::create(modelArch, moddelWeights);

    vector<Rect> bbox;
    vector<float> outProbabillities;
    textSpotter->detect(image, bbox, outProbabillities);
    std::vector<int> indexes;
    cv::dnn::NMSBoxes(bbox, outProbabillities, 0.4f, 0.5f, indexes);

    Mat image_copy = image.clone();
    textbox_draw(image_copy, bbox, outProbabillities, indexes);
    imshow("Text detection", image_copy);
    image_copy = image.clone();

    Ptr<text::OCRHolisticWordRecognizer> wordSpotter =
            text::OCRHolisticWordRecognizer::create("dictnet_vgg_deploy.prototxt", "dictnet_vgg.caffemodel", "dictnet_vgg_labels.txt");

    for(size_t i = 0; i < indexes.size(); i++)
    {
        Mat wordImg;
        cvtColor(image(bbox[indexes[i]]), wordImg, COLOR_BGR2GRAY);
        string word;
        vector<float> confs;
        wordSpotter->run(wordImg, word, NULL, NULL, &confs);

        Rect currrentBox = bbox[indexes[i]];
        rectangle(image_copy, currrentBox, Scalar( 0, 255, 255 ), 2, LINE_AA);

        int baseLine = 0;
        Size labelSize = getTextSize(word, FONT_HERSHEY_PLAIN, 1, 1, &baseLine);
        int yLeftBottom = std::max(currrentBox.y, labelSize.height);
        rectangle(image_copy, Point(currrentBox.x, yLeftBottom - labelSize.height),
                  Point(currrentBox.x + labelSize.width, yLeftBottom + baseLine), Scalar( 255, 255, 255 ), FILLED);

        putText(image_copy, word, Point(currrentBox.x, yLeftBottom), FONT_HERSHEY_PLAIN, 1, Scalar( 0,0,0 ), 1, LINE_AA);

    }
    imshow("Text recognition", image_copy);
    cout << "Recognition finished. Press any key to exit.\n";
    waitKey();
    return 0;
}