textbox_demo.cpp 3.28 KB
Newer Older
1 2 3
#include <opencv2/text.hpp>
#include <opencv2/highgui.hpp>
#include <opencv2/imgproc.hpp>
4
#include <opencv2/dnn.hpp>
sghoshcvc's avatar
sghoshcvc committed
5 6 7 8 9

#include  <sstream>
#include  <iostream>
#include  <fstream>

10
using namespace cv;
sghoshcvc's avatar
sghoshcvc committed
11

12 13
namespace
{
14
std::string getHelpStr(const std::string& progFname)
15 16 17
{
    std::stringstream out;
    out << "    Demo of text detection CNN for text detection." << std::endl
18
        << "    Minghui Liao, Baoguang Shi, Xiang Bai, Xinggang Wang, Wenyu Liu: TextBoxes: A Fast Text Detector with a Single Deep Neural Network, AAAI2017\n\n"
19
        << "    Usage: " << progFname << " <output_file> <input_image>" << std::endl
20 21
        << "    Caffe Model files  (textbox.prototxt, TextBoxes_icdar13.caffemodel)"<<std::endl
        << "      must be in the current directory. See the documentation of text::TextDetectorCNN class to get download links." << std::endl;
sghoshcvc's avatar
sghoshcvc committed
22 23 24
    return out.str();
}

25
bool fileExists (const std::string& filename)
26
{
sghoshcvc's avatar
sghoshcvc committed
27 28 29
    std::ifstream f(filename.c_str());
    return f.good();
}
30

31
void textbox_draw(Mat src, std::vector<Rect>& groups, std::vector<float>& probs, std::vector<int>& indexes)
sghoshcvc's avatar
sghoshcvc committed
32
{
33
    for (size_t i = 0; i < indexes.size(); i++)
sghoshcvc's avatar
sghoshcvc committed
34
    {
35
        if (src.type() == CV_8UC3)
sghoshcvc's avatar
sghoshcvc committed
36
        {
37 38 39 40 41 42 43 44 45 46 47 48
            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);
sghoshcvc's avatar
sghoshcvc committed
49
        }
50 51
        else
            rectangle(src, groups[i], Scalar( 255 ), 3, 8 );
sghoshcvc's avatar
sghoshcvc committed
52 53 54
    }
}

55
}
sghoshcvc's avatar
sghoshcvc committed
56

57 58 59 60 61 62
int main(int argc, const char * argv[])
{
    if (argc < 2)
    {
        std::cout << getHelpStr(argv[0]);
        std::cout << "Insufiecient parameters. Aborting!" << std::endl;
sghoshcvc's avatar
sghoshcvc committed
63 64 65
        exit(1);
    }

66 67 68 69
    const std::string modelArch = "textbox.prototxt";
    const std::string moddelWeights = "TextBoxes_icdar13.caffemodel";

    if (!fileExists(modelArch) || !fileExists(moddelWeights))
70
    {
71
        std::cout << getHelpStr(argv[0]);
72
        std::cout << "Model files not found in the current directory. Aborting!" << std::endl;
sghoshcvc's avatar
sghoshcvc committed
73 74 75
        exit(1);
    }

76
    Mat image = imread(String(argv[1]), IMREAD_COLOR);
sghoshcvc's avatar
sghoshcvc committed
77

78 79
    std::cout << "Starting Text Box Demo" << std::endl;
    Ptr<text::TextDetectorCNN> textSpotter =
80
            text::TextDetectorCNN::create(modelArch, moddelWeights);
sghoshcvc's avatar
sghoshcvc committed
81

82
    std::vector<Rect> bbox;
sghoshcvc's avatar
sghoshcvc committed
83
    std::vector<float> outProbabillities;
84
    textSpotter->detect(image, bbox, outProbabillities);
sghoshcvc's avatar
sghoshcvc committed
85

86 87 88 89
    std::vector<int> indexes;
    cv::dnn::NMSBoxes(bbox, outProbabillities, 0.3f, 0.4f, indexes);

    textbox_draw(image, bbox, outProbabillities, indexes);
sghoshcvc's avatar
sghoshcvc committed
90

91
    imshow("TextBox Demo",image);
sghoshcvc's avatar
sghoshcvc committed
92 93
    std::cout << "Done!" << std::endl << std::endl;
    std::cout << "Press any key to exit." << std::endl << std::endl;
94 95
    waitKey();
    return 0;
sghoshcvc's avatar
sghoshcvc committed
96
}