Git链接
安装按照上面的说明,说下使用。
把tools下面的test做了一点修改,可以读取一张图片,把里面的单个字符都检测和识别出来。
然后绘制到屏幕上。
import torch
from charnet.modeling.model import CharNet
import cv2, os
import numpy as np
import argparse
from charnet.config import cfgdef 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]=areturn dict_chardef resize(im, size):h, w, _ = im.shapescale = 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_heightscale_w = float(w) / image_resize_widthim = cv2.resize(im, (image_resize_width, image_resize_height), interpolation=cv2.INTER_LINEAR)return im, scale_w, scale_h, w, hif __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_SIMPLEXcv2.putText(im_original, label, (points[0][0],points[0][1]), font, 1, (0, 255, 0), 1)cv2.imshow("img",im_original)cv2.waitKey(0)