forked from justadudewhohacks/opencv4nodejs
-
Notifications
You must be signed in to change notification settings - Fork 0
/
EASTTextDetection.js
107 lines (86 loc) · 3.01 KB
/
EASTTextDetection.js
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
const path = require('path');
const fs = require('fs');
const { cv, drawBlueRect } = require('./utils');
const { extractResults } = require('./dnn/ssdUtils');
if (!cv.xmodules.dnn) {
throw new Error('exiting: opencv4nodejs compiled without dnn module');
}
const modelPath = path.resolve(__dirname,
'../data/text-models/frozen_east_text_detection.pb');
const imgPath = path.resolve(__dirname, '../data/text-data/detection.png');
if (!fs.existsSync(modelPath)) {
console.log('could not find EAST model');
console.log('download the model from: https://github.com/oyyd/frozen_east_text_detection.pb/blob/71415464412c55bb1d135fcdeda498e29a67effa/frozen_east_text_detection.pb?raw=true'
+ ' or create a .pb model from https://github.com/argman/EAST');
throw new Error('exiting: could not find EAST model');
}
const MIN_CONFIDENCE = 0.5;
const NMS_THRESHOLD = 0.4;
const SIZE = 320;
function decode(scores, geometry, confThreshold) {
const [numRows, numCols] = scores.sizes.slice(2);
const boxes = [];
const confidences = [];
for (let y = 0; y < numRows; y += 1) {
for (let x = 0; x < numCols; x += 1) {
const score = scores.at([0, 0, y, x]);
if (score < MIN_CONFIDENCE) {
continue;
}
const offsetX = x * 4;
const offsetY = y * 4;
const angle = geometry.at([0, 4, y, x]);
const cos = Math.cos(angle);
const sin = Math.sin(angle);
const h = geometry.at([0, 0, y, x]) + geometry.at([0, 2, y, x]);
const w = geometry.at([0, 1, y, x]) + geometry.at([0, 3, y, x]);
const endX = offsetX + (cos * geometry.at([0, 1, y, x])) + (sin * geometry.at([0, 2, y, x]));
const endY = offsetY - (sin * geometry.at([0, 1, y, x])) + (cos * geometry.at([0, 2, y, x]));
const startX = endX - w;
const startY = endY - h;
boxes.push(new cv.Rect(
startX,
startY,
endX - startX,
endY - startY,
));
confidences.push(score);
}
}
return [boxes, confidences];
}
function detection(modelAbsPath, imgAbsPath) {
const net = cv.readNetFromTensorflow(modelPath);
const img = cv.imread(imgAbsPath);
const [imgHeight, imgWidth] = img.sizes;
const widthRatio = imgWidth / SIZE;
const heightRatio = imgHeight / SIZE;
const inputBlob = cv.blobFromImage(img, 1,
new cv.Size(SIZE, SIZE), new cv.Vec3(123.68, 116.78, 103.94), true, false);
net.setInput(inputBlob);
const outBlobNames = [
'feature_fusion/Conv_7/Sigmoid',
'feature_fusion/concat_3',
];
const [scores, geometry] = net.forward(outBlobNames);
const [boxes, confidences] = decode(scores, geometry, MIN_CONFIDENCE);
const indices = cv.NMSBoxes(
boxes,
confidences, MIN_CONFIDENCE, NMS_THRESHOLD
);
indices.forEach((i) => {
const rect = boxes[i];
const imgRect = new cv.Rect(
rect.x * widthRatio,
rect.y * heightRatio,
rect.width * widthRatio,
rect.height * heightRatio,
)
drawBlueRect(img, imgRect);
});
cv.imshowWait('EAST text detection', img);
}
detection(
modelPath,
imgPath
);