Spaces:
Running
Running
| 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 | |