HYTOS

"""

Testing

pid_images에 테스트 할 도면을 넣고 돌리면 결과를 pid_results에서 확인할 수 있습니다.

pre_train_model path를 입력하고 돌리면 됩니다.

"""

import os

import glob

import argparse

import pickle

import cv2

import numpy as np

from src.utils import *

from src.yolo_net import Yolo

from src.yolo_doftech import YoloD

from PIL import Image

import time

DOFTECH_CLASSES = ['gate', 'globe', 'butterfly', 'check', 'ball', 'relief',

                  '3way_solenoid', 'gate_pressure', 'globe_pressure', 'butterfly_pressure', 'ball_shutoff', 'ball_pressure','ball_motor', 'plug_pressure',

                  'inst', 'func_valve', 'inst_console', 'inst_console_dcs', 'inst_console_sih', 'logic_dcs', 'utility', 'specialty_items', 'logic', 'logic_local_console_dcs',

                  'reducer', 'blind_spectacle_open', 'blind_insertion_open', 'blind_spectacle_close', 'blind_insertion_close',

                  'strainer_basket', 'strainer_conical', 'fitting_capillary_tubing', 'meter_ultrasonic', 'strainer_y', 'tube_pitot',

                  'opc']

class Patch:

    def __init__(self, start_h, start_w, img, check_line):

        self.start_h = start_h

        self.start_w = start_w

        self.img = img

        self.check_line = check_line

        self.obj_num = 0

    def setObjectNum(self, obj_num):

        self.obj_num = obj_num

class Symbol:

    def __init__(self, start_h, start_w, end_h, end_w, iou, class_info):

        self.start_h = start_h

        self.start_w = start_w

        self.end_h = end_h

        self.end_w = end_w

        self.iou = iou

        self.class_info = class_info

    def to_stream(self):

        # y, x, height, width : member name is incorrect

        #return [self.class_info, self.start_w, self.start_h, self.end_w - self.start_w, self.end_h - self.start_h, self.iou]

        # x, y, width, height

        return [self.class_info, self.start_h, self.start_w, self.end_h - self.start_h, self.end_w - self.start_w, self.iou]

def get_args():

    parser = argparse.ArgumentParser("You Only Look Once: Unified, Real-Time Object Detection")

    parser.add_argument("--image_size", type=int, default=448, help="The common width and height for all images")

    parser.add_argument("--conf_threshold", type=float, default=0.5)

    parser.add_argument("--nms_threshold", type=float, default=0.6)

    parser.add_argument("--pre_trained_model_type", type=str, choices=["model", "params"], default="params")

    parser.add_argument("--pre_trained_model_path", type=str, default="E:\Projects\DTIPID_GERRIT\DTI_PID\WebServer\symbol_recognition\MODEL/f_doftech_all_class_only_params.pth")

    parser.add_argument("--pre_trained_model_path2", type=str, default="E:\Projects\DTIPID_GERRIT\DTI_PID\WebServer\symbol_recognition\MODEL/doftech_all_class_only_params_opc.pth")

    #parser.add_argument("--pre_trained_model_path3", type=str, default="trained_models/only_params_trained_yolo_voc")

    parser.add_argument("--input", type=str, default="test_images")

    parser.add_argument("--output", type=str, default="result_images")

    parser.add_argument("--data_path_test", type=str, default="E:\Projects\DTIPID_GERRIT\DTI_PID\WebServer\symbol_recognition\pid_images/", help="the root folder of dataset") #Doftech VOCdevkit

    args = parser.parse_args()

    return args

