ytfeng's picture
Limit combinations of backends and targets in demos and benchmark (#145)
a07f7bd
# This file is part of OpenCV Zoo project.
# It is subject to the license terms in the LICENSE file found in the same directory.
#
# Copyright (C) 2021, Shenzhen Institute of Artificial Intelligence and Robotics for Society, all rights reserved.
# Third party copyrights are property of their respective owners.
import numpy as np
import cv2 as cv
class CRNN:
def __init__(self, modelPath, backendId=0, targetId=0):
self._model_path = modelPath
self._backendId = backendId
self._targetId = targetId
self._model = cv.dnn.readNet(self._model_path)
self._model.setPreferableBackend(self._backendId)
self._model.setPreferableTarget(self._targetId)
# load charset by the name of model
if '_EN_' in self._model_path:
self._charset = self._load_charset(self.CHARSET_EN_36)
elif '_CH_' in self._model_path:
self._charset = self._load_charset(self.CHARSET_CH_94)
elif '_CN_' in self._model_path:
self._charset = self._load_charset(self.CHARSET_CN_3944)
else:
print('Charset not supported! Exiting ...')
exit()
self._inputSize = [100, 32] # Fixed
self._targetVertices = np.array([
[0, self._inputSize[1] - 1],
[0, 0],
[self._inputSize[0] - 1, 0],
[self._inputSize[0] - 1, self._inputSize[1] - 1]
], dtype=np.float32)
@property
def name(self):
return self.__class__.__name__
def _load_charset(self, charset):
return ''.join(charset.splitlines())
def setBackendAndTarget(self, backendId, targetId):
self._backendId = backendId
self._targetId = targetId
self._model.setPreferableBackend(self._backendId)
self._model.setPreferableTarget(self._targetId)
def _preprocess(self, image, rbbox):
# Remove conf, reshape and ensure all is np.float32
vertices = rbbox.reshape((4, 2)).astype(np.float32)
rotationMatrix = cv.getPerspectiveTransform(vertices, self._targetVertices)
cropped = cv.warpPerspective(image, rotationMatrix, self._inputSize)
# 'CN' can detect digits (0\~9), upper/lower-case letters (a\~z and A\~Z), and some special characters
# 'CH' can detect digits (0\~9), upper/lower-case letters (a\~z and A\~Z), some Chinese characters and some special characters
if 'CN' in self._model_path or 'CH' in self._model_path:
pass
else:
cropped = cv.cvtColor(cropped, cv.COLOR_BGR2GRAY)
return cv.dnn.blobFromImage(cropped, size=self._inputSize, mean=127.5, scalefactor=1 / 127.5)
def infer(self, image, rbbox):
# Preprocess
inputBlob = self._preprocess(image, rbbox)
# Forward
self._model.setInput(inputBlob)
outputBlob = self._model.forward()
# Postprocess
results = self._postprocess(outputBlob)
return results
def _postprocess(self, outputBlob):
'''Decode charaters from outputBlob
'''
text = ''
for i in range(outputBlob.shape[0]):
c = np.argmax(outputBlob[i][0])
if c != 0:
text += self._charset[c - 1]
else:
text += '-'
# adjacent same letters as well as background text must be removed to get the final output
char_list = []
for i in range(len(text)):
if text[i] != '-' and (not (i > 0 and text[i] == text[i - 1])):
char_list.append(text[i])
return ''.join(char_list)
CHARSET_EN_36 = '''0
1
2
3
4
5
6
7
8
9
a
b
c
d
e
f
g
h
i
j
k
l
m
n
o
p
q
r
s
t
u
v
w
x
y
z'''
CHARSET_CH_94 = '''
0
1
2
3
4
5
6
7
8
9
a
b
c
d
e
f
g
h
i
j
k
l
m
n
o
p
q
r
s
t
u
v
w
x
y
z
A
B
C
D
E
F
G
H
I
J
K
L
M
N
O
P
Q
R
S
T
U
V
W
X
Y
Z
!
"
#
$
%
&
'
(
)
*
+
,
-
.
/
:
;
<
=
>
?
@
[
\
]
^
_
`
{
|
}
~'''
CHARSET_CN_3944 = '''
H
O
K
I
T
E
A
J
N
G
Y
C
U
Q
1
7
3
9
8
线
便
4
0
-
6
5
2
西
l
i
k
n
g
c
o
f
e
w
h
t
L
:
.
B
/
绿
d
D
y
F
r
P
a
u
M
R
S
®
鱿
*
V
X
W
j
Z
·
x
v
s
m
广
p
&
b
'
寿
|
仿
é
[
]
z
怀
ǐ
à
¥
©
ä
°
_
饿
!
鹿
屿
粿
q
滿
(
耀
漿
,
槿
%
<
>
宿
+
#
~
=
)
\
×
ā
¦
@
椿
亿
鸿
Λ
访
綿
Á
ō
驿
­
ǒ
穿
使
殿
尿
廿
Ξ
φ
á
ǎ
а
伿
湿
祿
稿
û
í
ó
Θ
{
π
`
姿
ī
ò
"
ē
退
ε
ě
}
ǔ
è
´
Ē
?
ʌ
É
齿
鴿
ú
ˊ
$
;
^'''