Spaces:
Sleeping
Sleeping
File size: 8,913 Bytes
f1554a2 1c42b13 f1554a2 1c42b13 f1554a2 | 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 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 | import cv2
from os.path import join as pjoin
import time
import CDM.detect_compo.lib_ip.ip_preprocessing as pre
import CDM.detect_compo.lib_ip.ip_draw as draw
import CDM.detect_compo.lib_ip.ip_detection as det
import CDM.detect_compo.lib_ip.file_utils as file
import CDM.detect_compo.lib_ip.Component as Compo
from CDM.config.CONFIG_UIED import Config
C = Config()
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.time()
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']))
full_size_org, full_size_grey = pre.read_img(input_img_path)
ratio = full_size_org.shape[0] / org.shape[0]
# *** 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(full_size_org, ratio, uicompos, show=show, name='merged compo', write_path=pjoin(ip_root, name + '.jpg'), wait_key=wai_key)
# # classify icons
# model = models.resnet18().to('cpu')
# in_feature_num = model.fc.in_features
# model.fc = nn.Linear(in_feature_num, 99)
# # model.conv1 = nn.Conv2d(in_channels=1, out_channels=64, kernel_size=(3,3), padding=(3,3), stride=(2,2), bias=False)
# model.conv1 = nn.Conv2d(in_channels=1, out_channels=64, kernel_size=(5, 5), padding=(3, 3), stride=(2, 2),
# bias=False)
# # PATH = "C:/ANU/2022 s2/honours project/code/UIED-master/model/model-99-resnet18.pkl"
# PATH = "./model/model-99-resnet18.pkl"
# # trained_model = model()
# model.load_state_dict(torch.load(PATH, map_location=torch.device('cpu')))
#
# model.eval()
#
# # ----------------- try on semantics dataset---------------------
#
# # sample_data = np.load('C:/ANU/2022 s2/honours project/code/semantic-icon-classifier-master/data/training_x.npy')
# #
# # array = np.reshape(sample_data[0, :, :, :], [32, 32])
# #
# # print("array: ", array)
# #
# # cv2.imshow("array", array)
# # cv2.waitKey(0)
# #
# # array = array.astype('float32')
# # array = array / 255
# # array = (array - array.mean()) / array.std()
# #
# # print("array mean: ", array.mean())
# # print("array std: ", array.std())
# #
# # array = array.reshape(1, 1, 32, 32)
# #
# # array = torch.tensor(array)
# # print("array_tensor: ", array)
# # array_pred_label = model(array)
# # print("output: ", array_pred_label)
#
# # ----------------- end trying ---------------------
#
# grey = grey.astype('float32')
# grey = grey / 255
# # grey = grey / np.linalg.norm(grey)
#
# grey = (grey-grey.mean())/grey.std()
# print("grey mean: ", grey.mean())
# print("grey std: ", grey.std())
#
# # grey = grey.to(torch.float32)
#
# # plt.imshow(Image.fromarray(binary))
# # plt.show()
# # cv2.imshow("grey", grey)
#
# privacy_compos = []
# for comp in uicompos:
#
# # cv2.imshow("comp", grey[comp.bbox.row_min:comp.bbox.row_max, comp.bbox.col_min:comp.bbox.col_max])
# # cv2.waitKey(0)
#
# # col_mid = int((comp.bbox.col_min+comp.bbox.col_max)/2)
# # row_mid = int((comp.bbox.row_min+comp.bbox.row_max)/2)
# # comp_crop = grey[max(0, row_mid-16):min(grey.shape[1], row_mid+16), max(0, col_mid-16):min(grey.shape[0], col_mid+16)]
# #
# # if comp_crop.shape[0] != 32 or comp_crop.shape[1] != 32:
# # print("A component is not classified, size: ", comp_crop.shape)
# # print("col_mid: ", col_mid)
# # print("row_mid: ", row_mid)
# # print("shape[0]: ", comp_crop.shape[0])
# # print("shape[1]: ", comp_crop.shape[1])
# # print("max(0, row_mid-16) and min(binary.shape[1], row_mid+16): ", max(0, row_mid-16), min(grey.shape[1], row_mid+16))
#
# comp_grey = grey[comp.bbox.row_min:comp.bbox.row_max, comp.bbox.col_min:comp.bbox.col_max]
#
# # cv2.imshow("comp_grey", comp_grey)
# # cv2.waitKey(0)
#
# # print("comp_crop: ", comp_crop)
# # comp_crop = comp_grey.reshape(1, 1, 32, 32)
# comp_crop = cv2.resize(comp_grey, (32, 32))
# print("comp_crop: ", comp_crop)
#
# # cv2.imshow("comp_crop", comp_crop)
# # cv2.waitKey(0)
#
# comp_crop = comp_crop.reshape(1, 1, 32, 32)
#
# comp_tensor = torch.tensor(comp_crop)
# comp_tensor = comp_tensor.permute(0, 1, 3, 2)
# print("comp_tensor: ", comp_tensor)
# # comp_float = comp_tensor.to(torch.float32)
# # print("comp_float: ", comp_float)
# # pred_label = model(comp_float)
# pred_label = model(comp_tensor)
# print("output: ", pred_label)
# print("label: ", np.argmax(pred_label.cpu().data.numpy(), axis=1))
# if np.argmax(pred_label.cpu().data.numpy(), axis=1) in [72.0, 42.0, 77.0, 91.0, 6.0, 89.0, 40.0, 43.0, 82.0, 3.0, 68.0,
# 49.0, 56.0, 89.0]:
# privacy_compos.append(comp)
#
# draw.draw_bounding_box(org, privacy_compos, 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)
file.save_corners_json(pjoin(ip_root, name + '.json'), uicompos)
# file.save_corners_json(pjoin(ip_root, name + '.json'), uicompos, full_size_org, ratio)
cd_time = time.time() - start
print("[Compo Detection Completed in %.3f s] Input: %s Output: %s" % (cd_time, input_img_path, pjoin(ip_root, name + '.json')))
return cd_time
|