def get_symbol(imgs, trained_model1=None, trained_model2=None):

    global colors

    parser = argparse.ArgumentParser("You Only Look Once: Unified, Real-Time Object Detection")

    parser.add_argument("--image_size", type=int, default=448, help="The common width and height for all images")

    parser.add_argument("--conf_threshold", type=float, default=0.5)

    parser.add_argument("--nms_threshold", type=float, default=0.6)

    parser.add_argument("--pre_trained_model_type", type=str, choices=["model", "params"], default="params")

    parser.add_argument("--pre_trained_model_path", type=str, default=trained_model1)

    parser.add_argument("--pre_trained_model_path2", type=str, default=trained_model2)

    parser.add_argument("--data_path_test", type=str, default="E:\Projects\DTIPID_GERRIT\DTI_PID\WebServer\symbol_recognition\pid_images/", help="the root folder of dataset") #Doftech VOCdevkit

    opt = parser.parse_args()

    if torch.cuda.is_available():

        model1 = Yolo(35).cuda()

        model1.load_state_dict(torch.load(opt.pre_trained_model_path))

        model2 = Yolo(35).cuda()

        model2.load_state_dict(torch.load(opt.pre_trained_model_path2))

    colors = pickle.load(open(os.path.dirname(os.path.realpath(__file__)) + "/src/pallete", "rb"))

    img_list = imgs

    start = time.time()

    total_symbole_lists = []

    # ----------------------------------------------

    # Get Patch arguments : (img list, patch_size, overlap_size)

    # ----------------------------------------------

    for idx_img in range(len(img_list)):

        print("=========> CROP PATCH")

        small_object_patch_list = []

        large_object_patch_list = []

        total_symbole_list = []

        s_image = img_list[idx_img].copy()

        l_image = img_list[idx_img].copy()

        small_object_patch_list.append(get_patch(s_image, 500, 250))

        large_object_patch_list.append(get_patch(l_image, 800, 200))

        img_name = str(idx_img)

        print("=========> "+ img_name)

        save_dir = 'E:\Projects\DTIPID_GERRIT\DTI_PID\WebServer\symbol_recognition\save/' + img_name + '/'

        if not os.path.isdir(save_dir):

            os.mkdir(save_dir)

            os.mkdir(save_dir+"a/")

            os.mkdir(save_dir+"b/")

        text_file = open(save_dir+'test_result.txt', mode='wt', encoding='utf-8')

        text_file.write(img_name+" -> Detection Reesult\n")

        text_file.write("================================\n")

            # Text File 만들기

        # ----------------------------------------------

        # Small Object Detection (Valve, Sensor, etc.)

        # ---------------------------------------------

        print("=========> SMALL OBJECT DETECTION")

        for idx_small in range(len(small_object_patch_list)):

            patchs = small_object_patch_list[idx_small]

            total_symbole_list.append(detection_object(patchs, model1, save_dir+"a/", opc=False, opt=opt))

        # ----------------------------------------------

        # Large Object Detection (OPC etc.)

        # ----------------------------------------------

        print("=========> LARGE OBJECT DETECTION")

        for idx_large in range(len(large_object_patch_list)):

            patchs = large_object_patch_list[idx_large]

            total_symbole_list.append(detection_object(patchs, model2, save_dir+"b/", opc=True, opt=opt))

        t_image = img_list[idx_img].copy()

        count_result = merge_fn(t_image, total_symbole_list, save_dir)

        for idx, value in enumerate(count_result):

            text_file.write(DOFTECH_CLASSES[idx]+ " : "+ str(value) + "\n")

        text_file.close()

        res = []

        for symbol in total_symbole_list[0] + total_symbole_list[1]:

            res.append(symbol.to_stream())

        total_symbole_lists.append(res)

    print("time :", time.time() - start)  # 현재시각 - 시작시간 = 실행 시간

    return total_symbole_lists

