def crnnOcr(image):
"""
crnn模型，ocr识别
image:PIL.Image.convert("L")
"""
scale = image.size[1]*1.0 / 32
w = image.size[0] / scale
w = int(w)
transformer = resizeNormalize((w, 32))
image = transformer(image)
image = image.astype(np.float32)
image = torch.from_numpy(image)
if torch.cuda.is_available() and GPU:
image = image.cuda()
else:
image = image.cpu()
image = image.view(1,1, *image.size())
image = Variable(image)
preds = model(image)
_, preds = preds.max(2)
preds = preds.transpose(1, 0).contiguous().view(-1)
sim_pred = converter.decode(preds)
return sim_pred
def predict(self, image):
image = resizeNormalize(image, 32)
image = image.astype(np.float32)
image = np.array([[image]])
self.model.setInput(image)
preds = self.model.forward()
preds = preds.transpose(0, 2, 3, 1)
preds = preds[0]
preds = np.argmax(preds, axis=2).reshape((-1, ))
raw = strLabelConverter(preds, self.alphabet)
return raw
def predict(self, image):
image = resizeNormalize(image, 32)
image = image.astype(np.float32)
image = np.array([[image]])
global graph
with graph.as_default():
preds = self.model.predict(image)
# preds = preds[0]
preds = np.argmax(preds, axis=2).reshape((-1, ))
raw = strLabelConverter(preds, self.alphabet)
return raw
def predict(self, image):
image = resizeNormalize(image, 32)
image = image.astype(np.float32)
image = torch.from_numpy(image)
if torch.cuda.is_available() and self.GPU:
image = image.cuda()
else:
image = image.cpu()
image = image.view(1, 1, *image.size())
image = Variable(image)
preds = self(image)
_, preds = preds.max(2)
preds = preds.transpose(1, 0).contiguous().view(-1)
raw = strLabelConverter(preds, self.alphabet)
return raw
def predict_batch(self, boxes, batch_size=1):
"""
predict on batch
"""
N = len(boxes)
res = []
imgW = 0
batch = N // batch_size
if batch * batch_size != N:
batch += 1
for i in range(batch):
tmpBoxes = boxes[i * batch_size:(i + 1) * batch_size]
imageBatch = []
imgW = 0
for box in tmpBoxes:
img = box['img']
image = resizeNormalize(img, 32)
h, w = image.shape[:2]
imgW = max(imgW, w)
imageBatch.append(np.array([image]))
imageArray = np.zeros((len(imageBatch), 1, 32, imgW),
dtype=np.float32)
n = len(imageArray)
for j in range(n):
_, h, w = imageBatch[j].shape
imageArray[j][:, :, :w] = imageBatch[j]
image = torch.from_numpy(imageArray)
image = Variable(image)
if torch.cuda.is_available() and self.GPU:
image = image.cuda()
else:
image = image.cpu()
preds = self(image)
preds = preds.argmax(2)
n = preds.shape[1]
for j in range(n):
res.append(strLabelConverter(preds[:, j], self.alphabet))
for i in range(N):
boxes[i]['text'] = res[i]
return boxes
def crnnOcr(image):
"""
crnn模型，ocr识别
image:PIL.Image.convert("L")
"""
scale = image.size[1]*1.0 / 32
w = image.size[0] / scale
w = int(w)
transformer = resizeNormalize((w, 32))
image = transformer(image)
image = image.astype(np.float32)
image = np.array([[image]])
global graph
with graph.as_default():
preds = model.predict(image)
preds = preds[0]
preds = np.argmax(preds,axis=2).reshape((-1,))
sim_pred = converter.decode(preds)
return sim_pred
def predict_batch(self, boxes, batch_size=1):
"""
predict on batch
"""
N = len(boxes)
res = []
imgW = 0
batch = N // batch_size
if batch * batch_size != N:
batch += 1
for i in range(batch):
tmpBoxes = boxes[i * batch_size:(i + 1) * batch_size]
imageBatch = []
imgW = 0
for box in tmpBoxes:
img = box['img']
image = resizeNormalize(img, 32)
h, w = image.shape[:2]
imgW = max(imgW, w)
imageBatch.append(np.array([image]))
imageArray = np.zeros((len(imageBatch), 1, 32, imgW),
dtype=np.float32)
n = len(imageArray)
for j in range(n):
_, h, w = imageBatch[j].shape
imageArray[j][:, :, :w] = imageBatch[j]
global graph
with graph.as_default():
preds = self.model.predict(imageArray, batch_size=batch_size)
preds = preds.argmax(axis=2)
n = preds.shape[0]
for j in range(n):
res.append(
strLabelConverter(preds[j, ].tolist(), self.alphabet))
for i in range(N):
boxes[i]['text'] = res[i]
return boxes
