def compo_in_img(org, bin, corners):
def reverse(img):
rec, b = cv2.threshold(img, 1, 255, cv2.THRESH_BINARY_INV)
return b
corners_compo = []
pad = 2
for corner in corners:
(top_left, bottom_right) = corner
(col_min, row_min) = top_left
(col_max, row_max) = bottom_right
col_min = max(col_min - pad, 0)
col_max = min(col_max + pad, org.shape[1])
row_min = max(row_min - pad, 0)
row_max = min(row_max + pad, org.shape[0])
clip_bin = bin[row_min:row_max, col_min:col_max]
clip_bin = reverse(clip_bin)
boundary_all, boundary_rec, boundary_nonrec = det.boundary_detection(
clip_bin,
C.THRESHOLD_OBJ_MIN_AREA,
C.THRESHOLD_OBJ_MIN_PERIMETER, # size of area
C.THRESHOLD_LINE_THICKNESS, # line check
C.THRESHOLD_REC_MIN_EVENNESS_STRONG,
C.THRESHOLD_IMG_MAX_DENT_RATIO) # rectangle check
corners_rec = det.get_corner(boundary_rec)
corners_rec = utils.corner_cvt_relative_position(
corners_rec, col_min, row_min)
corners_compo += corners_rec
draw.draw_bounding_box(org, corners_compo, show=True)
return corners_compo
def save(index, org, binary, corners_block, corners_img, corners_compo,
compos_class, corners_text):
out_img_gradient = pyjoin(C.ROOT_IMG_GRADIENT, index + '.png')
out_img_draw = pyjoin(C.ROOT_IMG_DRAWN, index + '.png')
out_img_clean = pyjoin(C.ROOT_IMG_CLEAN, index + '.png')
out_label = pyjoin(C.ROOT_LABEL, index + '.json')
# *** Step 7 *** post-processing: remove img elements from original image and segment into smaller size
img_clean = draw.draw_bounding_box(org,
corners_img,
color=(255, 255, 255),
line=-1)
# draw results
draw_bounding = draw.draw_bounding_box_class(org, corners_compo,
compos_class)
draw_bounding = draw.draw_bounding_box_class(
draw_bounding, corners_block,
['block' for i in range(len(corners_block))])
draw_bounding = draw.draw_bounding_box_class(
draw_bounding, corners_img, ['img' for i in range(len(corners_img))])
draw_bounding = draw.draw_bounding_box(draw_bounding, corners_text, line=1)
# save results
cv2.imwrite(out_img_gradient, binary)
cv2.imwrite(out_img_draw, draw_bounding)
cv2.imwrite(out_img_clean, img_clean)
file.save_corners_json(out_label, corners_block,
['div' for i in range(len(corners_block))])
file.save_corners_json(out_label, corners_img,
['div' for i in range(len(corners_img))])
file.save_corners_json(out_label, corners_compo, compos_class)
if is_clip:
file.save_clipping(org, C.ROOT_IMG_COMPONENT, corners_compo,
compos_class, compo_index)
def save(org, binary, corners_block, corners_img, corners_compo, compos_class,
corners_text):
# *** Step 7 *** post-processing: remove img elements from original image and segment into smaller size
# draw results
draw_bounding = draw.draw_bounding_box_class(org, corners_compo,
compos_class)
draw_bounding = draw.draw_bounding_box_class(
draw_bounding, corners_block,
['block' for i in range(len(corners_block))])
draw_bounding = draw.draw_bounding_box_class(
draw_bounding, corners_img, ['img' for i in range(len(corners_img))])
draw_bounding = draw.draw_bounding_box(draw_bounding, corners_text, line=1)
# save results
binary_r = pre.reverse_binary(binary)
cv2.imwrite('output/org.png', org)
cv2.imwrite('output/gradient_r.png', binary_r)
cv2.imwrite('output/gradient.png', binary)
cv2.imwrite('output/labeled.png', draw_bounding)
file.save_corners_json('output/compo.json', corners_block,
['div' for i in range(len(corners_block))])
file.save_corners_json('output/compo.json', corners_img,
['div' for i in range(len(corners_img))])
file.save_corners_json('output/compo.json', corners_compo, compos_class)
if is_clip:
file.save_clipping(org, 'output/components', corners_compo,
compos_class, {})
def processing(org, binary, main=True):
if main:
# *** Step 2 *** object detection: get connected areas -> get boundary -> get corners
boundary_rec, boundary_non_rec = det.boundary_detection(binary,
show=False)
corners_rec = det.get_corner(boundary_rec)
corners_non_rec = det.get_corner(boundary_non_rec)
# *** Step 3 *** data processing: identify blocks and compos from rectangles -> identify irregular compos
corners_block, corners_img, corners_compo = det.block_or_compo(
org, binary, corners_rec)
det.compo_irregular(org, corners_non_rec, corners_img, corners_compo)
corners_img, _ = det.rm_text(org, corners_img,
['img' for i in range(len(corners_img))])
# *** Step 4 *** classification: clip and classify the components candidates -> ignore noises -> refine img
compos = seg.clipping(org, corners_compo)
compos_class = CNN.predict(compos)
# corners_compo, compos_class = det.compo_filter(org, corners_compo, compos_class, is_icon)
corners_compo, compos_class = det.strip_img(corners_compo,
compos_class, corners_img)
# *** Step 5 *** result refinement
if is_merge:
corners_img, _ = det.merge_corner(
corners_img, ['img' for i in range(len(corners_img))])
corners_block, _ = det.rm_text(
org, corners_block, ['block' for i in range(len(corners_block))])
corners_img, _ = det.rm_text(org, corners_img,
['img' for i in range(len(corners_img))])
corners_compo, compos_class = det.rm_text(org, corners_compo,
compos_class)
if is_shrink_img:
corners_img = det.img_shrink(org, binary, corners_img)
# *** Step 6 *** text detection from cleaned image
img_clean = draw.draw_bounding_box(org,
corners_img,
color=(255, 255, 255),
line=-1)
corners_word = ocr.text_detection(org, img_clean)
corners_text = ocr.text_merge_word_into_line(org, corners_word)
# *** Step 7 *** img inspection: search components in img element
if is_img_inspect:
corners_block, corners_img, corners_compo, compos_class = det.compo_in_img(
processing, org, binary, corners_img, corners_block,
corners_compo, compos_class)
return corners_block, corners_img, corners_compo, compos_class, corners_text
# *** used for img inspection ***
# only consider rectangular components
else:
boundary_rec, boundary_non_rec = det.boundary_detection(binary)
corners_rec = det.get_corner(boundary_rec)
corners_block, corners_img, corners_compo = det.block_or_compo(
org, binary, corners_rec)
compos = seg.clipping(org, corners_compo)
compos_class = CNN.predict(compos)
corners_compo, compos_class = det.compo_filter(org, corners_compo,
compos_class, is_icon)
corners_compo, compos_class = det.strip_img(corners_compo,
compos_class, corners_img)
corners_block, _ = det.rm_text(
org, corners_block, ['block' for i in range(len(corners_block))])
corners_compo, compos_class = det.rm_text(org, corners_compo,
compos_class)
return corners_block, corners_compo, compos_class
corners_img += det.img_irregular(org, corners_nonrec)
# ignore too large and highly likely text areas
corners_img = det.img_refine(org, corners_img)
# merge overlapped corners, and remove nested corners
if is_merge_img:
corners_img = det.merge_corners(corners_img)
# detect components on img
corners_compo += det.uicomponent_in_img(org, binary, corners_img)
# remove text misrecognition
corners_block = det.rm_text(org, corners_block)
corners_img = det.rm_text(org, corners_img)
corners_compo = det.rm_text(org, corners_compo)
# *** Step 5 *** text detection from cleaned image
img_clean = draw.draw_bounding_box(org,
corners_img,
color=(255, 255, 255),
line=-1)
if is_ocr:
draw_bounding, word = ocr.text_detection(org, img_clean)
else:
draw_bounding = org
img_clean = draw.draw_bounding_box(img_clean,
corners_compo,
color=(255, 255, 255),
line=-1)
# *** Step 6 *** post-processing: remove img elements from original image and segment into smaller size
if is_segment:
seg.segment_img(img_clean, 600, 'output/segment')
# draw results
draw_bounding = draw.draw_bounding_box_class(
boundary_rec, boundary_all = det.boundary_detection(
binary, C.THRESHOLD_MIN_OBJ_AREA, C.THRESHOLD_MIN_REC_PARAMETER,
C.THRESHOLD_MIN_REC_EVENNESS, C.THRESHOLD_MAX_LINE_THICKNESS,
C.THRESHOLD_MIN_LIN_LENGTH, C.THRESHOLD_MAX_IMG_DENT_RATIO)
# get corner of boundaries -> img or block check
corners_rec = det.get_corner(boundary_rec)
corners_block, corners_img = det.block_or_img(
binary, corners_rec, C.THRESHOLD_MAX_BLOCK_BORDER_THICKNESS,
C.THRESHOLD_MAX_BLOCK_CROSS_POINT)
# refine img component
corners_img = det.img_refine2(corners_img, C.THRESHOLD_MAX_IMG_EDGE_RATIO,
C.THRESHOLD_MUST_IMG_HEIGHT,
C.THRESHOLD_MUST_IMG_WIDTH)
# *** Step 3 *** post-processing: remove img elements from original image and segment into smaller size
img_clean = draw.draw_bounding_box(corners_img, org, (255, 255, 255), -1)
seg.segment_img(img_clean, 600, out_img_segment, 0)
# draw results
draw_bounding = draw.draw_bounding_box(corners_block, org, (0, 255, 0))
draw_bounding = draw.draw_bounding_box(corners_img, draw_bounding,
(0, 0, 255))
# draw_boundary = draw.draw_boundary(boundary_rec, org.shape)
# save results
if is_save:
cv2.imwrite(out_img_draw, draw_bounding)
cv2.imwrite(out_img_gradient, binary)
cv2.imwrite(out_img_clean, img_clean)
file.save_corners(out_label, corners_block, 'div')
file.save_corners(out_label, corners_img, 'img', False)
end = file.timer(start)
# remove pure text element
corners_block = det.rm_text(org, corners_block)
corners_img = det.rm_text(org, corners_img)
corners_compo = det.rm_text(org, corners_compo)
# *** Step 5 *** classification: clip and classify the potential components
if is_classify:
CNN.load()
compos = seg.clipping(org, corners_compo)
compos_classes = CNN.predict(compos)
else:
compos_classes = None
# *** Step 6 *** text detection from cleaned image
img_clean = draw.draw_bounding_box(org,
corners_img,
color=(255, 255, 255),
line=-1)
if is_ocr:
corners_word = ocr.text_detection(org, img_clean)
corners_line = ocr.text_merge_into_line(org, corners_word)
draw_bounding = draw.draw_bounding_box(org, corners_line, line=1)
else:
draw_bounding = org
# *** Step 7 *** post-processing: remove img elements from original image and segment into smaller size
if is_segment:
seg.segment_img(img_clean, 600, 'output/segment')
# draw results
draw_bounding = draw.draw_bounding_box_class(
draw_bounding, corners_block,
['block' for i in range(len(corners_block))])