def test(opt):

    global colors

    if torch.cuda.is_available():

        model1 = Yolo(35).cuda()

        model1.load_state_dict(torch.load(opt.pre_trained_model_path))

        model2 = Yolo(35).cuda()

        model2.load_state_dict(torch.load(opt.pre_trained_model_path2))

    else:

        return None

    colors = pickle.load(open("E:\Projects\DTIPID_GERRIT\DTI_PID\WebServer\symbol_recognition/src/pallete", "rb"))

    print(opt.input)

    img_list = sorted(glob.glob(os.path.join(opt.data_path_test, '*.png')))

    start = time.time()

    print(img_list)

    # ----------------------------------------------

    # Get Patch arguments : (img list, patch_size, overlap_size)

    # ----------------------------------------------

    for idx_img in range(len(img_list)):

        print("=========> CROP PATCH")

        small_object_patch_list = []

        large_object_patch_list = []

        total_symbole_list = []

        small_object_patch_list.append(get_patch(img_list[idx_img], 500, 250))

        large_object_patch_list.append(get_patch(img_list[idx_img], 800, 200))

        img_name = img_list[idx_img].split('\\')[1].split('.')[0]

        print("=========> "+ img_name)

        save_dir = 'E:\Projects\DTIPID_GERRIT\DTI_PID\WebServer\symbol_recognition\save/' + img_name + '/'

        if not os.path.isdir(save_dir):

            os.mkdir(save_dir)

            os.mkdir(save_dir+"a/")

            os.mkdir(save_dir+"b/")

        text_file = open(save_dir+'test_result.txt', mode='wt', encoding='utf-8')

        text_file.write(img_name+" -> Detection Reesult\n")

        text_file.write("================================\n")

            # Text File 만들기

        # ----------------------------------------------

        # Small Object Detection (Valve, Sensor, etc.)

        # ---------------------------------------------

        print("=========> SMALL OBJECT DETECTION")

        for idx_small in range(len(small_object_patch_list)):

            patchs = small_object_patch_list[idx_small]

            total_symbole_list.append(detection_object(patchs, model1, save_dir+"a/", opc=False, opt=opt))

        # ----------------------------------------------

        # Large Object Detection (OPC etc.)

        # ----------------------------------------------

        print("=========> LARGE OBJECT DETECTION")

        for idx_large in range(len(large_object_patch_list)):

            patchs = large_object_patch_list[idx_large]

            total_symbole_list.append(detection_object(patchs, model2, save_dir+"b/", opc=True, opt=opt))

        count_result = merge_fn(img_list[idx_img], total_symbole_list, save_dir)

        for idx, value in enumerate(count_result):

            text_file.write(DOFTECH_CLASSES[idx]+ " : "+ str(value) + "\n")

        text_file.close()

    print("time :", time.time() - start)  # 현재시각 - 시작시간 = 실행 시간

