Spaces:
Running
Running
File size: 5,601 Bytes
a383d0e 0246ff9 a383d0e 0246ff9 a383d0e |
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 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 |
import cv2
from os.path import join as pjoin
import time
import json
import numpy as np
import detect_compo.lib_ip.ip_preprocessing as pre
import detect_compo.lib_ip.ip_draw as draw
import detect_compo.lib_ip.ip_detection as det
import detect_compo.lib_ip.file_utils as file
import detect_compo.lib_ip.Component as Compo
from config.CONFIG_UIED import Config
C = Config()
def resolve_uicompo_containment(uicompos):
"""
Resolves containment issues among UI components.
If a component's bounding box is fully contained within another's, it is removed.
"""
def contains(bbox_a, bbox_b):
"""Checks if bbox_a completely contains bbox_b."""
return bbox_a.col_min <= bbox_b.col_min and \
bbox_a.row_min <= bbox_b.row_min and \
bbox_a.col_max >= bbox_b.col_max and \
bbox_a.row_max >= bbox_b.row_max
compos_to_remove = set()
for i, compo1 in enumerate(uicompos):
for j, compo2 in enumerate(uicompos):
if i == j:
continue
# Check if compo1 contains compo2
if contains(compo1.bbox, compo2.bbox):
compos_to_remove.add(j)
# Filter out the contained components
final_compos = [compo for i, compo in enumerate(uicompos) if i not in compos_to_remove]
if len(final_compos) < len(uicompos):
print(f"Containment resolved: Removed {len(uicompos) - len(final_compos)} contained components.")
return final_compos
def nesting_inspection(org, grey, compos, ffl_block):
'''
Inspect all big compos through block division by flood-fill
:param ffl_block: gradient threshold for flood-fill
:return: nesting compos
'''
nesting_compos = []
for i, compo in enumerate(compos):
if compo.height > 50:
replace = False
clip_grey = compo.compo_clipping(grey)
n_compos = det.nested_components_detection(clip_grey, org, grad_thresh=ffl_block, show=False)
Compo.cvt_compos_relative_pos(n_compos, compo.bbox.col_min, compo.bbox.row_min)
for n_compo in n_compos:
if n_compo.redundant:
compos[i] = n_compo
replace = True
break
if not replace:
nesting_compos += n_compos
return nesting_compos
def compo_detection(input_img_path, output_root, uied_params,
resize_by_height=800, classifier=None, show=False, wai_key=0):
start = time.perf_counter()
name = input_img_path.split('/')[-1][:-4] if '/' in input_img_path else input_img_path.split('\\')[-1][:-4]
ip_root = file.build_directory(pjoin(output_root, "ip"))
# *** Step 1 *** pre-processing: read img -> get binary map
org, grey = pre.read_img(input_img_path, resize_by_height)
binary = pre.binarization(org, grad_min=int(uied_params['min-grad']))
# *** Step 2 *** element detection
det.rm_line(binary, show=show, wait_key=wai_key)
uicompos = det.component_detection(binary, min_obj_area=int(uied_params['min-ele-area']))
# *** Step 3 *** results refinement
uicompos = det.compo_filter(uicompos, min_area=int(uied_params['min-ele-area']), img_shape=binary.shape)
uicompos = det.merge_intersected_compos(uicompos)
det.compo_block_recognition(binary, uicompos)
if uied_params['merge-contained-ele']:
uicompos = det.rm_contained_compos_not_in_block(uicompos)
Compo.compos_update(uicompos, org.shape)
Compo.compos_containment(uicompos)
# *** Step 4 ** nesting inspection: check if big compos have nesting element
uicompos += nesting_inspection(org, grey, uicompos, ffl_block=uied_params['ffl-block'])
Compo.compos_update(uicompos, org.shape)
draw.draw_bounding_box(org, uicompos, show=show, name='merged compo', write_path=pjoin(ip_root, name + '.jpg'), wait_key=wai_key)
# *** Step 5 *** image inspection: recognize image -> remove noise in image -> binarize with larger threshold and reverse -> rectangular compo detection
# if classifier is not None:
# classifier['Image'].predict(seg.clipping(org, uicompos), uicompos)
# draw.draw_bounding_box_class(org, uicompos, show=show)
# uicompos = det.rm_noise_in_large_img(uicompos, org)
# draw.draw_bounding_box_class(org, uicompos, show=show)
# det.detect_compos_in_img(uicompos, binary_org, org)
# draw.draw_bounding_box(org, uicompos, show=show)
# if classifier is not None:
# classifier['Noise'].predict(seg.clipping(org, uicompos), uicompos)
# draw.draw_bounding_box_class(org, uicompos, show=show)
# uicompos = det.rm_noise_compos(uicompos)
# *** Step 6 *** element classification: all category classification
# if classifier is not None:
# classifier['Elements'].predict([compo.compo_clipping(org) for compo in uicompos], uicompos)
# draw.draw_bounding_box_class(org, uicompos, show=show, name='cls', write_path=pjoin(ip_root, 'result.jpg'))
# draw.draw_bounding_box_class(org, uicompos, write_path=pjoin(output_root, 'result.jpg'))
# *** Step 7 *** save detection result
Compo.compos_update(uicompos, org.shape)
# *** Step 8 *** Resolve containment before saving
uicompos = resolve_uicompo_containment(uicompos)
file.save_corners_json(pjoin(ip_root, name + '.json'), uicompos)
print("[Compo Detection Completed in %.3f s] Input: %s Output: %s" % (time.perf_counter() - start, input_img_path, pjoin(ip_root, name + '.json')))
return uicompos
|