CDM/detect_compo/ip_region_proposal.py

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