115 lines
4.2 KiB
Python
115 lines
4.2 KiB
Python
"""Modified from https://github.com/JaidedAI/EasyOCR/blob/803b907/easyocr/detection.py.
|
|
1. Disable DataParallel.
|
|
"""
|
|
import torch
|
|
import torch.backends.cudnn as cudnn
|
|
from torch.autograd import Variable
|
|
from PIL import Image
|
|
from collections import OrderedDict
|
|
|
|
import cv2
|
|
import numpy as np
|
|
from .craft_utils import getDetBoxes, adjustResultCoordinates
|
|
from .imgproc import resize_aspect_ratio, normalizeMeanVariance
|
|
from .craft import CRAFT
|
|
|
|
def copyStateDict(state_dict):
|
|
if list(state_dict.keys())[0].startswith("module"):
|
|
start_idx = 1
|
|
else:
|
|
start_idx = 0
|
|
new_state_dict = OrderedDict()
|
|
for k, v in state_dict.items():
|
|
name = ".".join(k.split(".")[start_idx:])
|
|
new_state_dict[name] = v
|
|
return new_state_dict
|
|
|
|
def test_net(canvas_size, mag_ratio, net, image, text_threshold, link_threshold, low_text, poly, device, estimate_num_chars=False):
|
|
if isinstance(image, np.ndarray) and len(image.shape) == 4: # image is batch of np arrays
|
|
image_arrs = image
|
|
else: # image is single numpy array
|
|
image_arrs = [image]
|
|
|
|
img_resized_list = []
|
|
# resize
|
|
for img in image_arrs:
|
|
img_resized, target_ratio, size_heatmap = resize_aspect_ratio(img, canvas_size,
|
|
interpolation=cv2.INTER_LINEAR,
|
|
mag_ratio=mag_ratio)
|
|
img_resized_list.append(img_resized)
|
|
ratio_h = ratio_w = 1 / target_ratio
|
|
# preprocessing
|
|
x = [np.transpose(normalizeMeanVariance(n_img), (2, 0, 1))
|
|
for n_img in img_resized_list]
|
|
x = torch.from_numpy(np.array(x))
|
|
x = x.to(device)
|
|
|
|
# forward pass
|
|
with torch.no_grad():
|
|
y, feature = net(x)
|
|
|
|
boxes_list, polys_list = [], []
|
|
for out in y:
|
|
# make score and link map
|
|
score_text = out[:, :, 0].cpu().data.numpy()
|
|
score_link = out[:, :, 1].cpu().data.numpy()
|
|
|
|
# Post-processing
|
|
boxes, polys, mapper = getDetBoxes(
|
|
score_text, score_link, text_threshold, link_threshold, low_text, poly, estimate_num_chars)
|
|
|
|
# coordinate adjustment
|
|
boxes = adjustResultCoordinates(boxes, ratio_w, ratio_h)
|
|
polys = adjustResultCoordinates(polys, ratio_w, ratio_h)
|
|
if estimate_num_chars:
|
|
boxes = list(boxes)
|
|
polys = list(polys)
|
|
for k in range(len(polys)):
|
|
if estimate_num_chars:
|
|
boxes[k] = (boxes[k], mapper[k])
|
|
if polys[k] is None:
|
|
polys[k] = boxes[k]
|
|
boxes_list.append(boxes)
|
|
polys_list.append(polys)
|
|
|
|
return boxes_list, polys_list
|
|
|
|
def get_detector(trained_model, device='cpu', quantize=True, cudnn_benchmark=False):
|
|
net = CRAFT()
|
|
|
|
if device == 'cpu':
|
|
net.load_state_dict(copyStateDict(torch.load(trained_model, map_location=device)))
|
|
if quantize:
|
|
try:
|
|
torch.quantization.quantize_dynamic(net, dtype=torch.qint8, inplace=True)
|
|
except:
|
|
pass
|
|
else:
|
|
net.load_state_dict(copyStateDict(torch.load(trained_model, map_location=device)))
|
|
# net = torch.nn.DataParallel(net).to(device)
|
|
net = net.to(device)
|
|
cudnn.benchmark = cudnn_benchmark
|
|
|
|
net.eval()
|
|
return net
|
|
|
|
def get_textbox(detector, image, canvas_size, mag_ratio, text_threshold, link_threshold, low_text, poly, device, optimal_num_chars=None, **kwargs):
|
|
result = []
|
|
estimate_num_chars = optimal_num_chars is not None
|
|
bboxes_list, polys_list = test_net(canvas_size, mag_ratio, detector,
|
|
image, text_threshold,
|
|
link_threshold, low_text, poly,
|
|
device, estimate_num_chars)
|
|
if estimate_num_chars:
|
|
polys_list = [[p for p, _ in sorted(polys, key=lambda x: abs(optimal_num_chars - x[1]))]
|
|
for polys in polys_list]
|
|
|
|
for polys in polys_list:
|
|
single_img_result = []
|
|
for i, box in enumerate(polys):
|
|
poly = np.array(box).astype(np.int32).reshape((-1))
|
|
single_img_result.append(poly)
|
|
result.append(single_img_result)
|
|
|
|
return result
|