def test_net(
net,
image,
text_threshold,
link_threshold,
low_text,
cuda,
poly,
canvas_size=1280,
mag_ratio=1.5,
):
# resize
img_resized, target_ratio, size_heatmap = imgproc.resize_aspect_ratio(
image,
canvas_size,
interpolation=cv2.INTER_LINEAR,
mag_ratio=mag_ratio)
ratio_h = ratio_w = 1 / target_ratio
# preprocessing
x = imgproc.normalizeMeanVariance(img_resized)
x = torch.from_numpy(x).permute(2, 0, 1) # [h, w, c] to [c, h, w]
x = Variable(x.unsqueeze(0)) # [c, h, w] to [b, c, h, w]
if cuda:
x = x.cuda()
# forward pass
with torch.no_grad():
y, feature = net(x)
# make score and link map
score_text = y[0, :, :, 0].cpu().data.numpy().astype(np.float32)
score_link = y[0, :, :, 1].cpu().data.numpy().astype(np.float32)
# NOTE
score_text = score_text[:size_heatmap[0], :size_heatmap[1]]
score_link = score_link[:size_heatmap[0], :size_heatmap[1]]
# Post-processing
boxes, polys = getDetBoxes(score_text, score_link, text_threshold,
link_threshold, low_text, poly)
# coordinate adjustment
boxes = adjustResultCoordinates(boxes, ratio_w, ratio_h)
polys = adjustResultCoordinates(polys, ratio_w, ratio_h)
for k in range(len(polys)):
if polys[k] is None:
polys[k] = boxes[k]
# render results (optional)
score_text = score_text.copy()
render_score_text = imgproc.cvt2HeatmapImg(score_text)
render_score_link = imgproc.cvt2HeatmapImg(score_link)
render_img = [render_score_text, render_score_link]
# ret_score_text = imgproc.cvt2HeatmapImg(render_img)
return boxes, polys, render_img
def test_net(net, image, text_threshold, link_threshold, low_text, poly,
refine_net):
# t0 = time.time()
# resize
img_resized, target_ratio, size_heatmap = resize_aspect_ratio(
image,
args.canvas_size,
interpolation=cv2.INTER_LINEAR,
mag_ratio=args.mag_ratio)
ratio_h = ratio_w = 1 / target_ratio
# preprocessing
x = img_normalize(img_resized)
x = np.expand_dims(x, axis=0) # [h, w, c] to [b, h, w, c]
# forward pass
y, feature = net(x)
# make score and link map
score_text = y[0, :, :, 0]
score_link = y[0, :, :, 1]
# refine link
if refine_net is not None:
# TODO
pass
# t0 = time.time() - t0
# t1 = time.time()
# Post-processing
boxes, polys = getDetBoxes(score_text, score_link, text_threshold,
link_threshold, low_text, poly)
# coordinate adjustment
boxes = adjustResultCoordinates(boxes, ratio_w, ratio_h)
polys = adjustResultCoordinates(polys, ratio_w, ratio_h)
for k in range(len(polys)):
if polys[k] is None:
polys[k] = boxes[k]
# t1 = time.time() - t1
# render results (optional)
render_img = score_text.numpy()
render_img = np.hstack((render_img, np.ones(
(np.shape(render_img)[0], 5)), score_link))
ret_score_text = score_to_heat_map(render_img)
# if args.show_time : print("\ninfer/postproc time : {:.3f}/{:.3f}".format(t0, t1))
return boxes, polys, ret_score_text
def test_net(net, image, text_threshold, link_threshold, low_text, cuda, poly, refine_net=None):
t0 = time.time()
# resize
img_resized, target_ratio, size_heatmap = imgproc.resize_aspect_ratio(image, args.canvas_size, interpolation=cv2.INTER_LINEAR, mag_ratio=args.mag_ratio)
ratio_h = ratio_w = 1 / target_ratio
# preprocessing
x = imgproc.normalizeMeanVariance(img_resized)
x = torch.from_numpy(x).permute(2, 0, 1) # [h, w, c] to [c, h, w]
x = Variable(x.unsqueeze(0)) # [c, h, w] to [b, c, h, w]
if cuda:
x = x.cuda()
# forward pass
with torch.no_grad():
y, feature = net(x)
# make score and link map
score_text = y[0,:,:,0].cpu().data.numpy()
score_link = y[0,:,:,1].cpu().data.numpy()
# refine link
if refine_net is not None:
with torch.no_grad():
y_refiner = refine_net(y, feature)
score_link = y_refiner[0,:,:,0].cpu().data.numpy()
t0 = time.time() - t0
t1 = time.time()
# Post-processing
boxes, polys = craft_utils.getDetBoxes(score_text, score_link, text_threshold, link_threshold, low_text, poly)
# coordinate adjustment
boxes = craft_utils.adjustResultCoordinates(boxes, ratio_w, ratio_h)
polys = craft_utils.adjustResultCoordinates(polys, ratio_w, ratio_h)
for k in range(len(polys)):
if polys[k] is None: polys[k] = boxes[k]
t1 = time.time() - t1
# render results (optional)
render_img = score_text.copy()
render_img = np.hstack((render_img, score_link))
ret_score_text = imgproc.cvt2HeatmapImg(render_img)
if args.show_time : print("\ninfer/postproc time : {:.3f}/{:.3f}".format(t0, t1))
return boxes, polys, ret_score_text