1/*
2 Text detection model (EAST): https://github.com/argman/EAST
3 Download link for EAST model: https://www.dropbox.com/s/r2ingd0l3zt8hxs/frozen_east_text_detection.tar.gz?dl=1
4
5 DB detector model:
6 https://drive.google.com/uc?export=download&id=17_ABp79PlFt9yPCxSaarVc_DKTmrSGGf
7
8 CRNN Text recognition model sourced from: https://github.com/meijieru/crnn.pytorch
9 How to convert from .pb to .onnx:
10 Using classes from: https://github.com/meijieru/crnn.pytorch/blob/master/models/crnn.py
11
12 Additional converted ONNX text recognition models available for direct download:
13 Download link: https://drive.google.com/drive/folders/1cTbQ3nuZG-EKWak6emD_s8_hHXWz7lAr?usp=sharing
14 These models are taken from: https://github.com/clovaai/deep-text-recognition-benchmark
15
16 Importing and using the CRNN model in PyTorch:
17 import torch
18 from models.crnn import CRNN
19
20 model = CRNN(32, 1, 37, 256)
21 model.load_state_dict(torch.load('crnn.pth'))
22 dummy_input = torch.randn(1, 1, 32, 100)
23 torch.onnx.export(model, dummy_input, "crnn.onnx", verbose=True)
24
25 Usage: ./example_dnn_text_detection DB
26*/
27#include <iostream>
28#include <fstream>
29
30#include <opencv2/geometry.hpp>
31#include <opencv2/imgproc.hpp>
32#include <opencv2/highgui.hpp>
33#include <opencv2/dnn.hpp>
34
35#include "common.hpp"
36
37using namespace cv;
38using namespace std;
39using namespace cv::dnn;
40
41const string about = "Use this script for Text Detection and Recognition using OpenCV. \n\n"
42 "Firstly, download required models using `download_models.py` (if not already done). Set environment variable OPENCV_DOWNLOAD_CACHE_DIR to point to the directory where models are downloaded. Also, point OPENCV_SAMPLES_DATA_PATH to opencv/samples/data.\n"
43 "To run:\n"
44 "\t Example: ./example_dnn_text_detection modelName(i.e. DB or East) --ocr_model=<path to VGG_CTC.onnx>\n\n"
45 "Detection model path can also be specified using --model argument. \n\n"
46 "Download ocr model using: python download_models.py OCR \n\n";
47
48// Command-line keys to parse the input arguments
49string keys =
50 "{ help h | | Print help message. }"
51 "{ input i | right.jpg | Path to an input image. }"
52 "{ @alias | | An alias name of model to extract preprocessing parameters from models.yml file. }"
53 "{ zoo | ../dnn/models.yml | An optional path to file with preprocessing parameters }"
54 "{ ocr_model | | Path to a binary .onnx model for recognition. }"
55 "{ model | | Path to detection model file. }"
56 "{ thr | 0.5 | Confidence threshold for EAST detector. }"
57 "{ nms | 0.4 | Non-maximum suppression threshold for EAST detector. }"
58 "{ binaryThreshold bt | 0.3 | Confidence threshold for the binary map in DB detector. }"
59 "{ polygonThreshold pt | 0.5 | Confidence threshold for polygons in DB detector. }"
60 "{ maxCandidate max | 200 | Max candidates for polygons in DB detector. }"
61 "{ unclipRatio ratio | 2.0 | Unclip ratio for DB detector. }"
62 "{ vocabularyPath vp | alphabet_36.txt | Path to vocabulary file. }";
63
64// Function prototype for the four-point perspective transform
65static void fourPointsTransform(const Mat& frame, const Point2f vertices[], Mat& result);
66static void processFrame(
67 const Mat& frame,
68 const vector<vector<Point>>& detResults,
69 const std::string& ocr_model,
70 bool imreadRGB,
71 Mat& board,
72 FontFace& fontFace,
73 int fontSize,
74 int fontWeight,
75 const vector<std::string>& vocabulary
76);
77
78int main(int argc, char** argv) {
79 // Setting up command-line parser with the specified keys
80 CommandLineParser parser(argc, argv, keys);
81
82 if (!parser.has("@alias") || parser.has("help"))
83 {
84 cout << about << endl;
85 parser.printMessage();
86 return -1;
87 }
88 const string modelName = parser.get<String>("@alias");
89 const string zooFile = findFile(parser.get<String>("zoo"));
90
91 keys += genPreprocArguments(modelName, zooFile, "");
92 keys += genPreprocArguments(modelName, zooFile, "ocr_");
93 parser = CommandLineParser(argc, argv, keys);
94 parser.about(about);
95
96 // Parsing command-line arguments
97
98 String sha1 = parser.get<String>("sha1");
99 String ocr_sha1 = parser.get<String>("ocr_sha1");
100 String detModelPath = findModel(parser.get<String>("model"), sha1);
101 String ocr = findModel(parser.get<String>("ocr_model"), ocr_sha1);
102 int height = parser.get<int>("height");
103 int width = parser.get<int>("width");
104 bool imreadRGB = parser.get<bool>("rgb");
105 String vocPath = parser.get<String>("vocabularyPath");
106 float binThresh = parser.get<float>("binaryThreshold");
107 float polyThresh = parser.get<float>("polygonThreshold");
108 double unclipRatio = parser.get<double>("unclipRatio");
109 uint maxCandidates = parser.get<uint>("maxCandidate");
110 float confThreshold = parser.get<float>("thr");
111 float nmsThreshold = parser.get<float>("nms");
112 Scalar mean = parser.get<Scalar>("mean");
113
114 // Ensuring the provided arguments are valid
115 if (!parser.check()) {
116 parser.printErrors();
117 return 1;
118 }
119
120 // Asserting detection model path is provided
121 CV_Assert(!detModelPath.empty());
122
123 vector<vector<Point>> detResults;
124 // Reading the input image
125 Mat frame = imread(samples::findFile(parser.get<String>("input")));
126 Mat board(frame.size(), frame.type(), Scalar(255, 255, 255));
127 int stdSize = 20;
128 int stdWeight = 400;
129 int stdImgSize = 512;
130 int imgWidth = min(frame.rows, frame.cols);
131 int size = (stdSize*imgWidth)/stdImgSize;
132 int weight = (stdWeight*imgWidth)/stdImgSize;
133 FontFace fontFace("sans");
134
135 // Initializing and configuring the text detection model based on the provided config
136 if (modelName == "East") {
137 // EAST Detector initialization
138 TextDetectionModel_EAST detector(detModelPath);
139 detector.setConfidenceThreshold(confThreshold)
140 .setNMSThreshold(nmsThreshold);
141 // Setting input parameters specific to EAST model
142 detector.setInputParams(1.0, Size(width, height), mean, true);
143 // Performing text detection
144 detector.detect(frame, detResults);
145 }
146 else if (modelName == "DB") {
147 // DB Detector initialization
148 TextDetectionModel_DB detector(detModelPath);
149 detector.setBinaryThreshold(binThresh)
150 .setPolygonThreshold(polyThresh)
151 .setUnclipRatio(unclipRatio)
152 .setMaxCandidates(maxCandidates);
153 // Setting input parameters specific to DB model
154 detector.setInputParams(1.0 / 255.0, Size(width, height), mean);
155 // Performing text detection
156 detector.detect(frame, detResults);
157 }
158 else {
159 cout << "[ERROR]: Unsupported file config for the detector model. Valid values: east/db" << endl;
160 return 1;
161 }
162
163 // Reading and storing vocabulary for text recognition
164 CV_Assert(!vocPath.empty());
165 ifstream vocFile;
166 vocFile.open(samples::findFile(vocPath));
167 CV_Assert(vocFile.is_open());
168 std::string vocLine;
169 vector<std::string> vocabulary;
170 while (getline(vocFile, vocLine)) {
171 vocabulary.push_back(vocLine);
172 }
173
174 processFrame(frame, detResults, ocr, imreadRGB, board, fontFace, size, weight, vocabulary);
175 return 0;
176}
177
178// Performs a perspective transform for a four-point region
179static void fourPointsTransform(const Mat& frame, const Point2f vertices[], Mat& result) {
180 const Size outputSize = Size(100, 32);
181 // Defining target vertices for the perspective transform
182 Point2f targetVertices[4] = {
183 Point(0, outputSize.height - 1),
184 Point(0, 0),
185 Point(outputSize.width - 1, 0),
186 Point(outputSize.width - 1, outputSize.height - 1)
187 };
188 // Computing the perspective transform matrix
189 Mat rotationMatrix = getPerspectiveTransform(vertices, targetVertices);
190 // Applying the perspective transform to the region
191 warpPerspective(frame, result, rotationMatrix, outputSize);
192}
193
194void processFrame(
195 const Mat& frame,
196 const vector<vector<Point>>& detResults,
197 const std::string& ocr_model,
198 bool imreadRGB,
199 Mat& board,
200 FontFace& fontFace,
201 int fontSize,
202 int fontWeight,
203 const vector<std::string>& vocabulary
204) {
205 if (detResults.size() > 0) {
206 // Text Recognition
207 Mat recInput;
208 if (!imreadRGB) {
209 cvtColor(frame, recInput, cv::COLOR_BGR2GRAY);
210 } else {
211 recInput = frame;
212 }
213
214 vector<vector<Point>> contours;
215 for (uint i = 0; i < detResults.size(); i++) {
216 const auto& quadrangle = detResults[i];
217 CV_CheckEQ(quadrangle.size(), (size_t)4, "");
218
219 contours.emplace_back(quadrangle);
220
221 vector<Point2f> quadrangle_2f;
222 for (int j = 0; j < 4; j++)
223 quadrangle_2f.emplace_back(detResults[i][j]);
224
225 // Cropping the detected text region using a four-point transform
226 Mat cropped;
227 fourPointsTransform(recInput, &quadrangle_2f[0], cropped);
228
229 if(!ocr_model.empty()){
230 TextRecognitionModel recognizer(ocr_model);
231 recognizer.setVocabulary(vocabulary);
232 recognizer.setDecodeType("CTC-greedy");
233
234 // Setting input parameters for the recognition model
235 double recScale = 1.0 / 127.5;
236 Scalar recMean = Scalar(127.5);
237 Size recInputSize = Size(100, 32);
238 recognizer.setInputParams(recScale, recInputSize, recMean);
239 // Recognizing text from the cropped image
240 string recognitionResult = recognizer.recognize(cropped);
241 cout << i << ": '" << recognitionResult << "'" << endl;
242
243 // Displaying the recognized text on the image
244 putText(board, recognitionResult, Point(detResults[i][1].x, detResults[i][0].y), Scalar(0, 0, 0), fontFace, fontSize, fontWeight);
245 }
246 else{
247 cout << "[WARN] Please pass the path to the ocr model using --ocr_model to get the recognised text." << endl;
248 }
249 }
250 // Drawing detected text regions on the image
251 polylines(board, contours, true, Scalar(200, 255, 200), 1);
252 polylines(frame, contours, true, Scalar(0, 255, 0), 1);
253 } else {
254 cout << "No Text Detected." << endl;
255 }
256
257 // Displaying the final image with detected and recognized text
258 Mat stacked;
259 hconcat(frame, board, stacked);
260 imshow("Text Detection and Recognition", stacked);
261 waitKey(0);
262}
