安装按照上面的说明,说下使用。
把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)