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
def saveInput(imagename, vis_dir, image, region_scores, affinity_scores,
confidence_mask):
image = np.uint8(image.copy())
image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
boxes, polys = craft_utils.getDetBoxes(region_scores, affinity_scores,
0.85, 0.2, 0.5, False)
if image.shape[0] / region_scores.shape[0] >= 2:
boxes = np.array(boxes, np.int32) * 2
else:
boxes = np.array(boxes, np.int32)
if len(boxes) > 0:
np.clip(boxes[:, :, 0], 0, image.shape[1])
np.clip(boxes[:, :, 1], 0, image.shape[0])
for box in boxes:
cv2.polylines(image, [np.reshape(box, (-1, 1, 2))], True,
(0, 0, 255))
target_gaussian_heatmap_color = imgproc.cvt2HeatmapImg(region_scores)
target_gaussian_affinity_heatmap_color = imgproc.cvt2HeatmapImg(
affinity_scores)
confidence_mask_gray = imgproc.cvt2HeatmapImg(confidence_mask)
# overlay
height, width, channel = image.shape
overlay_region = cv2.resize(target_gaussian_heatmap_color, (width, height))
overlay_aff = cv2.resize(target_gaussian_affinity_heatmap_color,
(width, height))
confidence_mask_gray = cv2.resize(confidence_mask_gray, (width, height),
interpolation=cv2.INTER_NEAREST)
overlay_region = cv2.addWeighted(image, 0.4, overlay_region, 0.6, 5)
overlay_aff = cv2.addWeighted(image, 0.4, overlay_aff, 0.7, 6)
gt_scores = np.concatenate([overlay_region, overlay_aff], axis=1)
output = np.concatenate([gt_scores, confidence_mask_gray], axis=1)
output = np.hstack([image, output])
# synthtext
if type(imagename) is not str:
imagename = imagename[0].split("/")[-1][:-4]
outpath = vis_dir + f"/{imagename}_input.jpg"
if not os.path.exists(os.path.dirname(outpath)):
os.makedirs(os.path.dirname(outpath), exist_ok=True)
cv2.imwrite(outpath, output)
def detect(self, image, set_batch):
img_resized, img_srcs, img_ratio = imgproc.resize_aspect_ratio(
image, 512, 512, set_batch=set_batch)
img_resized = torch.from_numpy(img_resized)
with torch.no_grad():
# torch.onnx.export(self.net, img_resized, "craft_dynamic.onnx", verbose=True, input_names=['input'], opset_version=12,
# output_names=['output1', 'output2'],
# dynamic_axes={"input": {0: "-1", 2: "-1", 3: "-1"}, "output1": {0: "-1"}, "output2": {0: "-1"}})
y_refiner, score_text = self.net(img_resized)
score_links = y_refiner[:, 0, :, :].cpu().data.numpy()
score_texts = score_text[:, 0, :, :].cpu().data.numpy()
polys_list = []
for idx, score_link in enumerate(score_links):
boxes, polys = craft_utils.getDetBoxes(score_texts[idx],
score_link,
self.text_threshold,
self.link_threshold,
self.low_text_score,
poly=self.poly)
for k in range(len(polys)):
if polys[k] is None:
polys[k] = boxes[k]
polys_list.append(polys)
point_list2 = []
for idx, image_src in enumerate(img_srcs):
img = np.array(image_src[:, :, ::-1])
with open('./result/' + "gt_" + str(idx) + '.txt', 'w') as f:
for i, box in enumerate(polys_list[idx]):
box[:, 0] *= img_ratio[idx][1] * 2
box[:, 1] *= img_ratio[idx][0] * 2
# warp_img = mask_crop.imgWrapAffine(img, np.int0(box))
# cv2.imwrite(''.join(['./result/warp_img_', str(i), '.jpg']), warp_img)
poly = np.array(box).astype(np.int32).reshape((-1))
point_list2.append(','.join([str(p)
for p in poly]) + ',1\n')
f.write(','.join([str(p) for p in poly]) + ',1\n')
cv2.polylines(img, [poly.reshape((-1, 1, 2))],
True,
color=(0, 0, 255),
thickness=2)
cv2.imwrite(''.join(['./result/', str(time.time()), '.jpg']), img)