def detection_object(patchs, model, save_root, opc, opt):

    symbol_list = []

    for idx in range(len(patchs)):

        pil_image = patchs[idx]

        # pil_image.img.show()

        np_img = np.asarray(pil_image.img)

        image = cv2.cvtColor(np_img, cv2.COLOR_RGB2BGR)

        height, width = image.shape[:2]

        image = cv2.resize(image, (opt.image_size, opt.image_size))

        image = np.transpose(np.array(image, dtype=np.float32), (2, 0, 1))

        image = image[None, :, :, :]

        width_ratio = float(opt.image_size) / width

        height_ratio = float(opt.image_size) / height

        data = Variable(torch.FloatTensor(image))

        if torch.cuda.is_available():

            data = data.cuda()

        with torch.no_grad():

            logits = model(data)

            predictions = post_processing(logits, opt.image_size, DOFTECH_CLASSES, model.anchors, opt.conf_threshold,

                                          opt.nms_threshold)

        if len(predictions) != 0:

            predictions = predictions[0]

            output_image = cv2.cvtColor(np.array(pil_image.img), cv2.COLOR_RGB2BGR)

            for pred in predictions:

                if opc == True :

                    if pred[5] == "opc" and pred[4] > 0.4:  # Classification threshold

                        xmin = int(max(pred[0] / width_ratio, 0))

                        ymin = int(max(pred[1] / height_ratio, 0))

                        xmax = int(min((pred[0] + pred[2]) / width_ratio, width))

                        ymax = int(min((pred[1] + pred[3]) / height_ratio, height))

                        color = colors[DOFTECH_CLASSES.index(pred[5])]

                        symbol = Symbol(xmin + pil_image.start_w, ymin + pil_image.start_h, xmax + pil_image.start_w,

                                        ymax + pil_image.start_h, pred[4], pred[5])  # To save symbole information

                        check = check_duplication(symbol_list, symbol, True)

                        if check is True:

                            symbol_list.append(symbol)

                        cv2.rectangle(output_image, (xmin, ymin), (xmax, ymax), color, 2)

                        text_size = cv2.getTextSize(pred[5] + ' : %.2f' % pred[4], cv2.FONT_HERSHEY_PLAIN, 1, 1)[0]

                        cv2.rectangle(output_image, (xmin, ymin), (xmin + text_size[0] + 3, ymin + text_size[1] + 4),

                                      color,-1)

                        cv2.putText(

                            output_image, pred[5] + ' : %.2f' % pred[4],

                            (xmin, ymin + text_size[1] + 4), cv2.FONT_HERSHEY_PLAIN, 1,

                            (255, 255, 255), 1)

                else :

                    if pred[4] > 0.1:  # Classification threshold

                        xmin = int(max(pred[0] / width_ratio, 0))

                        ymin = int(max(pred[1] / height_ratio, 0))

                        xmax = int(min((pred[0] + pred[2]) / width_ratio, width))

                        ymax = int(min((pred[1] + pred[3]) / height_ratio, height))

                        color = colors[DOFTECH_CLASSES.index(pred[5])]

                        symbol = Symbol(xmin + pil_image.start_w, ymin + pil_image.start_h, xmax + pil_image.start_w,

                                        ymax + pil_image.start_h, pred[4], pred[5])  # To save symbole information

                        check = check_duplication(symbol_list, symbol)

                        if check is True:

                            symbol_list.append(symbol)

                        cv2.rectangle(output_image, (xmin, ymin), (xmax, ymax), color, 2)

                        text_size = cv2.getTextSize(pred[5] + ' : %.2f' % pred[4], cv2.FONT_HERSHEY_PLAIN, 1, 1)[0]

                        cv2.rectangle(output_image, (xmin, ymin), (xmin + text_size[0] + 3, ymin + text_size[1] + 4), color,

                                      -1)

                        cv2.putText(

                            output_image, pred[5] + ' : %.2f' % pred[4],

                            (xmin, ymin + text_size[1] + 4), cv2.FONT_HERSHEY_PLAIN, 1,

                            (255, 255, 255), 1)

            # -----------------------fill king image--------------------------------------

            cv2.imwrite(save_root+"{}_prediction.png".format(idx), output_image)

            pil_image.setObjectNum(len(predictions))

    return symbol_list

def merge_fn(total_img, total_symbole_list, save_root):

    """ total_img:path -> img file """

    print("=========> MERGE RESULT")

    if type(total_img) is str:

        total_img = Image.open(total_img)

    king_image = total_img.convert('RGB')

    temp_king_img = np.array(king_image)#Image.new("RGB", (king_width, king_height))

    temp_king_img = temp_king_img[:,:,::-1].copy()

    count_result = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]

    for idx_out in range(len(total_symbole_list)):

        for idx in range(len(total_symbole_list[idx_out])):

            color = colors[DOFTECH_CLASSES.index(total_symbole_list[idx_out][idx].class_info)]

            cv2.rectangle(temp_king_img, (total_symbole_list[idx_out][idx].start_h, total_symbole_list[idx_out][idx].start_w), (total_symbole_list[idx_out][idx].end_h, total_symbole_list[idx_out][idx].end_w), color, 2)

            text_size = cv2.getTextSize(total_symbole_list[idx_out][idx].class_info + ' : %.2f' % total_symbole_list[idx_out][idx].iou, cv2.FONT_HERSHEY_PLAIN, 1, 1)[0]

            cv2.rectangle(temp_king_img, (total_symbole_list[idx_out][idx].start_h, total_symbole_list[idx_out][idx].start_w), (total_symbole_list[idx_out][idx].start_h + text_size[0] + 3,                                                                                                total_symbole_list[idx_out][idx].start_w + text_size[1] + 4), color, -1)

            cv2.putText(

                temp_king_img, total_symbole_list[idx_out][idx].class_info + ' : %.2f' % total_symbole_list[idx_out][idx].iou,

                (total_symbole_list[idx_out][idx].start_h, total_symbole_list[idx_out][idx].start_w+ text_size[1] + 4), cv2.FONT_HERSHEY_PLAIN, 1,

                (255, 255, 255), 1)

            # COUNTING

            for index, name in enumerate(DOFTECH_CLASSES):

                if name == total_symbole_list[idx_out][idx].class_info :

                    count_result[index] = count_result[index]+1

    cv2.imwrite(save_root+"entire.png", temp_king_img)

    return count_result

