PaddleOCR文字识别
·
make_data.py (生成train.txt文档)
# coding=utf-8
import os
import cv2
import numpy as np
import uuid
import json
import shutil
def read_json(path, encoding="gbk"):
with open(path, "r", encoding=encoding) as _fp:
content = _fp.read()
if len(content) < 3:
return ""
return json.loads(content)
# path = '/home/data1/data_yy/行驶证/身份证error'
# out = '/home/data1/data_yy/行驶证/身份证'
# for root, dirs, files in os.walk(out):
# for file in files:
# if file.endswith('.jpg'):
# saveimg = os.path.join(root, file)
# savelabel = saveimg.replace(".jpg", ".txt")
# imgpath = os.path.join(path, file)
# labelpath = imgpath.replace(".jpg", ".txt")
# if not os.path.exists(labelpath):
# continue
# shutil.copy(imgpath, saveimg)
# shutil.copy(labelpath, savelabel)
# print(file)
# input = '/home/data1/data_yy/行驶证/train.txt'
# out = '/home/data1/data_yy/行驶证/train1.txt'
# num = 0
# with open(input, "r", encoding="utf-8") as fp:
# lines = fp.readlines()
# for line in lines:
# imgpath, label = line.strip('\n').split('\t')
# # if './自制数据/' in imgpath or 'syn_mingpai_black' in imgpath or 'syn_mingpai_white' in imgpath or 'vindata_7w' in imgpath or 'vindata' in imgpath:
# # continue
# # if '身份证' in imgpath or '中文街景文字识别' in imgpath or '表单' in imgpath or 'RCTW' in imgpath or 'ICDAR2019-LSVT' in imgpath or 'ICDAR2019-ART' in imgpath:
# # continue
# if '12L_img_error' in imgpath:
# continue
# # fname= imgpath.split('/')[-1]
# num += 1
# # imgpath = '/home/data1/data_yy/mingpaidata'+ imgpath[1:]
# # saveimg = '/home/data1/data_yy/mingpaidata/test/'+ fname
# # shutil.copy(imgpath, saveimg)
# # txtpath = '/home/data1/data_yy/mingpaidata/test/'+ fname.replace('.jpg', '.txt')
# # with open(txtpath, 'w', encoding="utf-8")as fw:
# # fw.write(label)
# # print(label)
# with open(out, 'a', encoding="utf-8")as fw:
# fw.write(line)
# print(line)
# print(num)
# print("done.")
# input = '/home/data1/data_yy/行驶证/通用文字/train4.txt'
# out = '/home/data1/data_yy/行驶证/通用文字/train.txt'
# num = 0
# with open(input, "r", encoding="utf-8") as fr:
# lines = fr.readlines()
# for line in lines:
# # line = line.replace('./身份证/公民身份号码/公民身份号码/', './身份证/公民身份号码/')
# # if './通用文字/HanaMinB_datacrop/' in line:
# # continue
# imgpath=line.split('\t')[0].replace('./','/home/data1/data_yy/行驶证/')
# if not os.path.exists(imgpath):
# continue
# with open(out, "a", encoding="utf-8") as fw:
# fw.write(line)
# print(line)
# num += 1
#
# print(num)
# input = '/home/data1/data_yy/mingpaidata/nameplate_rec/QDCY36F0105_all_good.txt'
# out = '/home/data1/data_yy/mingpaidata/test1.txt'
# num = 0
# i = 1
# with open(input, "r", encoding="utf-8") as fr:
# lines = fr.readlines()
# for line in lines:
# imgpath, label = line.split('|')
# line = line.replace('/home/data1/dzk/', './').replace('|', '\t')
# with open(out, "a", encoding="utf-8") as fw:
# fw.write(line)
# print(line)
# num += 1
# print(num)
# path1 = '/home/gonglei/yuyang/project/PaddleOCR-release-2.6/ppocr/utils/tongyongzi.txt'
# list = []
# with open(path1, 'r', encoding='utf8')as f:
# list = f.read().split('\n')
# input = '/home/gonglei/yuyang/project/PaddleOCR-release-2.6/ppocr/utils/生僻字.txt'
# with open(input, "r", encoding="utf-8") as fr:
# lines = fr.read()
# n = len(lines)
# for i in range(0, n-1):
# if lines[i] not in list:
# list.append(lines[i])
# print(lines[i])
# out = '/home/gonglei/yuyang/project/PaddleOCR-release-2.6/ppocr/utils/tongyongzi2.txt'
# with open(out, 'w', encoding='utf8')as f:
# content = '\n'.join(list)
# f.write(content)
# list = []
# with open('/home/data1/data_yy/行驶证/train_error.txt', 'r', encoding='utf8')as fr:
# lines = fr.readlines()
# for line in lines:
# imgpath, label = line.split('\t')
# list.append(imgpath.split('/')[-1])
#
# with open('ppocr/utils/文字.txt', 'r')as fr:
# content = fr.read()
# dics = content.split('\n')
#
# tianjia = []
# input = '/home/data1/data_yy/行驶证/通用文字/train.txt'
# out = '/home/data1/data_yy/行驶证/通用文字/train1.txt'
# num = 0
# with open(input, "r", encoding="utf-8") as fr:
# lines = fr.readlines()
# for line in lines:
# if './通用文字/HanaMinB_datacrop/' in line:
# continue
# with open(out, "a", encoding="utf-8") as fw:
# fw.write(line)
# num += 1
# print(num)
# datadir = '/home/data1/data_yy/行驶证/通用文字/题目区域'
# out = '/home/data1/data_yy/行驶证/通用文字/题目区域.txt'
#
# list = []
# num = 0
# for root, dirs, files in os.walk(datadir):
# for file in files:
# if file.split('.')[-1] == 'jpg':
# dir1 = root.split('/')[-1]
# dir2 = root.split('/')[-2]
# img_path = os.path.join(root, file)
# # label = file.split('.jpg')[0]
# label_path = img_path.replace(".jpg", ".txt")
# if not os.path.exists(label_path):
# continue
# try:
# with open(label_path, 'r', encoding='utf-8')as fr:
# line = fr.readline().rstrip('\n')
# except:
# with open(label_path, 'r', encoding='gbk')as fr:
# line = fr.readline().rstrip('\n')
# # if len(line.split('|')) == 2:
# # label = line.split('|')[-1]
# # else:
# # box = line.split('|')[0]
# # label = line.split(box)[-1][1:]
# # if label=='-1' or label == '' or label == ' ':
# # continue
#
# label = line
# # line = './身份证/身份证号码错误_crop/'+str(dir2) +'/'+str(dir1)+'/'+file+'\t'+label+'\n'
# line = '/home/data1/data_yy/行驶证/通用文字/题目区域/' + file + '\t' + label + '\n'
# print(line)
# list.append(line)
# num += 1
#
# np.random.shuffle(list)
# print(len(list))
# content = "".join(list)
# with open(out, "a", encoding="utf-8") as fp:
# fp.write(content)
# print(num)
# print("done.")
datadir = '/home/data1/data_yy/行驶证/通用文字/通用生成数据datacrop'
out = '/home/data1/data_yy/行驶证/通用文字/train2.txt'
list = []
num = 0
for root, dirs, files in os.walk(datadir):
for file in files:
if file.split('.')[-1] == 'jpg':
dir1 = root.split('/')[-1]
dir2 = root.split('/')[-2]
img_path = os.path.join(root, file)
label_path = img_path.replace(".jpg", ".txt")
if not os.path.exists(label_path):
continue
# if num > 2000:
# break
try:
with open(label_path, 'r', encoding='utf-8')as fr:
label = fr.readline().rstrip('\n')
except:
print('!!!!!1')
continue
if label=='-1' or label == '' or label == ' ':
continue
line = './通用文字/通用生成数据datacrop/'+file+'\t'+label+'\n'
print(line)
list.append(line)
num += 1
num = 0
np.random.shuffle(list)
print(len(list))
content = "".join(list)
with open(out, "a", encoding="utf-8") as fp:
# fp.write('\n')
# fp.write(content)
for list1 in list:
# if num > 2000:
# break
fp.write(list1)
num+=1
print(num)
print("done.")
# datadir = '/home/data1/data_yy/行驶证/地址'
# out = '/home/data1/data_yy/行驶证/train.txt'
#
# list = []
# num = 0
# for root, dirs, files in os.walk(datadir):
# for file in files:
# if file.split('.')[-1] == 'jpg':
# dir = root.split('/')[-1]
# img_path = os.path.join(root, file)
# label = file.split('.jpg')[0]
# label = file.split('.jpg')[0]
# if len(label) != 17:
# continue
# line = './地址/'+file+'\t'+label+'\n'
# print(line)
# list.append(line)
# num += 1
#
# np.random.shuffle(list)
# print(len(list))
# content = "".join(list)
# with open(out, "a", encoding="utf-8") as fp:
# # fp.write('\n')
# fp.write(content)
# print(num)
# print("done.")
# list = ['×','√','A','B','C','D','AB','AC','AD','BC','BD','CD','ABC','ABD','ACD','BCD','ABCD']
# datadir = '/home/data1/data_yy/行驶证/通用文字/通用生成数据datacrop'
# out = '/home/data1/data_yy/行驶证/通用文字/train2.txt'
#
# list = []
# num = 0
# for root, dirs, files in os.walk(datadir):
# for file in files:
# if file.split('.')[-1] == 'jpg':
# dir = root.split('/')[-1]
# img_path = os.path.join(root, file)
# label_path = img_path.replace(".jpg", ".json")
# if not os.path.exists(label_path):
# continue
# try:
# label_content = read_json(label_path, "gbk")
# label = label_content["shapes"][0]["label"]
# except:
# continue
# if label == '':
# continue
# # if label not in list:
# # continue
# if label == '-1' or label == '' or label == '1':
# continue
# # if len(label) != 18:
# # print(file)
# if len(label_content["shapes"]) > 0:
# if '\n' in label_content["shapes"][0]["label"]:
# print(file)
# break
# for shape in label_content["shapes"]:
# label = shape["label"]
# line = './通用文字/通用生成数据datacrop/'+file+'\t'+label+'\n'
# list.append(line)
# num += 1
#
# np.random.shuffle(list)
# print(len(list))
# content = "".join(list)
# with open(out, "a", encoding="utf-8") as fp:
# # fp.write('\n')
# fp.write(content)
# print(num)
# print("done.")
#find error char
# with open('/home/gonglei/yuyang/project/PaddleOCR-release-2.6/ppocr/utils/ppocr_keys_v1.txt', 'r')as fr:
# chars = fr.read().split('\n')
# print(len(chars))
# datadir = '/home/data1/data_yy/mingpaidata/biaozhu_0626'
# for root, dirs, files in os.walk(datadir):
# for file in files:
# if file.split('.')[-1] == 'jpg':
# dir = root.split('/')[-1]
# img_path = os.path.join(root, file)
# label_path = img_path.replace(".jpg", ".json")
# if not os.path.exists(label_path):
# continue
# try:
# label_content = read_json(label_path, "gbk")
# except:
# continue
# if len(label_content["shapes"]) > 0:
# for shape in label_content["shapes"]:
# label = shape["label"].replace(' ', '')
# for lbl in label:
# if lbl not in chars:
# print(file)
# print(lbl)
export_jit.py
import os
import sys
__dir__ = os.path.dirname(os.path.abspath(__file__))
sys.path.append(__dir__)
sys.path.append(os.path.abspath(os.path.join(__dir__, '../..')))
from PIL import Image
import cv2
import numpy as np
import math
import time
import torch
from pytorchocr.base_ocr_v20 import BaseOCRV20
from pytorchocr.postprocess import build_post_process
class TextRecognizer(BaseOCRV20):
def __init__(self,rec_model_path,rec_char_dict_path, **kwargs):
rec_image_shape = "3, 32, 320"
self.rec_image_shape = [int(v) for v in rec_image_shape.split(",")]
self.character_type = 'ch'
self.rec_batch_num = 6
self.rec_algorithm = 'CRNN'
self.max_text_length = 25
self.use_space_char = False
self.rec_char_dict_path = rec_char_dict_path
postprocess_params = {
'name': 'CTCLabelDecode',
"character_type": self.character_type,
"character_dict_path": self.rec_char_dict_path,
"use_space_char": self.use_space_char
}
self.postprocess_op = build_post_process(postprocess_params)
use_gpu = True
self.use_gpu = torch.cuda.is_available() and use_gpu
self.limited_max_width = 1280
self.limited_min_width = 16
self.weights_path = rec_model_path
self.yaml_path = None
self.device = "cuda:0"
weights = self.read_pytorch_weights(self.weights_path)
self.out_channels = self.get_out_channels(weights)
print('out_channels: ', self.out_channels)
print(type(kwargs), kwargs)
if self.rec_algorithm == 'NRTR':
self.out_channels = list(weights.values())[-1].numpy().shape[0]
print('out out_channels: ', self.out_channels)
network_config = {'model_type': 'rec',
'algorithm': 'CRNN',
'Transform': None,
'Backbone': {'name': 'ResNet', 'layers': 34},
'Neck': {'name': 'SequenceEncoder', 'hidden_size': 256, 'encoder_type': 'rnn'},
'Head': {'name': 'CTCHead', 'fc_decay': 4e-05}}
kwargs['out_channels'] = self.out_channels
super(TextRecognizer, self).__init__(network_config, **kwargs)
self.load_state_dict(weights)
self.net.eval()
if self.use_gpu:
self.net.cuda()
# self.net = torch.jit.load('./weights_jit/carPaint_rec_use.jit')
# self.net.eval()
if self.use_gpu:
example = torch.rand(1, 3, 32, 320).cuda()
else:
example = torch.rand(1, 3, 32, 320).to(self.device)
traced_script_module = torch.jit.trace(self.net.to(self.device), example)
output = traced_script_module(example)
output1 = self.net(example)
traced_script_module.save('./weights_jit/carPaints_rec_250516_1.jit')
def resize_norm_img(self, img, max_wh_ratio):
imgC, imgH, imgW = self.rec_image_shape
if self.rec_algorithm == 'NRTR':
img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
# return padding_im
image_pil = Image.fromarray(np.uint8(img))
img = image_pil.resize([100, 32], Image.ANTIALIAS)
img = np.array(img)
norm_img = np.expand_dims(img, -1)
norm_img = norm_img.transpose((2, 0, 1))
return norm_img.astype(np.float32) / 128. - 1.
assert imgC == img.shape[2]
max_wh_ratio = max(max_wh_ratio, imgW / imgH)
imgW = int((32 * max_wh_ratio))
imgW = max(min(imgW, self.limited_max_width), self.limited_min_width)
h, w = img.shape[:2]
ratio = w / float(h)
ratio_imgH = math.ceil(imgH * ratio)
ratio_imgH = max(ratio_imgH, self.limited_min_width)
if ratio_imgH > imgW:
resized_w = imgW
else:
resized_w = int(ratio_imgH)
resized_image = cv2.resize(img, (resized_w, imgH))
resized_image = resized_image.astype('float32')
resized_image = resized_image.transpose((2, 0, 1)) / 255
resized_image -= 0.5
resized_image /= 0.5
padding_im = np.zeros((imgC, imgH, imgW), dtype=np.float32)
padding_im[:, :, 0:resized_w] = resized_image
return padding_im
def __call__(self, img_list):
img_num = len(img_list)
# Calculate the aspect ratio of all text bars
width_list = []
for img in img_list:
width_list.append(img.shape[1] / float(img.shape[0]))
# Sorting can speed up the recognition process
indices = np.argsort(np.array(width_list))
# rec_res = []
rec_res = [['', 0.0]] * img_num
batch_num = self.rec_batch_num
elapse = 0
for beg_img_no in range(0, img_num, batch_num):
end_img_no = min(img_num, beg_img_no + batch_num)
norm_img_batch = []
max_wh_ratio = 0
for ino in range(beg_img_no, end_img_no):
# h, w = img_list[ino].shape[0:2]
h, w = img_list[indices[ino]].shape[0:2]
wh_ratio = w * 1.0 / h
max_wh_ratio = max(max_wh_ratio, wh_ratio)
for ino in range(beg_img_no, end_img_no):
if self.rec_algorithm != "SRN":
norm_img = self.resize_norm_img(img_list[indices[ino]],
max_wh_ratio)
norm_img = norm_img[np.newaxis, :]
norm_img_batch.append(norm_img)
else:
norm_img = self.process_image_srn(img_list[indices[ino]],
self.rec_image_shape, 8,
self.max_text_length)
encoder_word_pos_list = []
gsrm_word_pos_list = []
gsrm_slf_attn_bias1_list = []
gsrm_slf_attn_bias2_list = []
encoder_word_pos_list.append(norm_img[1])
gsrm_word_pos_list.append(norm_img[2])
gsrm_slf_attn_bias1_list.append(norm_img[3])
gsrm_slf_attn_bias2_list.append(norm_img[4])
norm_img_batch.append(norm_img[0])
norm_img_batch = np.concatenate(norm_img_batch)
norm_img_batch = norm_img_batch.copy()
if self.rec_algorithm == "SRN":
starttime = time.time()
encoder_word_pos_list = np.concatenate(encoder_word_pos_list)
gsrm_word_pos_list = np.concatenate(gsrm_word_pos_list)
gsrm_slf_attn_bias1_list = np.concatenate(
gsrm_slf_attn_bias1_list)
gsrm_slf_attn_bias2_list = np.concatenate(
gsrm_slf_attn_bias2_list)
with torch.no_grad():
inp = torch.from_numpy(norm_img_batch)
encoder_word_pos_inp = torch.from_numpy(encoder_word_pos_list)
gsrm_word_pos_inp = torch.from_numpy(gsrm_word_pos_list)
gsrm_slf_attn_bias1_inp = torch.from_numpy(gsrm_slf_attn_bias1_list)
gsrm_slf_attn_bias2_inp = torch.from_numpy(gsrm_slf_attn_bias2_list)
if self.use_gpu:
inp = inp.cuda()
encoder_word_pos_inp = encoder_word_pos_inp.cuda()
gsrm_word_pos_inp = gsrm_word_pos_inp.cuda()
gsrm_slf_attn_bias1_inp = gsrm_slf_attn_bias1_inp.cuda()
gsrm_slf_attn_bias2_inp = gsrm_slf_attn_bias2_inp.cuda()
backbone_out = self.net.backbone(inp) # backbone_feat
prob_out = self.net.head(backbone_out, [encoder_word_pos_inp, gsrm_word_pos_inp, gsrm_slf_attn_bias1_inp, gsrm_slf_attn_bias2_inp])
# preds = {"predict": prob_out[2]}
preds = {"predict": prob_out["predict"]}
else:
starttime = time.time()
with torch.no_grad():
inp = torch.from_numpy(norm_img_batch)
if self.use_gpu:
inp = inp.cuda()
prob_out = self.net(inp)
if isinstance(prob_out, list):
preds = [v.cpu().numpy() for v in prob_out]
else:
preds = prob_out.cpu().numpy()
rec_result = self.postprocess_op(preds)
for rno in range(len(rec_result)):
rec_res[indices[beg_img_no + rno]] = rec_result[rno]
elapse += time.time() - starttime
return rec_res, elapse
def main():
rec_model_path = './weights_pt/carPaints_rec_250515.pth' # './weights_pt/tongyongzi_250425.pth'
rec_char_dict_path = './pytorchocr/utils/ppocr_keys_v1.txt' # 'pytorchocr/utils/ppocr_keys_v3.txt' # './pytorchocr/utils/tongyongzi3.txt'
text_recognizer = TextRecognizer(rec_model_path,rec_char_dict_path)
if __name__ == '__main__':
main()
rec_postprocess.py
# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import numpy as np
import torch
class BaseRecLabelDecode(object):
""" Convert between text-label and text-index """
def __init__(self,
character_dict_path=None,
use_space_char=False):
self.beg_str = "sos"
self.end_str = "eos"
self.character_str = []
if character_dict_path is None:
self.character_str = "0123456789abcdefghijklmnopqrstuvwxyz"
dict_character = list(self.character_str)
else:
with open(character_dict_path, "rb") as fin:
lines = fin.readlines()
for line in lines:
line = line.decode('utf-8').strip("\n").strip("\r\n")
self.character_str.append(line)
if use_space_char:
self.character_str.append(" ")
dict_character = list(self.character_str)
dict_character = self.add_special_char(dict_character)
self.dict = {}
for i, char in enumerate(dict_character):
self.dict[char] = i
self.character = dict_character
def add_special_char(self, dict_character):
return dict_character
def decode(self, text_index, text_prob=None, is_remove_duplicate=False):
""" convert text-index into text-label. """
result_list = []
ignored_tokens = self.get_ignored_tokens()
batch_size = len(text_index)
for batch_idx in range(batch_size):
char_list = []
conf_list = []
for idx in range(len(text_index[batch_idx])):
if text_index[batch_idx][idx] in ignored_tokens:
continue
if is_remove_duplicate:
# only for predict
if idx > 0 and text_index[batch_idx][idx - 1] == text_index[
batch_idx][idx]:
continue
char_list.append(self.character[int(text_index[batch_idx][
idx])])
if text_prob is not None:
conf_list.append(text_prob[batch_idx][idx])
else:
conf_list.append(1)
text = ''.join(char_list)
result_list.append((text, np.mean(conf_list)))
return result_list
def get_ignored_tokens(self):
return [0] # for ctc blank
class CTCLabelDecode(BaseRecLabelDecode):
""" Convert between text-label and text-index """
def __init__(self,
character_dict_path=None,
use_space_char=False,
**kwargs):
super(CTCLabelDecode, self).__init__(character_dict_path,
use_space_char)
def __call__(self, preds, label=None, *args, **kwargs):
if isinstance(preds, torch.Tensor):
preds = preds.numpy()
preds_idx = preds.argmax(axis=2)
preds_prob = preds.max(axis=2)
text = self.decode(preds_idx, preds_prob, is_remove_duplicate=True)
if label is None:
return text
label = self.decode(label)
return text, label
def add_special_char(self, dict_character):
dict_character = ['blank'] + dict_character
return dict_character
predict.py
import os
import cv2
import numpy as np
import math
import torch
from rec_postprocess import CTCLabelDecode
class mingpai_recognization():
def __init__(self, model_path: str, alphabet_path: str, device: int = -1):
rec_image_shape = "3, 32, 320"
self.rec_image_shape = [int(v) for v in rec_image_shape.split(",")]
self.character_type = 'ch'
self.rec_batch_num = 6
self.rec_algorithm = 'CRNN'
self.max_text_length = 25
self.use_space_char = True
self.rec_char_dict_path = alphabet_path
postprocess_params = {
'name': 'CTCLabelDecode',
"character_type": self.character_type,
"character_dict_path": self.rec_char_dict_path,
"use_space_char": self.use_space_char
}
self.postprocess_op = CTCLabelDecode(**postprocess_params)
use_gpu = True
self.use_gpu = use_gpu
self.device = f"cuda:{device}" if device >= 0 and torch.cuda.is_available() else "cpu"
self.limited_max_width = 1280
self.limited_min_width = 16
self.net = torch.jit.load(model_path)
self.net.to(self.device)
self.net.eval()
def post_process(self,preds):
return self.postprocess_op(preds)
def pre_process(self,img):
# Calculate the aspect ratio of all text bars
width_list = []
width_list.append(img.shape[1] / float(img.shape[0]))
# rec_res = []
norm_img_batch = []
max_wh_ratio = 0
h, w = img.shape[0:2]
wh_ratio = w * 1.0 / h
max_wh_ratio = max(max_wh_ratio, wh_ratio)
norm_img = self.resize_norm_img(img,max_wh_ratio)
norm_img = norm_img[np.newaxis, :]
norm_img_batch.append(norm_img)
norm_img_batch = np.concatenate(norm_img_batch)
norm_img_batch = norm_img_batch.copy()
return norm_img_batch
def resize_norm_img(self, img, max_wh_ratio):
imgC, imgH, imgW = self.rec_image_shape
assert imgC == img.shape[2]
max_wh_ratio = max(max_wh_ratio, imgW / imgH)
imgW = int((32 * max_wh_ratio))
imgW = max(min(imgW, self.limited_max_width), self.limited_min_width)
h, w = img.shape[:2]
ratio = w / float(h)
ratio_imgH = math.ceil(imgH * ratio)
ratio_imgH = max(ratio_imgH, self.limited_min_width)
if ratio_imgH > imgW:
resized_w = imgW
else:
resized_w = int(ratio_imgH)
resized_image = cv2.resize(img, (resized_w, imgH))
resized_image = resized_image.astype('float32')
resized_image = resized_image.transpose((2, 0, 1)) / 255
resized_image -= 0.5
resized_image /= 0.5
padding_im = np.zeros((imgC, imgH, imgW), dtype=np.float32)
padding_im[:, :, 0:resized_w] = resized_image
return padding_im
def predict(self, img):
norm_img_batch =self.pre_process(img)
with torch.no_grad():
inp = torch.from_numpy(norm_img_batch).to(self.device)
prob_out = self.net(inp)
if isinstance(prob_out, list):
preds = [v.cpu().numpy() for v in prob_out]
else:
preds = prob_out.cpu().numpy()
rec_result = self.post_process(preds)
return rec_result[0]
def process(self, images):
result = []
for img in images:
res = self.predict(img)
result.append(res[0])
return result
if __name__ == '__main__':
rec_model_path = 'weights_jit/xingshizheng_1221.jit'
rec_char_dict_path = 'pytorchocr/utils/ppocr_keys_v3.txt'
text_recognizer = mingpai_recognization(rec_model_path, rec_char_dict_path,0)
datadir = 'mingpai_pic/'
names = os.listdir(datadir)
i = 0
for name in names:
i = i + 1
picpath = datadir + name
img = cv2.imdecode(np.fromfile(picpath, dtype=np.uint8), cv2.IMREAD_UNCHANGED)
pred,score = text_recognizer.predict(img)
print(picpath,i, pred)
cobver (mingpai_rec_converter.py paddleOCR模型转pytorch模型)
# https://zhuanlan.zhihu.com/p/335753926
import os, sys
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from collections import OrderedDict
import numpy as np
import cv2
import torch
from pytorchocr.base_ocr_v20 import BaseOCRV20
class ServerV20RecConverter(BaseOCRV20):
def __init__(self, config, paddle_pretrained_model_path, **kwargs):
para_state_dict, opti_state_dict = self.read_paddle_weights(paddle_pretrained_model_path)
out_channels = list(para_state_dict.values())[-1].shape[0]
print('out_channels: ', out_channels)
print(type(kwargs), kwargs)
kwargs['out_channels'] = out_channels
super(ServerV20RecConverter, self).__init__(config, **kwargs)
# self.load_paddle_weights(paddle_pretrained_model_path)
self.load_paddle_weights([para_state_dict, opti_state_dict])
print('model is loaded: {}'.format(paddle_pretrained_model_path))
self.net.eval()
def load_paddle_weights(self, paddle_weights):
para_state_dict, opti_state_dict = paddle_weights
for k,v in self.net.state_dict().items():
keyword = 'block_list.'
if keyword in k:
# replace: 'block_list.' -> ''
name = k.replace(keyword, '')
else:
name = k
if name.endswith('num_batches_tracked'):
continue
if name.endswith('running_mean'):
ppname = name.replace('running_mean', '_mean')
elif name.endswith('running_var'):
ppname = name.replace('running_var', '_variance')
elif name.endswith('bias') or name.endswith('weight'):
ppname = name
elif 'lstm' in name:
ppname = name
else:
print('Redundance:')
print(name)
raise ValueError
try:
if ppname.endswith('fc.weight'):
self.net.state_dict()[k].copy_(torch.Tensor(para_state_dict[ppname].T))
else:
self.net.state_dict()[k].copy_(torch.Tensor(para_state_dict[ppname]))
except Exception as e:
print('pytorch: {}, {}'.format(k, v.size()))
print('paddle: {}, {}'.format(ppname, para_state_dict[ppname].shape))
raise e
print('model is loaded.')
if __name__ == '__main__':
import argparse, json, textwrap, sys, os
parser = argparse.ArgumentParser()
parser.add_argument("--src_model_path", type=str, default='weights_paddle', help='Assign the paddleOCR trained model(best_accuracy)')
args = parser.parse_args()
cfg = {'model_type':'rec',
'algorithm':'CRNN',
'Transform':None,
'Backbone':{'name':'ResNet', 'layers':34},
'Neck':{'name':'SequenceEncoder', 'hidden_size':256, 'encoder_type':'rnn'},
'Head':{'name':'CTCHead', 'fc_decay': 4e-05}}
paddle_pretrained_model_path = os.path.join(os.path.abspath(args.src_model_path), 'latest')
converter = ServerV20RecConverter(cfg, paddle_pretrained_model_path)
# save
converter.save_pytorch_weights('weights_pt/tongyongzi_250425.pth')
print('done.')
obnx_predict.py
# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import os
import sys
from PIL import Image
__dir__ = os.path.dirname(os.path.abspath(__file__))
sys.path.append(__dir__)
sys.path.insert(0, os.path.abspath(os.path.join(__dir__, '../..')))
os.environ["FLAGS_allocator_strategy"] = 'auto_growth'
import cv2
import numpy as np
import math
import time
import traceback
import paddle
import onnxruntime as ort
from ppocr.utils.logging import get_logger
from ppocr.utils.utility import get_image_file_list, check_and_read
from ppocr.postprocess.rec_postprocess import CTCLabelDecode
logger = get_logger()
class TextRecognizer(object):
def __init__(self, model_path='rec_ppocr_v3_teacher_0801.onnx', dict_path='ppocr/utils/ppocr_keys_v1.txt',
device='cpu', image_shape=(3, 48, 320)):
self.rec_image_shape = image_shape
self.rec_batch_num = 6
self.postprocess_op = CTCLabelDecode(dict_path, use_space_char=True)
self.predictor = ort.InferenceSession(model_path)
self.input_tensor, self.output_tensors = self.predictor.get_inputs()[0], None
def __call__(self, img_list):
img_num = len(img_list)
# Calculate the aspect ratio of all text bars
width_list = []
for img in img_list:
width_list.append(img.shape[1] / float(img.shape[0]))
# Sorting can speed up the recognition process
indices = np.argsort(np.array(width_list))
rec_res = [['', 0.0]] * img_num
batch_num = self.rec_batch_num
st = time.time()
for beg_img_no in range(0, img_num, batch_num):
end_img_no = min(img_num, beg_img_no + batch_num)
norm_img_batch = []
imgC, imgH, imgW = self.rec_image_shape[:3]
max_wh_ratio = imgW / imgH
# max_wh_ratio = 0
for ino in range(beg_img_no, end_img_no):
h, w = img_list[indices[ino]].shape[0:2]
wh_ratio = w * 1.0 / h
max_wh_ratio = max(max_wh_ratio, wh_ratio)
for ino in range(beg_img_no, end_img_no):
norm_img,_ = self.resize_norm_img(img_list[indices[ino]],[3,48,320])
norm_img = norm_img[np.newaxis, :]
norm_img_batch.append(norm_img)
norm_img_batch = np.concatenate(norm_img_batch)
norm_img_batch = norm_img_batch.copy()
input_dict = {}
input_dict[self.input_tensor.name] = norm_img_batch
outputs = self.predictor.run(self.output_tensors,
input_dict)
preds = outputs[0]
rec_result = self.postprocess_op(preds)
for rno in range(len(rec_result)):
rec_res[indices[beg_img_no + rno]] = rec_result[rno]
return rec_res, time.time() - st
def resize_norm_img(self,img,
image_shape,
padding=True,
interpolation=cv2.INTER_LINEAR):
imgC, imgH, imgW = image_shape
h = img.shape[0]
w = img.shape[1]
if not padding:
resized_image = cv2.resize(
img, (imgW, imgH), interpolation=interpolation)
resized_w = imgW
else:
ratio = w / float(h)
if math.ceil(imgH * ratio) > imgW:
resized_w = imgW
else:
resized_w = int(math.ceil(imgH * ratio))
resized_image = cv2.resize(img, (resized_w, imgH))
resized_image = resized_image.astype('float32')
if image_shape[0] == 1:
resized_image = resized_image / 255
resized_image = resized_image[np.newaxis, :]
else:
resized_image = resized_image.transpose((2, 0, 1)) / 255
resized_image -= 0.5
resized_image /= 0.5
padding_im = np.zeros((imgC, imgH, imgW), dtype=np.float32)
padding_im[:, :, 0:resized_w] = resized_image
valid_ratio = min(1.0, float(resized_w / imgW))
return padding_im, valid_ratio
def main():
image_dir = 'test_imgs'
rec_model_dir = 'rec_ppocr_v3_teacher_0801.onnx'
# rec_model_dir = '/home/zl/work/自助终端/model_data/rec/inference.onnx'
rec_char_dict_path='ppocr/utils/ppocr_keys_v1.txt'
image_file_list = get_image_file_list(image_dir)
text_recognizer = TextRecognizer(model_path=rec_model_dir, dict_path=rec_char_dict_path)
valid_image_file_list = []
img_list = []
logger.info(
"In PP-OCRv3, rec_image_shape parameter defaults to '3, 48, 320', "
"if you are using recognition model with PP-OCRv2 or an older version, please set --rec_image_shape='3,32,320"
)
for image_file in image_file_list:
img, flag, _ = check_and_read(image_file)
if not flag:
img = cv2.imread(image_file)
if img is None:
logger.info("error in loading image:{}".format(image_file))
continue
valid_image_file_list.append(image_file)
img_list.append(img)
try:
rec_res, _ = text_recognizer(img_list)
except Exception as E:
logger.info(traceback.format_exc())
logger.info(E)
exit()
for ino in range(len(img_list)):
logger.info("Predicts of {}:{}".format(valid_image_file_list[ino],
rec_res[ino]))
if __name__ == "__main__":
main()
更多推荐
所有评论(0)