File size: 1,486 Bytes
a99b9cb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 |
import layoutparser as lp
import cv2
import matplotlib.pyplot as plt
image = cv2.imread("Projects/HandwritingOCR/captured_images/pasted_image.jpg")
#image = image[..., ::-1]
model = lp.Detectron2LayoutModel(
"lp://PubLayNet/faster_rcnn_R_50_FPN_3x/config",
extra_config=["MODEL.ROI_HEADS.SCORE_THRESH_TEST", 0.8],
label_map={0: "Text", 1: "Title", 2: "List", 3:"Table", 4:"Figure"},
)
# Phát hiện bố cục trên ảnh đã được cải thiện độ phân giải
layout = model.detect(image)
# Vẽ các hộp phát hiện lên ảnh
lp.draw_box(image, layout, box_width=3)
print(type(layout))
print(layout[0])
text_blocks = lp.Layout([b for b in layout if b.type=='Text'])
figure_blocks = lp.Layout([b for b in layout if b.type=='Figure'])
text_blocks = lp.Layout([b for b in text_blocks \
if not any(b.is_in(b_fig) for b_fig in figure_blocks)])
h, w = image.shape[:2]
left_interval = lp.Interval(0, w/2*1.05, axis='x').put_on_canvas(image)
left_blocks = text_blocks.filter_by(left_interval, center=True)
left_blocks.sort(key = lambda b:b.coordinates[1])
right_blocks = [b for b in text_blocks if b not in left_blocks]
right_blocks.sort(key = lambda b:b.coordinates[1])
# And finally combine the two list and add the index
# according to the order
text_blocks = lp.Layout([b.set(id = idx) for idx, b in enumerate(left_blocks + right_blocks)])
lp.draw_box(image, text_blocks,
box_width=3,
show_element_id=True) |