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()

更多推荐