mirror of
https://github.com/leigest519/ScreenCoder.git
synced 2026-02-13 18:22:50 +00:00
128 lines
5.5 KiB
Python
128 lines
5.5 KiB
Python
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 issues among UI components
|
|
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
|