textbox_demo.cpp 3.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596
  1. #include <opencv2/text.hpp>
  2. #include <opencv2/highgui.hpp>
  3. #include <opencv2/imgproc.hpp>
  4. #include <opencv2/dnn.hpp>
  5. #include <sstream>
  6. #include <iostream>
  7. #include <fstream>
  8. using namespace cv;
  9. namespace
  10. {
  11. std::string getHelpStr(const std::string& progFname)
  12. {
  13. std::stringstream out;
  14. out << " Demo of text detection CNN for text detection." << std::endl
  15. << " Minghui Liao, Baoguang Shi, Xiang Bai, Xinggang Wang, Wenyu Liu: TextBoxes: A Fast Text Detector with a Single Deep Neural Network, AAAI2017\n\n"
  16. << " Usage: " << progFname << " <output_file> <input_image>" << std::endl
  17. << " Caffe Model files (textbox.prototxt, TextBoxes_icdar13.caffemodel)"<<std::endl
  18. << " must be in the current directory. See the documentation of text::TextDetectorCNN class to get download links." << std::endl;
  19. return out.str();
  20. }
  21. bool fileExists (const std::string& filename)
  22. {
  23. std::ifstream f(filename.c_str());
  24. return f.good();
  25. }
  26. void textbox_draw(Mat src, std::vector<Rect>& groups, std::vector<float>& probs, std::vector<int>& indexes)
  27. {
  28. for (size_t i = 0; i < indexes.size(); i++)
  29. {
  30. if (src.type() == CV_8UC3)
  31. {
  32. Rect currrentBox = groups[indexes[i]];
  33. rectangle(src, currrentBox, Scalar( 0, 255, 255 ), 2, LINE_AA);
  34. String label = format("%.2f", probs[indexes[i]]);
  35. std::cout << "text box: " << currrentBox << " confidence: " << probs[indexes[i]] << "\n";
  36. int baseLine = 0;
  37. Size labelSize = getTextSize(label, FONT_HERSHEY_PLAIN, 1, 1, &baseLine);
  38. int yLeftBottom = std::max(currrentBox.y, labelSize.height);
  39. rectangle(src, Point(currrentBox.x, yLeftBottom - labelSize.height),
  40. Point(currrentBox.x + labelSize.width, yLeftBottom + baseLine), Scalar( 255, 255, 255 ), FILLED);
  41. putText(src, label, Point(currrentBox.x, yLeftBottom), FONT_HERSHEY_PLAIN, 1, Scalar( 0,0,0 ), 1, LINE_AA);
  42. }
  43. else
  44. rectangle(src, groups[i], Scalar( 255 ), 3, 8 );
  45. }
  46. }
  47. }
  48. int main(int argc, const char * argv[])
  49. {
  50. if (argc < 2)
  51. {
  52. std::cout << getHelpStr(argv[0]);
  53. std::cout << "Insufiecient parameters. Aborting!" << std::endl;
  54. exit(1);
  55. }
  56. const std::string modelArch = "textbox.prototxt";
  57. const std::string moddelWeights = "TextBoxes_icdar13.caffemodel";
  58. if (!fileExists(modelArch) || !fileExists(moddelWeights))
  59. {
  60. std::cout << getHelpStr(argv[0]);
  61. std::cout << "Model files not found in the current directory. Aborting!" << std::endl;
  62. exit(1);
  63. }
  64. Mat image = imread(String(argv[1]), IMREAD_COLOR);
  65. std::cout << "Starting Text Box Demo" << std::endl;
  66. Ptr<text::TextDetectorCNN> textSpotter =
  67. text::TextDetectorCNN::create(modelArch, moddelWeights);
  68. std::vector<Rect> bbox;
  69. std::vector<float> outProbabillities;
  70. textSpotter->detect(image, bbox, outProbabillities);
  71. std::vector<int> indexes;
  72. cv::dnn::NMSBoxes(bbox, outProbabillities, 0.3f, 0.4f, indexes);
  73. textbox_draw(image, bbox, outProbabillities, indexes);
  74. imshow("TextBox Demo",image);
  75. std::cout << "Done!" << std::endl << std::endl;
  76. std::cout << "Press any key to exit." << std::endl << std::endl;
  77. waitKey();
  78. return 0;
  79. }