def get_patch(image, patch_size, overlap_size):

    """ img_path -> img """

    crop_list = []

    if type(image) is str:

        image = Image.open(image)

    (img_h, img_w) = image.size

    # crop 할 사이즈 : grid_w, grid_h

    grid_w = patch_size  # crop width

    grid_h = patch_size  # crop

    range_w = (int)(img_w / grid_w)

    range_h = (int)(img_h / grid_h)

    # print(range_w, range_h)

    repeat_num_w = (int)(grid_w / overlap_size)

    repeat_num_h = (int)(grid_h / overlap_size)

    # print(repeat_num_w, repeat_num_h)

    for w in range(range_w * repeat_num_w):

        for h in range(range_h * repeat_num_h):

            if h == (range_h * repeat_num_h) - 1 and w == (range_w * repeat_num_w) - 1:

                start_h = img_h - grid_h

                start_w = img_w - grid_w

                end_h = img_h

                end_w = img_w

                bbox = (start_h, start_w, end_h, end_w)

                crop_img = image.crop(bbox)

                patch = Patch(start_w, start_h, crop_img, True)

                crop_list.append(patch)

            elif (h * overlap_size) + (grid_h) >= img_h:

                start_h = img_h - grid_h  # OK

                start_w = img_w - grid_w

                end_h = img_h

                end_w = start_w + grid_w

                bbox = (start_h, start_w, end_h, end_w)

                crop_img = image.crop(bbox)

                patch = Patch(start_w, start_h, crop_img, True)

                crop_list.append(patch)

            elif (w * overlap_size) + (grid_w) >= img_w:

                start_h = img_h - grid_h

                start_w = img_w - grid_w

                end_h = start_h + grid_h

                end_w = img_w

                bbox = (start_h, start_w, end_h, end_w)

                crop_img = image.crop(bbox)

                patch = Patch(start_w, start_h, crop_img, True)

                crop_list.append(patch)

            else:

                bbox = (h * overlap_size, w * overlap_size, (h * overlap_size) + (grid_h),

                        (w * overlap_size) + (grid_w))  # 좌, 상, 우, 하

                crop_img = image.crop(bbox)

                patch = Patch(w * overlap_size, h * overlap_size, crop_img, False)

                crop_list.append(patch)

            # print(bbox)

            # patch.img.save('./pid_results/_{}.png'.format(len(crop_list)))

    return crop_list

def check_duplication(symbol_list, check_symbol, opc=False):

    if check_symbol.class_info == "strainer_basket":

        threshold = 30

    elif check_symbol.class_info == "3way_solenoid":

        threshold = 30

    else :

        threshold = 20

    for idx in range(len(symbol_list)):

        if (abs(symbol_list[idx].start_h - check_symbol.start_h) <= threshold):

            if (abs(symbol_list[idx].start_w - check_symbol.start_w) <= threshold):

                if (symbol_list[idx].iou < check_symbol.iou):

                    symbol_list[idx] = check_symbol

                    return False

                else:

                    return False

    return True

if __name__ == "__main__":

    opt = get_args()

    test(opt)