【工程部署】在RK3588上部署OCR(文字检测识别)（DBNet+CRNN）
硬件平台：
1、firefly安装Ubuntu系统的RK3588；
2、安装Windows系统的电脑一台，其上安装Ubuntu18.04系统虚拟机。
参考手册：《00-Rockchip_RKNPU_User_Guide_RKNN_API_V1.3.0_CN》
《RKNN Toolkit Lite2 用户使用指南》
1、文字检测
项目地址：
GitHub - WenmuZhou/PytorchOCR: 基于Pytorch的OCR工具库，支持常用的文字检测和识别算法
DBNet（Dynamic-Link Bi-directional Network）是一种用于文本检测的深度学习模型。该模型于2019年由Minghui Liao等人提出，并在文本检测领域取得了显著的成果。DBNet的设计目标是在保持高精度的同时，提高文本检测的效率。传统的文本检测模型通常使用单向的横向连接或纵向连接来处理文本实例。然而，这种单向连接可能导致信息的不完整传递或信息冗余，影响了检测性能和速度。
为了解决这些问题，DBNet引入了双向动态连接机制，允许横向和纵向两个方向上的信息流动。具体来说，DBNet由两个关键组成部分构成：
(1) Bi-directional FFM（Feature Fusion Module）：这是DBNet的核心组件之一。它包括横向和纵向两个方向的子模块。在横向子模块中，DBNet通过可变形卷积（deformable convolution）从不同尺度的特征图中提取并融合文本实例的特征。而在纵向子模块中，DBNet使用自适应的特征选择机制，动态选择最具有代表性的特征。这些子模块的组合使得文本实例的特征能够全面而高效地进行建模。
(2) Aggregation Decoder：这是DBNet的另一个重要组件，用于从特征图中生成文本实例的边界框和对应的文本分数。该解码器结合了横向和纵向的特征，通过逐步聚合来预测文本的位置和形状。由于使用了双向动态连接，解码器能够更准确地还原文本实例的形态。
DBNet的训练过程包括前向传播和反向传播。在前向传播中，DBNet将图像输入网络，经过一系列卷积、特征融合和解码操作，得到文本检测的结果。然后，通过计算预测结果和真实标签之间的损失函数，使用反向传播算法来更新网络参数，从而不断优化模型的性能。
DBNet在文本检测任务中取得了非常好的效果。其双向动态连接机制允许更好地利用横向和纵向的信息，提高了文本检测的准确性和鲁棒性。此外，相比传统的文本检测模型，DBNet在保持高精度的情况下，大幅提升了检测速度，使得它在实际应用中更具可用性和实用性。因此，DBNet在文字检测、自动化办公、图像识别等领域都具有广泛的应用前景。论文地址：https://arxiv.org/abs/1911.08947
图1. DBNet网络结构
2、文字识别
项目地址：
GitHub - WenmuZhou/PytorchOCR: 基于Pytorch的OCR工具库，支持常用的文字检测和识别算法
CRNN（Convolutional Recurrent Neural Network）是一种深度学习模型，结合了卷积神经网络（CNN）和循环神经网络（RNN）的优势，广泛应用于图像文本识别（OCR）任务。CRNN模型于2015年由Baoguang Shi等人首次提出，并在OCR领域取得了显著的突破。
CRNN的设计思想是将卷积神经网络用于图像的特征提取，并利用循环神经网络来对序列建模，从而使得CRNN能够直接从图像级别到序列级别进行端到端的学习。
CRNN模型通常由以下几个部分组成：
(1) 卷积层（Convolutional Layers）：CRNN利用多个卷积层来提取图像中的局部特征。这些卷积层可以学习不同层次的图像表示，从低级特征（如边缘和纹理）到高级特征（如形状和模式）。
(2) RNN层（Recurrent Layers）：在卷积层后面，CRNN采用RNN层来处理序列数据。RNN能够捕捉序列的上下文信息，因此对于OCR任务而言，它可以有效地处理不同长度的文本序列。
(3) 转录层（Transcription Layer）：在RNN层之后，CRNN使用转录层来将RNN输出映射到字符类别。这通常是一个全连接层，将RNN输出映射到预定义的字符集合，从而实现对文本的识别。
CRNN的训练过程包括两个主要步骤：前向传播和反向传播。在前向传播中，CRNN将图像输入模型，经过卷积和循环层，最终得到文本序列的预测。然后，通过计算预测结果和真实标签之间的损失函数，使用反向传播算法来更新网络参数，从而使得模型的预测结果逐渐接近真实标签。
CRNN在OCR领域的应用广泛，能够识别不同尺寸、字体、颜色和背景的文本。它在识别长文本序列方面表现优秀，并且由于端到端的设计，避免了传统OCR系统中复杂的流水线处理。因此，CRNN在很多实际场景中都取得了很好的效果，如车牌识别、文字检测和手写体识别等。
总结来说，CRNN是一种将CNN和RNN结合起来的深度学习模型，用于图像文本识别任务。其端到端的设计、优秀的序列建模能力和在OCR领域的广泛应用，使得CRNN成为了一种重要的OCR模型，为自动化文本处理和识别带来了巨大的便利。论文地址：https://arxiv.org/abs/1507.05717
图2. CRNN结构
环境搭建
rknn-toolkit以及rknpu_sdk环境搭建
（手把手）rknn-toolkit以及rknpu_sdk环境搭建--以rk3588为例_warren@伟_的博客-CSDN博客
模型的导出与验证
文字检测
导出onnx模型
'''
Author: warren
Date: 2023-06-07 14:52:27
LastEditors: warren
LastEditTime: 2023-06-12 15:20:28
FilePath: /warren/VanillaNet1/export_onnx.py
Description: export onnx model
Copyright (c) 2023 by ${git_name_email}, All Rights Reserved.
'''
#!/usr/bin/env python3
import torch
from torchocr.networks import build_model
MODEL_PATH='./model/det_db_mbv3_new.pth'
DEVICE='cuda:0' if torch.cuda.is_available() else 'cpu'
print("-----------------------devices",DEVICE)
class DetInfer:
def __init__(self, model_path):
ckpt = torch.load(model_path, map_location=DEVICE)
cfg = ckpt['cfg']
self.model = build_model(cfg['model'])
state_dict = {}
for k, v in ckpt['state_dict'].items():
state_dict[k.replace('module.', '')] = v
self.model.load_state_dict(state_dict)
self.device = torch.device(DEVICE)
self.model.to(self.device)
self.model.eval()
checkpoint = torch.load(MODEL_PATH, map_location=DEVICE)
# Prepare input tensor
input = torch.randn(1, 3, 640, 640, requires_grad=False).float().to(torch.device(DEVICE))
# Export the torch model as onnx
print("-------------------export")
torch.onnx.export(self.model,
input,
'detect_model_small.onnx', # name of the exported onnx model
export_params=True,
opset_version=12,
do_constant_folding=False)
# Load the pretrained model and export it as onnx
model = DetInfer(MODEL_PATH)验证
import numpy as np
import cv2
import torch
from torchvision import transforms
# from label_convert import CTCLabelConverter
import cv2
import numpy as np
import pyclipper
from shapely.geometry import Polygon
import onnxruntime
class DBPostProcess():
def __init__(self, thresh=0.3, box_thresh=0.7, max_candidates=1000, unclip_ratio=2):
self.min_size = 3
self.thresh = thresh
self.box_thresh = box_thresh
self.max_candidates = max_candidates
self.unclip_ratio = unclip_ratio
def __call__(self, pred, h_w_list, is_output_polygon=False):
'''
batch: (image, polygons, ignore_tags
h_w_list: 包含[h,w]的数组
pred:
binary: text region segmentation map, with shape (N, 1,H, W)
'''
pred = pred[:, 0, :, :]
segmentation = self.binarize(pred)
boxes_batch = []
scores_batch = []
for batch_index in range(pred.shape[0]):
height, width = h_w_list[batch_index]
boxes, scores = self.post_p(pred[batch_index], segmentation[batch_index], width, height,
is_output_polygon=is_output_polygon)
boxes_batch.append(boxes)
scores_batch.append(scores)
return boxes_batch, scores_batch
def binarize(self, pred):
return pred > self.thresh
def post_p(self, pred, bitmap, dest_width, dest_height, is_output_polygon=False):
'''
_bitmap: single map with shape (H, W),
whose values are binarized as {0, 1}
'''
height, width = pred.shape
boxes = []
new_scores = []
# bitmap = bitmap.cpu().numpy()
if cv2.__version__.startswith('3'):
_, contours, _ = cv2.findContours((bitmap * 255).astype(np.uint8), cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
if cv2.__version__.startswith('4'):
contours, _ = cv2.findContours((bitmap * 255).astype(np.uint8), cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
for contour in contours[:self.max_candidates]:
epsilon = 0.005 * cv2.arcLength(contour, True)
approx = cv2.approxPolyDP(contour, epsilon, True)
points = approx.reshape((-1, 2))
if points.shape[0] < 4:
continue
score = self.box_score_fast(pred, contour.squeeze(1))
if self.box_thresh > score:
continue
if points.shape[0] > 2:
box = self.unclip(points, unclip_ratio=self.unclip_ratio)
if len(box) > 1:
continue
else:
continue
four_point_box, sside = self.get_mini_boxes(box.reshape((-1, 1, 2)))
if sside < self.min_size + 2:
continue
if not isinstance(dest_width, int):
dest_width = dest_width.item()
dest_height = dest_height.item()
if not is_output_polygon:
box = np.array(four_point_box)
else:
box = box.reshape(-1, 2)
box[:, 0] = np.clip(np.round(box[:, 0] / width * dest_width), 0, dest_width)
box[:, 1] = np.clip(np.round(box[:, 1] / height * dest_height), 0, dest_height)
boxes.append(box)
new_scores.append(score)
return boxes, new_scores
def unclip(self, box, unclip_ratio=1.5):
poly = Polygon(box)
distance = poly.area * unclip_ratio / poly.length
offset = pyclipper.PyclipperOffset()
offset.AddPath(box, pyclipper.JT_ROUND, pyclipper.ET_CLOSEDPOLYGON)
expanded = np.array(offset.Execute(distance))
return expanded
def get_mini_boxes(self, contour):
bounding_box = cv2.minAreaRect(contour)
points = sorted(list(cv2.boxPoints(bounding_box)), key=lambda x: x[0])
index_1, index_2, index_3, index_4 = 0, 1, 2, 3
if points[1][1] > points[0][1]:
index_1 = 0
index_4 = 1
else:
index_1 = 1
index_4 = 0
if points[3][1] > points[2][1]:
index_2 = 2
index_3 = 3
else:
index_2 = 3
index_3 = 2
box = [points[index_1], points[index_2], points[index_3], points[index_4]]
return box, min(bounding_box[1])
def box_score_fast(self, bitmap, _box):
# bitmap = bitmap.detach().cpu().numpy()
h, w = bitmap.shape[:2]
box = _box.copy()
xmin = np.clip(np.floor(box[:, 0].min()).astype(np.int), 0, w - 1)
xmax = np.clip(np.ceil(box[:, 0].max()).astype(np.int), 0, w - 1)
ymin = np.clip(np.floor(box[:, 1].min()).astype(np.int), 0, h - 1)
ymax = np.clip(np.ceil(box[:, 1].max()).astype(np.int), 0, h - 1)
mask = np.zeros((ymax - ymin + 1, xmax - xmin + 1), dtype=np.uint8)
box[:, 0] = box[:, 0] - xmin
box[:, 1] = box[:, 1] - ymin
cv2.fillPoly(mask, box.reshape(1, -1, 2).astype(np.int32), 1)
return cv2.mean(bitmap[ymin:ymax + 1, xmin:xmax + 1], mask)[0]
def narrow_224_32(image, expected_size=(224,32)):
ih, iw = image.shape[0:2]
ew, eh = expected_size
# scale = eh / ih
scale = min((eh/ih),(ew/iw))
# scale = eh / max(iw,ih)
nh = int(ih * scale)
nw = int(iw * scale)
image = cv2.resize(image, (nw, nh), interpolation=cv2.INTER_CUBIC)
top = 0
bottom = eh - nh
left = 0
right = ew - nw
new_img = cv2.copyMakeBorder(image, top, bottom, left, right, cv2.BORDER_CONSTANT, value=(114, 114, 114))
return image,new_img
def draw_bbox(img_path, result, color=(0, 0, 255), thickness=2):
import cv2
if isinstance(img_path, str):
img_path = cv2.imread(img_path)
# img_path = cv2.cvtColor(img_path, cv2.COLOR_BGR2RGB)
img_path = img_path.copy()
for point in result:
point = point.astype(int)
cv2.polylines(img_path, [point], True, color, thickness)
return img_path
if __name__ == '__main__':
onnx_model = onnxruntime.InferenceSession("detect_model_small.onnx")
input_name = onnx_model.get_inputs()[0].name
# Set inputs
img = cv2.imread('./pic/6.jpg')
img0 , image= narrow_224_32(img,expected_size=(640,640))
transform_totensor = transforms.ToTensor()
tensor=transform_totensor(image)
tensor_nor=transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
tensor=tensor_nor(tensor)
tensor = np.array(tensor,dtype=np.float32).reshape(1,3,640,640)
post_proess = DBPostProcess()
is_output_polygon = False
#run
outputs = onnx_model.run(None, {input_name:tensor})
#post process
feat_2 = torch.from_numpy(outputs[0])
print(feat_2.size())
box_list, score_list = post_proess(outputs[0], [image.shape[:2]], is_output_polygon=is_output_polygon)
box_list, score_list = box_list[0], score_list[0]
if len(box_list) > 0:
idx = [x.sum() > 0 for x in box_list]
box_list = [box_list[i] for i, v in enumerate(idx) if v]
score_list = [score_list[i] for i, v in enumerate(idx) if v]
else:
box_list, score_list = [], []
print("-----------------box list",box_list)
img = draw_bbox(image, box_list)
img = img[0:img0.shape[0],0:img0.shape[1]]
print("============save pic")
img1=np.array(img,dtype=np.uint8).reshape(640,640,3)
cv2.imwrite("img.jpg",img1)
cv2.waitKey()
文字识别
onnx模型导出
#!/usr/bin/env python3
import os
import sys
import pathlib
# 将 torchocr路径加到python路径里
__dir__ = pathlib.Path(os.path.abspath(__file__))
import numpy as np
sys.path.append(str(__dir__))
sys.path.append(str(__dir__.parent.parent))
import torch
from torchocr.networks import build_model
MODEL_PATH='./model/ch_rec_moblie_crnn_mbv3.pth'
DEVICE='cuda:0' if torch.cuda.is_available() else 'cpu'
print("-----------------------devices",DEVICE)
class RecInfer:
def __init__(self, model_path, batch_size=1):
ckpt = torch.load(model_path, map_location=DEVICE)
cfg = ckpt['cfg']
self.model = build_model(cfg['model'])
state_dict = {}
for k, v in ckpt['state_dict'].items():
state_dict[k.replace('module.', '')] = v
self.model.load_state_dict(state_dict)
self.batch_size = batch_size
self.device = torch.device(DEVICE)
self.model.to(self.device)
self.model.eval()
# Prepare input tensor
input = torch.randn(1, 3, 32, 224, requires_grad=False).float().to(torch.device(DEVICE))
# Export the torch model as onnx
print("-------------------export")
torch.onnx.export(self.model,
input,
'rego_model_small.onnx',
export_params=True,
opset_version=12,
do_constant_folding=False)
# Load the pretrained model and export it as onnx
model = RecInfer(MODEL_PATH)验证
import onnxruntime
import numpy as np
import cv2
import torch
DEVICE='cuda:0' if torch.cuda.is_available() else 'cpu'
IMG_WIDTH=448
ONNX_MODEL='./onnx_model/repvgg_s.onnx'
LABEL_FILE='/root/autodl-tmp/warren/PytorchOCR_OLD/torchocr/datasets/alphabets/dict_text.txt'
#ONNX_MODEL='./onnx_model/rego_model_small.onnx'
#LABEL_FILE='/root/autodl-tmp/warren/PytorchOCR_OLD/torchocr/datasets/alphabets/ppocr_keys_v1.txt'
PIC='./pic/img.jpg'
class CTCLabelConverter(object):
""" Convert between text-label and text-index """
def __init__(self, character):
# character (str): set of the possible characters.
dict_character = []
with open(character, "rb") as fin:
lines = fin.readlines()
for line in lines:
line = line.decode('utf-8').strip("\n").strip("\r\n")
dict_character += list(line)
self.dict = {}
for i, char in enumerate(dict_character):
# NOTE: 0 is reserved for 'blank' token required by CTCLoss
self.dict[char] = i + 1
#TODO replace ‘ ’ with special symbol
self.character = ['[blank]'] + dict_character+[' '] # dummy '[blank]' token for CTCLoss (index 0)
def decode(self, preds, raw=False):
""" convert text-index into text-label. """
preds_idx = preds.argmax(axis=2)
preds_prob = preds.max(axis=2)
result_list = []
for word, prob in zip(preds_idx, preds_prob):
if raw:
result_list.append((''.join([self.character[int(i)] for i in word]), prob))
else:
result = []
conf = []
for i, index in enumerate(word):
if word[i] != 0 and (not (i > 0 and word[i - 1] == word[i])):
result.append(self.character[int(index)])
conf.append(prob[i])
result_list.append((''.join(result), conf))
return result_list
def decode(preds, raw=False):
""" convert text-index into text-label. """
dict_character = []
dict = {}
character=LABEL_FILE
with open(character, "rb") as fin:
lines = fin.readlines()
for line in lines:
line = line.decode('utf-8').strip("\n").strip("\r\n")
dict_character += list(line)
for i, char in enumerate(dict_character):
# NOTE: 0 is reserved for 'blank' token required by CTCLoss
dict[char] = i + 1
#TODO replace ‘ ’ with special symbol
character = ['[blank]'] + dict_character+[' '] # dummy '[blank]' token for CTCLoss (index 0)
preds_idx = preds.argmax(axis=2)
preds_prob = preds.max(axis=2)
result_list = []
for word, prob in zip(preds_idx, preds_prob):
if raw:
result_list.append((''.join([character[int(i)] for i in word]), prob))
else:
result = []
conf = []
for i, index in enumerate(word):
if word[i] != 0 and (not (i > 0 and word[i - 1] == word[i])):
result.append(character[int(index)])
conf.append(prob[i])
result_list.append((''.join(result), conf))
return result_list
def width_pad_img(_img, _target_width, _pad_value=0):
_height, _width, _channels = _img.shape
to_return_img = np.ones([_height, _target_width, _channels], dtype=_img.dtype) * _pad_value
to_return_img[:_height, :_width, :] = _img
return to_return_img
def resize_with_specific_height(_img):
resize_ratio = 32 / _img.shape[0]
return cv2.resize(_img, (0, 0), fx=resize_ratio, fy=resize_ratio, interpolation=cv2.INTER_LINEAR)
def normalize_img(_img):
return (_img.astype(np.float32) / 255 - 0.5) / 0.5
if __name__ == '__main__':
onnx_model = onnxruntime.InferenceSession(ONNX_MODEL)
input_name = onnx_model.get_inputs()[0].name
# Set inputs
imgs = cv2.imread(PIC)
if not isinstance(imgs,list):
imgs = [imgs]
imgs = [normalize_img(resize_with_specific_height(img)) for img in imgs]
widths = np.array([img.shape[1] for img in imgs])
idxs = np.argsort(widths)
txts = []
label_convert=CTCLabelConverter(LABEL_FILE)
for idx in range(len(imgs)):
batch_idxs = idxs[idx:min(len(imgs), idx+1)]
batch_imgs = [width_pad_img(imgs[idx],IMG_WIDTH) for idx in batch_idxs]
batch_imgs = np.stack(batch_imgs)
print(batch_imgs.shape)
tensor =batch_imgs.transpose([0,3, 1, 2]).astype(np.float32)
out = onnx_model.run(None, {input_name:tensor})
tensor_out = torch.tensor(out)
tensor_out = torch.squeeze(tensor_out,dim=1)
softmax_output = tensor_out.softmax(dim=2)
print("---------------out shape is",softmax_output.shape)
txts.extend([label_convert.decode(np.expand_dims(txt, 0)) for txt in softmax_output])
idxs = np.argsort(idxs)
out_txts = [txts[idx] for idx in idxs]
import sys
import codecs
sys.stdout = codecs.getwriter("utf-8")(sys.stdout.detach())
print(out_txts)
至此 导出验证成功
rk3588板端部署
转化为rknn模型
from rknn.api import RKNN
ONNX_MODEL = 'xxx.onnx'
RKNN_MODEL = 'xxxx.rknn'
DATASET = './dataset.txt'
if __name__ == '__main__':
# Create RKNN object
rknn = RKNN(verbose=True)
# pre-process config
print('--> Config model')
ret=rknn.config(mean_values=[[0, 0, 0]], std_values=[[0, 0, 0]],target_platform='rk3588') #wzw
if ret != 0:
print('config model failed!')
exit(ret)
print('done')
# Load ONNX model
print('--> Loading model')
ret = rknn.load_onnx(model=ONNX_MODEL, outputs=['output', '345', '346'])
if ret != 0:
print('Load model failed!')
exit(ret)
print('done')
# Build model
print('--> Building model')
ret = rknn.build(do_quantization=True, dataset=DATASET)
#ret = rknn.build(do_quantization=False)
if ret != 0:
print('Build model failed!')
exit(ret)
print('done')
# Export RKNN model
print('--> Export rknn model')
ret = rknn.export_rknn(RKNN_MODEL)
if ret != 0:
print('Export rknn model failed!')
exit(ret)
print('done')
#release rknn
rknn.release()使用pyqt进行开发
PyQt软件设计
使用pyqt进行开发，ui界面如图所示
UI
图6. 基于PYQT的ui界面
该界面包含了三个功能按钮，其中包裹一个选择静态图片，一个使用相机，一个检测按钮，TextEdit用于显示识别结果，label用于显示处理完成后的图片。
软件流程图如下：
总体目录参照
下面依次介绍图片检测的相关代码：
import platform
import sys
import cv2
import numpy as np
import torch
import pyclipper
from shapely.geometry import Polygon
from torchvision import transforms
import time
import os
import glob
import threading
from PyQt5.QtGui import *
from PyQt5.QtWidgets import *
from PyQt5.QtCore import *
import platform
from rknnlite.api import RKNNLite
import os
os.environ.pop("QT_QPA_PLATFORM_PLUGIN_PATH")
DETECT_MODEL = './model/model_small.rknn'
REGO_MODEL='./model/repvgg_s.rknn'
LABEL_FILE='./dict/dict_text.txt'
LABEL_SIZE_PRIVIOUS=0
LABEL_SIZE_LATTER=0
# 文件夹路径
folder_path = './crop_pic'
# 使用 glob 来获取所有图片文件的路径
image_files = glob.glob(os.path.join(folder_path, '*.png')) + glob.glob(os.path.join(folder_path, '*.jpg'))
def resize_img_self(image,reszie_size=(0,0)):
ih,iw=image.shape[0:2]
ew,eh=reszie_size
scale=eh/ih
width=int(iw*scale)
height=int(ih*scale)
if height!=eh:
height=eh
image=cv2.resize(image,(width,height),interpolation=cv2.INTER_LINEAR)
top = 0
bottom = 0
left = 0
right = ew-width
new_img = cv2.copyMakeBorder(image, top, bottom, left, right, cv2.BORDER_CONSTANT, value=(114, 114, 114))
#print("new image shape",new_img.shape)
return new_img
def narrow_224_32(image, expected_size=(224,32)):
ih, iw = image.shape[0:2]
ew, eh = expected_size
# scale = eh / ih
scale = min((eh/ih),(ew/iw))
# scale = eh / max(iw,ih)
nh = int(ih * scale)
nw = int(iw * scale)
image = cv2.resize(image, (nw, nh), interpolation=cv2.INTER_CUBIC)
top = 0
bottom = eh - nh
left = 0
right = ew - nw
new_img = cv2.copyMakeBorder(image, top, bottom, left, right, cv2.BORDER_CONSTANT, value=(114, 114, 114))
return image,new_img
def draw_bbox(img_path, result, color=(0, 0, 255), thickness=2):
import cv2
if isinstance(img_path, str):
img_path = cv2.imread(img_path)
# img_path = cv2.cvtColor(img_path, cv2.COLOR_BGR2RGB)
img_path = img_path.copy()
for point in result:
point = point.astype(int)
cv2.polylines(img_path, [point], True, color, thickness)
return img_path
def delay_milliseconds(milliseconds):
seconds = milliseconds / 1000.0
time.sleep(seconds)
""" Convert between text-label and text-index """
class CTCLabelConverter(object):
def __init__(self, character):
# character (str): set of the possible characters.
dict_character = []
with open(character, "rb") as fin:
lines = fin.readlines()
for line in lines:
line = line.decode('utf-8').strip("\n").strip("\r\n")
dict_character += list(line)
self.dict = {}
for i, char in enumerate(dict_character):
# NOTE: 0 is reserved for 'blank' token required by CTCLoss
self.dict[char] = i + 1
#TODO replace ‘ ’ with special symbol
self.character = ['[blank]'] + dict_character+[' '] # dummy '[blank]' token for CTCLoss (index 0)
def decode(self, preds, raw=False):
""" convert text-index into text-label. """
preds_idx = preds.argmax(axis=2)
preds_prob = preds.max(axis=2)
result_list = []
for word, prob in zip(preds_idx, preds_prob):
if raw:
result_list.append((''.join([self.character[int(i)] for i in word]), prob))
else:
result = []
conf = []
for i, index in enumerate(word):
if word[i] != 0 and (not (i > 0 and word[i - 1] == word[i])):
result.append(self.character[int(index)])
#conf.append(prob[i])
#result_list.append((''.join(result), conf))
result_list.append((''.join(result)))
return result_list
class DBPostProcess():
def __init__(self, thresh=0.3, box_thresh=0.7, max_candidates=1000, unclip_ratio=2):
self.min_size = 3
self.thresh = thresh
self.box_thresh = box_thresh
self.max_candidates = max_candidates
self.unclip_ratio = unclip_ratio
def __call__(self, pred, h_w_list, is_output_polygon=False):
pred = pred[:, 0, :, :]
segmentation = self.binarize(pred)
boxes_batch = []
scores_batch = []
for batch_index in range(pred.shape[0]):
height, width = h_w_list[batch_index]
boxes, scores = self.post_p(pred[batch_index], segmentation[batch_index], width, height,
is_output_polygon=is_output_polygon)
boxes_batch.append(boxes)
scores_batch.append(scores)
return boxes_batch, scores_batch
def binarize(self, pred):
return pred > self.thresh
def post_p(self, pred, bitmap, dest_width, dest_height, is_output_polygon=False):
'''
_bitmap: single map with shape (H, W),
whose values are binarized as {0, 1}
'''
height, width = pred.shape
boxes = []
new_scores = []
# bitmap = bitmap.cpu().numpy()
if cv2.__version__.startswith('3'):
_, contours, _ = cv2.findContours((bitmap * 255).astype(np.uint8), cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
if cv2.__version__.startswith('4'):
contours, _ = cv2.findContours((bitmap * 255).astype(np.uint8), cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
for contour in contours[:self.max_candidates]:
epsilon = 0.005 * cv2.arcLength(contour, True)
approx = cv2.approxPolyDP(contour, epsilon, True)
points = approx.reshape((-1, 2))
if points.shape[0] < 4:
continue
score = self.box_score_fast(pred, contour.squeeze(1))
if self.box_thresh > score:
continue
if points.shape[0] > 2:
box = self.unclip(points, unclip_ratio=self.unclip_ratio)
if len(box) > 1:
continue
else:
continue
four_point_box, sside = self.get_mini_boxes(box.reshape((-1, 1, 2)))
if sside < self.min_size + 2:
continue
if not isinstance(dest_width, int):
dest_width = dest_width.item()
dest_height = dest_height.item()
if not is_output_polygon:
box = np.array(four_point_box)
else:
box = box.reshape(-1, 2)
box[:, 0] = np.clip(np.round(box[:, 0] / width * dest_width), 0, dest_width)
box[:, 1] = np.clip(np.round(box[:, 1] / height * dest_height), 0, dest_height)
boxes.append(box)
new_scores.append(score)
return boxes, new_scores
def unclip(self, box, unclip_ratio=1.5):
poly = Polygon(box)
distance = poly.area * unclip_ratio / poly.length
offset = pyclipper.PyclipperOffset()
offset.AddPath(box, pyclipper.JT_ROUND, pyclipper.ET_CLOSEDPOLYGON)
expanded = np.array(offset.Execute(distance))
return expanded
def get_mini_boxes(self, contour):
bounding_box = cv2.minAreaRect(contour)
points = sorted(list(cv2.boxPoints(bounding_box)), key=lambda x: x[0])
index_1, index_2, index_3, index_4 = 0, 1, 2, 3
if points[1][1] > points[0][1]:
index_1 = 0
index_4 = 1
else:
index_1 = 1
index_4 = 0
if points[3][1] > points[2][1]:
index_2 = 2
index_3 = 3
else:
index_2 = 3
index_3 = 2
box = [points[index_1], points[index_2], points[index_3], points[index_4]]
return box, min(bounding_box[1])
def box_score_fast(self, bitmap, _box):
# bitmap = bitmap.detach().cpu().numpy()
h, w = bitmap.shape[:2]
box = _box.copy()
xmin = np.clip(np.floor(box[:, 0].min()).astype(np.int), 0, w - 1)
xmax = np.clip(np.ceil(box[:, 0].max()).astype(np.int), 0, w - 1)
ymin = np.clip(np.floor(box[:, 1].min()).astype(np.int), 0, h - 1)
ymax = np.clip(np.ceil(box[:, 1].max()).astype(np.int), 0, h - 1)
mask = np.zeros((ymax - ymin + 1, xmax - xmin + 1), dtype=np.uint8)
box[:, 0] = box[:, 0] - xmin
box[:, 1] = box[:, 1] - ymin
cv2.fillPoly(mask, box.reshape(1, -1, 2).astype(np.int32), 1)
return cv2.mean(bitmap[ymin:ymax + 1, xmin:xmax + 1], mask)[0]
class Process_Class(QWidget):
detect_end = pyqtSignal(str)
clear_text = pyqtSignal()
def __init__(self):
super().__init__()
self.image = None
self.img=None
self.camera_status=False
self.result_string=None
self.cap = cv2.VideoCapture()
#detect
rknn_model_detect = DETECT_MODEL
self.rknn_lite_detect = RKNNLite()
self.rknn_lite_detect.load_rknn(rknn_model_detect)# load RKNN model
self.rknn_lite_detect.init_runtime(core_mask=RKNNLite.NPU_CORE_2)# init runtime environment
#rego
rknn_model_rego = REGO_MODEL
self.rknn_lite_rego = RKNNLite()
self.rknn_lite_rego.load_rknn(rknn_model_rego)# load RKNN model
self.rknn_lite_rego.init_runtime(core_mask=RKNNLite.NPU_CORE_0_1)# init runtime environment
self.detect_end.connect(self.update_text_box)
self.clear_text.connect(self.clear_text_box)
def cv2_to_qpixmap(self, cv_image):
height, width, channel = cv_image.shape
bytes_per_line = 3 * width
q_image = QImage(cv_image.data, width, height, bytes_per_line, QImage.Format_RGB888).rgbSwapped()
return QPixmap.fromImage(q_image)
def show_pic(self, cv_image):
pixmap = self.cv2_to_qpixmap(cv_image)
if MainWindow.pic_label is not None:
MainWindow.pic_label.setPixmap(pixmap)
QApplication.processEvents()
else:
print("wrong!!!!!!!")
def camera_open(self):
self.camera_status = not self.camera_status
print("------------camera status is",self.camera_status)
if self.camera_status:
self.cap.open(12)
if self.cap.isOpened():
print("run camera")
while(True):
frame = self.cap.read()
if not frame[0]:
print("read frame failed!!!!")
exit()
self.image=frame[1]
self.detect_pic()
if not self.camera_status:
break
else:
print("Cannot open camera")
exit()
else:
self.release_camera()
def release_camera(self):
if self.cap.isOpened():
self.cap.release()
self.camera_status = False
print("摄像头关闭")
def open_file(self):
# 获取图像的路径
img_path, _ = QFileDialog.getOpenFileName()
if img_path != '':
self.image = cv2.imread(img_path)
self.show_pic(self.image)
def crop_and_save_image(self,image, box_points):
global LABEL_SIZE_PRIVIOUS
global LABEL_SIZE_LATTER
i=-1
# 将box_points转换为NumPy数组，并取整数值
box_points = np.array(box_points, dtype=np.int32)
mask = np.zeros_like(image) # 创建与图像相同大小的全黑图像
print("LABEL_SIZE_PRIVIOUS ",LABEL_SIZE_PRIVIOUS,"LABEL_SIZE_LATTER ",LABEL_SIZE_LATTER)
if LABEL_SIZE_PRIVIOUS==LABEL_SIZE_LATTER:
LABEL_SIZE_PRIVIOUS=len(box_points)
for box_point in box_points:
i=i+1
cropped_image = image.copy()
# 使用OpenCV的函数裁剪图像
x, y, w, h = cv2.boundingRect(box_point)
cropped_image = image[y:y+h, x:x+w]
# 创建与图像大小相同的全黑掩码
mask = np.zeros_like(cropped_image)
# 在掩码上绘制多边形
cv2.fillPoly(mask, [box_point - (x, y)], (255, 255, 255))
# 使用 bitwise_and 进行图像裁剪
masked_cropped_image = cv2.bitwise_and(cropped_image, mask)
# 保存裁剪后的图像
output_path = f"{'./crop_pic/'}img_{i}.jpg"
cv2.imwrite(output_path, masked_cropped_image)
else:
#self.clear_text.emit()
LABEL_SIZE_LATTER=LABEL_SIZE_PRIVIOUS
current_directory = os.getcwd()+'/crop_pic' # Get the current directory
for filename in os.listdir(current_directory):
if filename.endswith(".jpg"):
file_path = os.path.join(current_directory, filename)
os.remove(file_path)
print(f"Deleted: {file_path}")
def detect_thread(self):
#detect inference
img0 , image= narrow_224_32(self.image,expected_size=(640,640))
outputs =self.rknn_lite_detect.inference(inputs=[image])
post_proess = DBPostProcess()
is_output_polygon = False
box_list, score_list = post_proess(outputs[0], [image.shape[:2]], is_output_polygon=is_output_polygon)
box_list, score_list = box_list[0], score_list[0]
if len(box_list) > 0:
idx = [x.sum() > 0 for x in box_list]
box_list = [box_list[i] for i, v in enumerate(idx) if v]
score_list = [score_list[i] for i, v in enumerate(idx) if v]
else:
box_list, score_list = [], []
self.image = draw_bbox(image, box_list)
self.crop_and_save_image(image,box_list)
self.image = self.image[0:img0.shape[0],0:img0.shape[1]]
self.show_pic(self.image)
def rego_thread(self):
label_convert=CTCLabelConverter(LABEL_FILE)
self.clear_text.emit()
for image_file in image_files:
if os.path.exists(image_file):
print('-----------image file',image_file,len(image_files))
self.img = cv2.imread(image_file)
image = resize_img_self(self.img,reszie_size=(448,32))
# Inference
outputs = self.rknn_lite_rego.inference(inputs=[image])
#post process
feat_2 = torch.tensor(outputs[0],dtype=torch.float32)
txt = label_convert.decode(feat_2.detach().numpy())
self.result_string = ' '.join(txt)
print(self.result_string)
self.detect_end.emit(self.result_string)
else:
print("-----------no crop image!!!")
def detect_pic(self):
self.detect_thread()
my_thread = threading.Thread(target=self.rego_thread)
# 启动线程
my_thread.start()
# 等待线程结束
my_thread.join()
def update_text_box(self, text):
# 在主线程中更新文本框的内容
MainWindow.text_box.append(text)
def clear_text_box(self):
print("clear--------------------------------")
# 在主线程中更新文本框的内容
MainWindow.text_box.clear()
class MainWindow(QMainWindow):
#pic_label = None
def __init__(self):
pic_label = None
text_box = None
super().__init__()
self.process_functions = Process_Class()
self.window = QWidget()
# 创建小部件
self.pic_label = QLabel('Show Window!', parent=self.window)
self.pic_label.setMinimumHeight(500) # 设置最小高度
self.pic_label.setMaximumHeight(500) # 设置最大高度
self.pic_button = QPushButton('Picture', parent=self.window)
self.pic_button.clicked.connect(self.process_functions.open_file)
self.camera_button = QPushButton('Camera', parent=self.window)
self.camera_button.clicked.connect(self.process_functions.camera_open)
self.detect_button = QPushButton('Detect', parent=self.window)
self.detect_button.clicked.connect(self.process_functions.detect_pic)
self.text_box = QTextEdit()
# 创建垂直布局管理器并将小部件添加到布局中
self.left_layout = QVBoxLayout()
self.right_layout = QVBoxLayout()
self.layout = QHBoxLayout()
self.create_ui()
self.window.closeEvent = self.closeEvent
def create_ui(self):
self.window.setWindowTitle('Scene_text_rego')
self.window.setGeometry(0, 0, 800, 600) # 设置窗口位置和大小
# 设置主窗口的布局
self.pic_label.setStyleSheet('border: 2px solid black; padding: 10px;')
self.left_layout.addWidget(self.pic_label)
self.left_layout.addWidget(self.text_box)
self.right_layout.addWidget(self.pic_button)
self.right_layout.addWidget(self.camera_button)
self.right_layout.addWidget(self.detect_button)
self.layout.addLayout(self.left_layout)
self.layout.addLayout(self.right_layout)
self.window.setLayout(self.layout)
self.window.show()
def closeEvent(self, event):
# 释放摄像头资源
self.process_functions.release_camera()
event.accept()
def main():
# 创建应用程序对象
app = QApplication(sys.argv)
win = MainWindow()
MainWindow.pic_label = win.pic_label # 设置类变量pic_label为MainWindow对象的pic_label
MainWindow.text_box = win.text_box # 设置类变量pic_label为MainWindow对象的pic_label
# 运行应用程序
sys.exit(app.exec_())
rknn_lite_detect.release()
if __name__ == '__main__':
main()运行结果
参考资料
博文：
【工程部署】手把手教你在RKNN上部署OCR服务（上）_rknn ocr_三叔家的猫的博客-CSDN博客