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 test_net(self, net, image, cuda):
# resize
img_resized, target_ratio, size_heatmap = imgproc.resize_aspect_ratio(
image, 1280, interpolation=cv2.INTER_LINEAR, mag_ratio=1.5)
# 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, _ = net(x)
# make score and link map
score_text = y[0, :, :, 0].cpu().data.numpy()
score_link = y[0, :, :, 1].cpu().data.numpy()
return score_text, score_link
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)
def text_detect(self, image):
# if not os.path.exists(image_path):
# print("Not exists path")
# return []
# image = imgproc.loadImage(image_path) # numpy array img (RGB order)
# image = cv2.imread()
time0 = time.time()
# resize
img_resized, target_ratio, size_heatmap = imgproc.resize_aspect_ratio(
image,
pr.canvas_size,
interpolation=cv2.INTER_LINEAR,
mag_ratio=pr.mag_ratio)
print(img_resized.shape)
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 pr.cuda:
x = x.cuda()
# forward pass
with torch.no_grad():
y, feature = self.model(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 self.refine_model is not None:
with torch.no_grad():
y_refiner = self.refine_model(y, feature)
score_link = y_refiner[0, :, :, 0].cpu().data.numpy()
time0 = time.time() - time0
time1 = time.time()
# Post-processing
boxes, polys = craft_utils.getDetBoxes(score_text, score_link,
pr.text_threshold,
pr.link_threshold, pr.low_text,
pr.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]
# expand box: poly = np.array([[l, t], [r, t], [r, b], [l, b]], dtype=np.float32)
if pr.horizontal_mode:
if self.check_horizontal(polys):
height, width, channel = image.shape
new_polys = []
for box in polys:
[[l1, t1], [r1, t2], [r2, b1], [l2, b2]] = box
if t1 < t2:
l, r, t, b = l2, r1, t1, b1
elif t1 > t2:
l, r, t, b = l1, r2, t2, b2
else:
l, r, t, b = l1, r1, t1, b1
h_box = abs(b - t)
t = max(0, t - h_box * pr.expand_ratio)
b = min(b + h_box * pr.expand_ratio, height)
x_min, y_min, x_max, y_max = l, t, r, b
new_box = [x_min, y_min, x_max, y_max]
new_polys.append(new_box)
polys = np.array(new_polys, dtype=np.float32)
# for box in polys:
time1 = time.time() - time1
total_time = round(time0 + time1, 2)
# render results (optional)
render_img = score_text.copy()
render_img = np.hstack((render_img, score_link))
ret_score_text = imgproc.cvt2HeatmapImg(render_img)
if pr.show_time:
print("\ninfer/postproc time : {:.3f}/{:.3f}".format(time0, time1))
if pr.folder_test:
return boxes, polys, ret_score_text
if pr.visualize:
img_draw = displayResult(img=image[:, :, ::-1], boxes=polys)
plt.imshow(cv2.cvtColor(img_draw, cv2.COLOR_RGB2BGR))
plt.show()
result_boxes = []
for box in polys:
result_boxes.append(box.tolist())
return result_boxes, total_time
def text_detect(self, image, have_cmnd=True):
time0 = time.time()
# resize
img_resized, target_ratio, size_heatmap = imgproc.resize_aspect_ratio(
image,
pr.canvas_size,
interpolation=cv2.INTER_LINEAR,
mag_ratio=pr.mag_ratio)
print(img_resized.shape)
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 pr.cuda:
x = x.cuda()
# forward pass
with torch.no_grad():
y, feature = self.model(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 self.refine_model is not None:
with torch.no_grad():
y_refiner = self.refine_model(y, feature)
score_link = y_refiner[0, :, :, 0].cpu().data.numpy()
# Post-processing
boxes, polys = craft_utils.getDetBoxes(score_text, score_link,
pr.text_threshold,
pr.link_threshold, pr.low_text,
pr.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]
# get box + extend
list_box = []
for box in polys:
[[l1, t1], [r1, t2], [r2, b1], [l2, b2]] = box
if t1 < t2:
l, r, t, b = l2, r1, t1, b1
elif t1 > t2:
l, r, t, b = l1, r2, t2, b2
else:
l, r, t, b = l1, r1, t1, b1
xmin, ymin, xmax, ymax = l, t, r, b
xmin, ymin, xmax, ymax = max(0, xmin - int((xmax - xmin) * pr.expand_ratio)),\
max(0, ymin - int((ymax - ymin) * pr.expand_ratio)),\
xmax + int((xmax - xmin) * pr.expand_ratio),\
ymax + int((ymax - ymin) * pr.expand_ratio)
list_box.append([xmin, ymin, xmax, ymax])
# sort line
dict_cum_sorted = self.sort_line_cmnd(list_box)
list_box_optim = []
for cum in dict_cum_sorted:
for box in cum:
list_box_optim.append(box)
# draw box on image
img_res = image.copy()
img_res = np.ascontiguousarray(img_res)
for box in list_box_optim:
xmin, ymin, xmax, ymax = box
cv2.rectangle(img_res, (int(xmin), int(ymin)),
(int(xmax), int(ymax)), (29, 187, 255), 2, 2)
# crop image
result_list_img_cum = []
image_PIL = Image.fromarray(image)
for cum in dict_cum_sorted:
list_img = []
for box in cum:
xmin, ymin, xmax, ymax = box
list_img.append(image_PIL.copy().crop(
(xmin, ymin, xmax, ymax)))
result_list_img_cum.append(list_img)
return result_list_img_cum, img_res, None