Git链接

安装按照上面的说明，说下使用。

把tools下面的test做了一点修改，可以读取一张图片，把里面的单个字符都检测和识别出来。

然后绘制到屏幕上。

import torch
from charnet.modeling.model import CharNet
import cv2, os
import numpy as np
import argparse
from charnet.config import cfg

def loadDict():
    fn_dict="tools\char_dict.txt"
    with open(fn_dict, 'r') as file:
        lines = file.readlines()
    # 去除每行末尾的换行符
    lines = [line.strip() for line in lines]
    dict_char={}
    for line in lines:
        line=line.replace("\x1f","")
        num_line=len(line)
        a=line[0]
        index=line[1:]
        index=int(index)
        dict_char[index]=a
    return dict_char

def resize(im, size):
    h, w, _ = im.shape
    scale = max(h, w) / float(size)
    image_resize_height = int(round(h / scale / cfg.SIZE_DIVISIBILITY) * cfg.SIZE_DIVISIBILITY)
    image_resize_width = int(round(w / scale / cfg.SIZE_DIVISIBILITY) * cfg.SIZE_DIVISIBILITY)
    scale_h = float(h) / image_resize_height
    scale_w = float(w) / image_resize_width
    im = cv2.resize(im, (image_resize_width, image_resize_height), interpolation=cv2.INTER_LINEAR)
    return im, scale_w, scale_h, w, h

if __name__ == '__main__':

    dict_char=loadDict()
    parser = argparse.ArgumentParser(description="Test")
    fn_conf=r"configs\icdar2015_hourglass88.yaml"
    fn_weight=r"configs\icdar2015_hourglass88.pth"
    args = parser.parse_args()
    cfg.merge_from_file(fn_conf)
    cfg.freeze()

    charnet = CharNet()
    charnet.load_state_dict(torch.load(fn_weight))
    charnet.eval()
    charnet.cuda()
    im_file=r"data\2.jpg"
    im_original = cv2.imread(im_file)
    im, scale_w, scale_h, original_w, original_h = resize(im_original, size=cfg.INPUT_SIZE)
    with torch.no_grad():
        char_bboxes, char_scores, word_instances = charnet(im, scale_w, scale_h, original_w, original_h)
        for ic,box in enumerate(char_bboxes):
            print(box)
            score=char_scores[ic]
            max_index = np.argmax(score)
            label=dict_char[max_index]
            points = np.array(box[0:8]).reshape(-1, 2).astype(np.int32)
            cv2.polylines(im_original, [points], isClosed=True, color=(0, 0, 255), thickness=1)
            font = cv2.FONT_HERSHEY_SIMPLEX
            cv2.putText(im_original, label, (points[0][0],points[0][1]), font, 1, (0, 255, 0), 1)
            cv2.imshow("img",im_original)
            cv2.waitKey(0)