VerificationCode/puzzle_detect/pre_process.py

import copy
import json
import os
import random
import re
import sys
import time
import traceback
from glob import glob
from itertools import combinations, product

import chardet
import cv2
import jieba
import numpy as np
from PIL import ImageFont, ImageDraw, Image
from captcha.image import ImageCaptcha
from keras_preprocessing.sequence import pad_sequences
from matplotlib.colors import rgb_to_hsv, hsv_to_rgb

sys.path.append(os.path.dirname(os.path.abspath(__file__)) + "/../")
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
from click_captcha.utils import np2pil, pil2np, pil_resize, pil_rotate, pil2np_a, np2pil_a, pil_resize_a, pil_rotate_a


def gen_siamese(paths, batch_size=32, shape=(40, 40), cls_num=1):
    num = len(paths)
    data_path = os.path.dirname(os.path.abspath(__file__)) + "/../data/click/"

    i = 0
    while True:
        if i >= num:
            i = 0
            random.shuffle(paths)

        height, width = shape[:2]
        X1 = np.zeros((batch_size, height, width, 1))
        X2 = np.zeros((batch_size, height, width, 1))
        Y = np.zeros((batch_size, 2))

        for j in range(batch_size):
            # 生成标注数据
            img1, img2, label = paths[i][:-1].split("\t")
            # print(img1, img2, label)
            img1 = cv2.imread(data_path + img1)
            img1 = pil_resize(img1, shape[0], shape[1])
            img1 = cv2.cvtColor(img1, cv2.COLOR_BGR2GRAY)
            img1 = np.expand_dims(img1, axis=-1)
            img2 = cv2.imread(data_path + img2)
            img2 = pil_resize(img2, shape[0], shape[1])
            img2 = cv2.cvtColor(img2, cv2.COLOR_BGR2GRAY)
            img2 = np.expand_dims(img2, axis=-1)
            if label == "1":
                label = np.array([0, 1])
            else:
                label = np.array([1, 0])
            X1[j] = img1
            X2[j] = img2
            Y[j] = label

        yield {"input_1": X1, "input_2": X2}, {"output": Y}


def gen_mobile(paths, batch_size=32, shape=(40, 40), cls_num=5710, data_path="click"):
    num = len(paths)
    data_path = os.path.dirname(os.path.abspath(__file__)) + "/../data/" + data_path + "/"

    i = 0
    random.shuffle(paths)
    while True:
        if i >= num:
            i = 0
            random.shuffle(paths)

        height, width = shape[:2]
        if len(shape) > 2:
            channel = 3
        else:
            channel = 1
        X = np.zeros((batch_size, height, width, channel))
        Y = np.zeros((batch_size, cls_num))

        j = 0
        error_num = 0
        while j < batch_size:
        # for j in range(batch_size):
            if i >= num:
                random.shuffle(paths)
                i = 0
            path = paths[i].split(os.sep)[-1]
            char_index = int(path.split("_")[0])
            label = np.zeros(cls_num)
            # print("char_index", char_index)
            label[char_index] = 1
            # print("label", np.argmax(label), char_index)


            img1 = cv2.imread(data_path + path)
            img1 = pil_resize(img1, shape[0], shape[1])
            img1 = img1 / 255.
            # img1 = cv2.cvtColor(img1, cv2.COLOR_BGR2GRAY)
            # img1 = np.expand_dims(img1, axis=-1)

            X[j] = img1
            Y[j] = label
            i += 1
            j += 1
        # print("error_num", error_num)
        yield X, Y


def gen_yolo_char(paths, batch_size, input_shape, anchors, num_classes, box_num=6):
    """data generator for fit_generator"""
    n = len(paths)
    data_path = os.path.dirname(os.path.abspath(__file__)) + "/../data/detect/"
    i = 0
    while True:
        image_data = []
        box_data = []

        batch_cnt = 0
        while batch_cnt < batch_size:
            try:
                if i == 0:
                    np.random.shuffle(paths)

                ss = paths[i][:-1].split(" ")
                image_path = ss[0]
                image = cv2.imread(data_path+image_path)
                image = pil_resize(image, input_shape[0], input_shape[1])
                image_show = copy.deepcopy(image)
                image = image / 255.
                box = np.array([np.array(list(map(int, box.split(',')))) for box in ss[1:]])

                # box数不同，复制
                if box.shape[0] < box_num:
                    box = np.concatenate([box, box[:2, :]], axis=0)

                # show
                # box_show = box.tolist()
                # for b in box_show:
                #     print("box", b)
                #     cv2.rectangle(image_show, (b[0], b[1]), (b[2], b[3]), (255, 0, 0), 2)
                # cv2.imshow("image_show", image_show)
                # cv2.waitKey(0)

                image_data.append(image)
                box_data.append(box)
                i = (i+1) % n
                batch_cnt += 1
            except:
                i = (i+1) % n
                continue

            # print
            # print(image.shape)
            # image_show = (image*255).astype(np.uint8)
            # print("annotation_lines[i]", annotation_lines[i])
            # for _b in box:
            #     print(_b)
            #     cv2.rectangle(image_show, (int(_b[0]), int(_b[1])), (int(_b[2]), int(_b[3])), (0, 255, 0), 1)
            # cv2.imshow("image", image_show)
            # cv2.waitKey(0)

        image_data = np.array(image_data)
        box_data = np.array(box_data)
        # print(image_data.shape, box_data.shape)
        y_true = preprocess_true_boxes(box_data, input_shape, anchors, num_classes)
        yield [image_data, *y_true], np.zeros(batch_size)


def gen_yolo_puzzle(paths, batch_size, input_shape, anchors, num_classes, box_num=1):
    """data generator for fit_generator"""
    n = len(paths)
    data_path = os.path.dirname(os.path.abspath(__file__)) + "/../data/detect2/"
    i = 0
    while True:
        image_data = []
        box_data = []

        batch_cnt = 0
        while batch_cnt < batch_size:
            try:
                if i == 0:
                    np.random.shuffle(paths)

                ss = paths[i][:-1].split(" ")
                image_path = ss[0]
                image = cv2.imread(data_path+image_path)
                image = pil_resize(image, input_shape[0], input_shape[1])
                image_show = copy.deepcopy(image)
                image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
                image = 255. - image
                image = np.uint8(image)
                # cv2.imshow("image", image)
                # cv2.waitKey(0)
                image = np.expand_dims(image, -1)
                image = image / 255.
                box = np.array([np.array(list(map(int, box.split(',')))) for box in ss[1:]])

                # box数不同，复制
                if box.shape[0] < box_num:
                    box = np.concatenate([box, box[:2, :]], axis=0)

                # show
                # box_show = box.tolist()
                # for b in box_show:
                #     print("box", b)
                #     cv2.rectangle(image_show, (b[0], b[1]), (b[2], b[3]), (0, 0, 255), 2)
                # cv2.imshow("image_show", image_show)
                # cv2.waitKey(0)

                image_data.append(image)
                box_data.append(box)
                i = (i+1) % n
                batch_cnt += 1
            except:
                i = (i+1) % n
                continue

            # print
            # print(image.shape)
            # image_show = (image*255).astype(np.uint8)
            # for _b in box:
            #     print(_b)
            #     cv2.rectangle(image_show, (int(_b[0]), int(_b[1])), (int(_b[2]), int(_b[3])), (0, 255, 0), 1)
            # cv2.imshow("image", image_show)
            # cv2.waitKey(0)

        image_data = np.array(image_data)
        box_data = np.array(box_data)
        # print(image_data.shape, box_data.shape)
        y_true = preprocess_true_boxes(box_data, input_shape, anchors, num_classes)
        yield [image_data, *y_true], np.zeros(batch_size)


def gen_drag(paths, batch_size=32, shape=(128, 256), cls_num=2):
    num = len(paths)
    data_path = os.path.dirname(os.path.abspath(__file__)) + "/../data/drag/"

    map_path = data_path+"map.txt"
    with open(map_path, "r") as f:
        _list = f.readlines()
    map_dict = {}
    for s in _list:
        ss = s[:-1].split(" ")
        map_dict[ss[0]] = ss[1]

    i = 0
    random.shuffle(paths)
    while True:
        if i >= num:
            i = 0
            random.shuffle(paths)

        height, width = shape[:2]
        if len(shape) > 2:
            channel = 3
        else:
            channel = 1
        X = np.zeros((batch_size, height, width, channel))
        Y = np.zeros((batch_size, height, width, 1))

        for j in range(batch_size):
            if i >= num:
                random.shuffle(paths)
                i = 0
            path = paths[i].split(os.sep)[-1]
            w_index = int(map_dict.get(path))
            # label = np.zeros(cls_num)
            # print("char_index", char_index)
            # label[w_index] = 1
            # print("label", np.argmax(label), char_index)

            img1 = cv2.imread(data_path + path)
            img1 = pil_resize(img1, shape[0], shape[1])
            # cv2.imshow("image", img1)
            label = np.full((shape[0], shape[1], 1), 0, dtype='uint8')
            label[:, w_index, 0] = 1
            # label[:, w_index, 1] = 1
            # cv2.imshow("label", np.expand_dims(label[..., 0], -1))
            # cv2.waitKey(0)
            img1 = img1 / 255.
            # img1 = cv2.cvtColor(img1, cv2.COLOR_BGR2GRAY)
            # img1 = np.expand_dims(img1, axis=-1)
            i += 1

            X[j] = img1
            Y[j] = label

        yield X, Y


def gen_phrase(map_list, batch_size=32, shape=(5707, 3)):
    voc_dim, timesteps = shape[:2]

    data_list = []
    for line in map_list:
        data_list.append([eval(line[:-3]), int(line[-2:-1])])
    num = len(data_list)

    i = 0
    random.shuffle(data_list)
    while True:
        X = np.zeros((batch_size, timesteps))
        Y = np.zeros((batch_size, 1))

        for j in range(batch_size):
            if i >= num:
                random.shuffle(data_list)
                i = 0
            data = data_list[i]

            d_list = [x for x in data[0]]
            d_list = d_list + [voc_dim]*(timesteps-len(d_list))
            X[j] = np.array(d_list)
            Y[j] = data[1]
            i += 1
        yield X, Y


def generate_data_siamese(char_num=6, char_shape=(40, 40)):
    bg_paths = glob("../data/base/*")
    char_path = "../data/chinese_2500.txt"
    with open(char_path, "r") as f:
        char_str = f.read()

    data_dir = "../data/click/"
    for i in range(1000):
        bg = cv2.imread(random.sample(bg_paths, 1)[0])
        char_list = [char_str[x] for x in random.sample(range(len(char_str)), char_num)]
        image_np, position_list, tips_image_list = char_on_image(bg, char_list, char_shape)

        # for j in range(len(char_list)):
        #     char = char_list[j]
        #     p = position_list[j]
        #     print(char)
        #     cv2.rectangle(image_np, [p[1], p[0]], [p[1]+char_shape[1], p[0]+char_shape[0]], (255, 0, 0), 2)
        #     cv2.imshow("generate_data", image_np)
        #     cv2.waitKey(0)

        # 保存tips图片
        tips_path_list = []
        for k in range(len(tips_image_list)):
            tips_image = tips_image_list[k]
            tips_path = str(i) + "_" + str(k) + ".jpg"
            tips_path_list.append(tips_path)
            cv2.imwrite(data_dir+tips_path, tips_image)
        # 保存文字区域图片
        char_path_list = []
        for j in range(len(char_list)):
            p = position_list[j]
            char_path = str(i) + "_" + str(len(tips_image_list)+j) + ".jpg"
            char_path_list.append(char_path)
            cv2.imwrite(data_dir+char_path, image_np[p[0]:p[0]+char_shape[0], p[1]:p[1]+char_shape[1], :])

        # 生成映射数据
        with open("../data/click/map.txt", "a") as f:
            for j in range(len(tips_path_list)):
                tips_path = tips_path_list[j]
                for k in range(len(char_path_list)):
                    char_path = char_path_list[k]
                    if j == k:
                        f.write(tips_path + "\t" + char_path + "\t" + str(1) + "\n")
                    else:
                        f.write(tips_path + "\t" + char_path + "\t" + str(0) + "\n")


def generate_data_mobile(char_num=6, char_shape=(40, 40), image_shape=(160, 260)):
    # (40,40) (160, 260)
    # (80,80) (360, 590)

    bg_paths = glob("../data/base/*")
    char_path = "../data/chinese_5710.txt"
    with open(char_path, "r") as f:
        char_str = f.read()

    data_dir = "../data/click/"
    # 每个字生成多张图片
    for i in range(0, len(char_str)):
        if i % 100 == 0:
            print("Loop", i)

        char = char_str[i]
        # 生成带背景图数
        image_cnt = 1
        char_list = [char] * image_cnt

        tips_cnt = 0
        image_cnt = 0
        tips_list = []
        for l in range(50):
            # 背景图
            bg = cv2.imread(random.sample(bg_paths, 1)[0])
            if random.choice([0, 0, 1]):
                bg = distort_image(bg)
            if random.choice([0, 0, 1]):
                bg = flip_image(bg)

            # 生成4张tips图，6张带背景的旋转图
            image_np, p_list, t_list = char_on_image(bg, char_list, char_shape, image_shape)
            tips_list += t_list
            for p in p_list:
                char_path = str(i) + "_" + str(image_cnt) + "_2" + ".jpg"
                cv2.imwrite(data_dir+char_path, image_np[p[0]:p[0]+char_shape[0], p[1]:p[1]+char_shape[1], :])
                image_cnt += 1
        for tips_image in tips_list[:4]:
            tips_path = str(i) + "_" + str(tips_cnt) + "_1" + ".jpg"
            cv2.imwrite(data_dir+tips_path, tips_image)
            tips_cnt += 1


def generate_data_yolo_char(char_num=6, char_shape=(40, 40), image_shape=(160, 256)):
    bg_paths = glob("../data/base/*")
    char_path = "../data/chinese_5710.txt"
    with open(char_path, "r") as f:
        char_str = f.read()

    data_dir = "../data/detect/"
    # with open(data_dir+"map.txt", "w") as f:
    #     f.write("")

    for i in range(40000, 60000):
        if i % 1000 == 0:
            print("Loop", i)

        bg = cv2.imread(random.sample(bg_paths, 1)[0])
        if random.choice([0, 0, 1]):
            bg = distort_image(bg)
        if random.choice([0, 0, 1]):
            bg = flip_image(bg)

        char_list = [char_str[x] for x in random.sample(range(len(char_str)), char_num)]
        image_np, position_list, tips_image_list = char_on_image(bg, char_list, char_shape, image_shape, 4)

        if i < 5000:
            tips_image_np, tips_position_list = get_tips_image(tips_image_list, char_shape, image_shape)

        image_np_path = str(i) + ".jpg"
        cv2.imwrite(data_dir+image_np_path, image_np)

        if i < 5000:
            tips_image_np_path = str(i) + "_0.jpg"
            cv2.imwrite(data_dir+tips_image_np_path, tips_image_np)

        # 生成映射数据
        with open(data_dir+"map.txt", "a") as f:
            box_str = ""
            for p in position_list:
                box_str += str(p[1]) + "," + str(p[0]) + "," + \
                           str(p[1]+char_shape[1]) + "," + str(p[0]+char_shape[0]) +\
                           "," + str(0) + " "
            #     cv2.rectangle(image_np, (p[1], p[0]), (p[1]+char_shape[1], p[0]+char_shape[0]), (255, 0, 0), 2)
            # cv2.imshow("image_np", image_np)
            box_str = box_str[:-1]
            f.write(image_np_path + " " + box_str + "\n")

            if i < 5000:
                box_str = ""
                for p in tips_position_list:
                    box_str += str(p[1]) + "," + str(p[0]) + "," + \
                               str(p[1]+char_shape[1]) + "," + str(p[0]+char_shape[0]) + \
                               "," + str(0) + " "
                #     cv2.rectangle(tips_image_np, (p[1], p[0]), (p[1]+char_shape[1], p[0]+char_shape[0]), (255, 0, 0), 2)
                # cv2.imshow("tips_image_np", tips_image_np)
                # cv2.waitKey(0)
                box_str = box_str[:-1]
                f.write(tips_image_np_path + " " + box_str + "\n")


def generate_data_yolo_puzzle(image_shape=(160, 256)):
    bg_paths = glob("../data/base/*.jpeg")

    data_dir = "../data/detect2/"
    with open(data_dir+"map.txt", "w") as f:
        f.write("")

    for i in range(0, 10000):
        if i % 1000 == 0:
            print("Loop", i)

        bg = cv2.imread(random.sample(bg_paths, 1)[0])
        if random.choice([0, 0, 1]):
            bg = distort_image(bg)
        if random.choice([0, 0, 1]):
            bg = flip_image(bg)

        r = random.randint(35, 60)
        puzzle_shape = (r, r)
        image_np, position_list = puzzle_on_image(bg, puzzle_shape, image_shape)

        image_np_path = str(i) + ".jpg"
        cv2.imwrite(data_dir+image_np_path, image_np)

        # 生成映射数据
        with open(data_dir+"map.txt", "a") as f:
            box_str = ""
            for p in position_list:
                box_str += str(p[1]) + "," + str(p[0]) + "," + \
                           str(p[1]+puzzle_shape[1]) + "," + str(p[0]+puzzle_shape[0]) + \
                           "," + str(0) + " "
            #     cv2.rectangle(image_np, (p[1], p[0]), (p[1]+puzzle_shape[1], p[0]+puzzle_shape[0]), (255, 0, 0), 2)
            # cv2.imshow("image_np", image_np)
            # cv2.waitKey(0)
            box_str = box_str[:-1]
            f.write(image_np_path + " " + box_str + "\n")


def generate_data_drag_image(image_shape=(160, 260)):
    bg_paths = glob("../data/base/*")

    data_dir = "../data/drag/"
    with open(data_dir+"map.txt", "w") as f:
        f.write("")

    for i in range(10000):
        if i % 1000 == 0:
            print("Loop", i)

        bg = cv2.imread(random.sample(bg_paths, 1)[0])
        bg = pil_resize(bg, image_shape[0], image_shape[1])
        if random.choice([0, 0, 1]):
            bg = distort_image(bg)
        if random.choice([0, 0, 1]):
            bg = flip_image(bg)

        image_np, clip_line = get_drag_image(bg)

        image_np_path = str(i) + ".jpg"
        cv2.imwrite(data_dir+image_np_path, image_np)

        # 生成映射数据
        with open(data_dir+"map.txt", "a") as f:
            f.write(image_np_path + " " + str(clip_line[0][0]) + "\n")


def generate_data_phrase():
    data_path = os.path.dirname(os.path.abspath(__file__)) + "/../data/phrase/"

    char_path = data_path+"char.txt"
    with open(char_path, "r") as f:
        char_list = f.readlines()
    char_dict = {}
    for i in range(len(char_list)):
        char_dict[char_list[i]] = i

    phrase_list = []
    phrase_path = data_path+"phrase3.txt"
    with open(phrase_path, "r") as f:
        phrase_list += f.readlines()
    phrase_path = data_path+"phrase4.txt"
    with open(phrase_path, "r") as f:
        phrase_list += f.readlines()
    phrase_path = data_path+"phrase5.txt"
    with open(phrase_path, "r") as f:
        phrase_list += f.readlines()
    phrase_set = set(phrase_list)

    map_path = data_path+"map3.txt"
    with open(map_path, "w") as f:
        f.write("")
    data_list = []
    start_time = time.time()
    i = 0
    negative_way_flag = False
    for phrase in phrase_list:
        if i % 500000 == 0:
            with open(map_path, "a") as f:
                f.writelines(data_list)
            data_list = []
            print("Loop", i, len(phrase_list), time.time()-start_time)
            start_time = time.time()
        i += 1

        # 正样本
        index_list = []
        for char in phrase[:-1]:
            index_list.append(char_dict.get(char+"\n"))
        data_list.append(str(index_list) + " 1\n")

        # 负样本
        if negative_way_flag:
            index1 = random.randint(0, len(index_list)-1)
            find_flag = False
            while not find_flag:
                index2 = random.randint(0, len(index_list)-1)
                if index1 != index2:
                    find_flag = True
            temp = index_list[index1]
            index_list[index1] = index_list[index2]
            index_list[index2] = temp
            if "".join([char_list[x][:-1] for x in index_list]) + "\n" not in phrase_set:
                data_list.append(str(index_list) + " 0\n")
        else:
            products = list(product(index_list, repeat=len(index_list)))
            random.shuffle(products)
            negative_cnt = 0
            for p in products:
                if negative_cnt >= 2:
                    break
                p = list(p)
                if len(set(p)) != len(p):
                    continue
                if p != index_list and "".join([char_list[x][:-1] for x in p]) + "\n" not in phrase_set:
                    data_list.append(str(p) + " 0\n")
                    negative_cnt += 1

    with open(map_path, "a") as f:
        f.writelines(data_list)


def generate_data_phrase_raw(word_len=5):
    paths = glob("D:/Chinese_corpus/answer/*/*.txt")

    phrase_path = "../data/phrase/phrase" + str(word_len) + "_new.txt"
    triple_list = []
    reg = "[^\u4e00-\u9fa5]"
    start_time = time.time()
    for i in range(len(paths)):
        if i % 1000 == 0:
            with open(phrase_path, "w") as f:
                f.writelines(triple_list)
            print("Loop", i, len(paths), time.time()-start_time)
            start_time = time.time()
            triple_list = []

        with open(paths[i], "rb") as f:
            _b = f.read()
        try:
            text = _b.decode("gbk")
        except:
            try:
                text = _b.decode("gb2312")
            except:
                try:
                    text = _b.decode("gb18030")
                except:
                    print(chardet.detect(_b), "is None")

        filter_word = ["的"]
        for word in filter_word:
            text = re.sub(word, "#"*len(word), text)

        word_list = jieba.lcut(text, cut_all=False, HMM=True)

        for j in range(1, len(word_list)):
            current = word_list[j]
            current_re = re.search(reg, current)
            last = word_list[j-1]
            last_re = re.search(reg, last)

            if current_re:
                continue
            if len(current) == word_len:
                triple_list.append(current + "\n")
            elif len(current) + len(last) == word_len and not last_re:
                triple_list.append(last+current + "\n")

        triple_list = list(set(triple_list))

    print("len(triple_list)", len(triple_list))
    with open(phrase_path, "w") as f:
        f.writelines(triple_list)


def char_on_image(image_np, char_list, char_shape, image_shape, tip_char_num=1):
    position_list = []
    for char in char_list:
        # 获取单字图片
        char_image_pil = get_char_image(char, char_shape)

        image_np = pil_resize(image_np, image_shape[0], image_shape[1])

        # h, w
        fg_w, fg_h = char_image_pil.size[:2]
        bg_h, bg_w = image_np.shape[:2]

        # 字体放置的位置，且位置不重叠
        find_flag = 0
        while not find_flag:
            position_h = random.randint(0, bg_h-fg_h)
            position_w = random.randint(0, bg_w-fg_w)
            if len(position_list) < 1:
                find_flag = 1
                break
            for p in position_list:
                if get_iou(position_w, position_h, position_w+fg_w, position_h+fg_h,
                           p[1], p[0], p[1]+fg_w, p[0]+fg_h) > 0:
                    find_flag = 0
                    break
                else:
                    find_flag = 1
        position_list.append([position_h, position_w])
        # 字体添加到背景图上
        # image_np = get_image_roi(image_np, char_image_np, position_h, position_w)
        image_np = get_image_paste(image_np, char_image_pil, position_h, position_w)

    # 生成提示图片
    image_list = []
    for char in char_list[:tip_char_num]:
        char_image_pil = get_char_image(char, char_shape, rotate=False, bg_color=(255, 255, 255, 255))
        char_image_np = pil2np_a(char_image_pil)
        # char_image_np = pil_resize(char_image_np, char_shape[0], char_shape[1])
        image_list.append(char_image_np)
    tips_image_np = np.concatenate(image_list, axis=1)
    # 加干扰
    tips_image_np = create_noise(tips_image_np)

    # 切割
    image_list = []
    for i in range(tip_char_num):
        image_list.append(tips_image_np[:, i*char_shape[1]:(i+1)*char_shape[1], :])
    return image_np, position_list, image_list


def get_char_image(char, char_shape, rotate=True, bg_color=(0, 0, 0, 0)):
    # 创建空图
    image_pil = Image.new('RGBA', (80, 80), bg_color)

    # 空图上写字
    # font_size = 35 # (40, 40)
    font_size = 75 # (80, 80)
    font_type_list = glob("../font/*")
    font_type = random.sample(font_type_list, 1)[0]
    font_config = ImageFont.truetype(font_type, int(font_size))
    dr = ImageDraw.Draw(image_pil)
    fill_color = random_color()
    fill_color = (fill_color[0], fill_color[1], fill_color[2])
    dr.text((3, -6), char, font=font_config, fill=fill_color)

    if rotate:
        if random.choice([0, 1]):
            angle = random.randint(0, 80)
        else:
            angle = random.randint(280, 360)
        image_pil = image_pil.rotate(angle, expand=False, fillcolor=bg_color)

    # image_pil.show("1")
    image_pil = image_pil.resize(char_shape)
    # cv2.imshow("get_char_image", pil2np(image_pil))
    # cv2.waitKey(0)
    return image_pil


def get_tips_image(tips_image_list, char_shape, image_shape, roatate=True):
    new_list = []
    for img in tips_image_list:
        if random.choice([0, 0, 1]):
            angle = random.randint(0, 360)
            img = pil_rotate(img, angle, (255, 255, 255))
        new_list.append(img)

    tips_image_np = np.concatenate(new_list, axis=1)

    new_image = np.full((image_shape[0], image_shape[1], 3), 0, np.uint8)
    new_image[:tips_image_np.shape[0], :tips_image_np.shape[1], :] = tips_image_np

    position_list = []
    for i in range(len(new_list)):
        h = 0
        w = i*char_shape[1]
        position_list.append([h, w])
    return new_image, position_list


def get_image_roi(image_bg, image_fg, roi_h, roi_w):
    # h, w
    fg_h, fg_w = image_fg.shape[:2]
    bg_h, bg_w = image_bg.shape[:2]

    # roi取值范围
    roi = image_bg[roi_h:roi_h+fg_h, roi_w:roi_w+fg_w]

    # 获取bg中非fg字体部分的掩码，相当于排除fg的字体部分，只保留bg的除fg字体外的部分
    img_fg_gray = cv2.cvtColor(image_fg, cv2.COLOR_BGR2GRAY)
    ret, mask = cv2.threshold(img_fg_gray, 0, 255, cv2.THRESH_OTSU)
    bg_roi = cv2.bitwise_and(roi, roi, mask=mask)

    # 获取fg中字体部分的掩码，相当于排除fg中的白色背景，只保留fg的字体部分
    mask_inv = cv2.bitwise_not(mask)
    fg_roi = cv2.bitwise_and(image_fg, image_fg, mask=mask_inv)
    # 膨胀腐蚀去掉白色颗粒
    # kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
    # fg_roi = cv2.erode(fg_roi, kernel)
    # fg_roi = cv2.dilate(fg_roi, kernel)

    # bg的除字体外背景部分 + fg的字体部分
    image_roi = cv2.add(bg_roi, fg_roi)

    # 将roi部分放回原bg
    image_bg[roi_h:roi_h+fg_h, roi_w:roi_w+fg_w, :] = image_roi

    # cv2.imshow("image_fg", image_fg)
    # cv2.imshow("get_image_roi", image_bg)
    # cv2.waitKey(0)
    return image_bg


def get_image_paste(image_bg, image_fg, roi_h, roi_w):
    fg_h, fg_w = image_fg.size[:2]
    image_bg = cv2.cvtColor(image_bg, cv2.COLOR_BGR2BGRA)
    image_bg = np2pil_a(image_bg)
    # image_fg = np2pil_a(image_fg)
    image_bg.paste(image_fg, (roi_w, roi_h), image_fg)
    image_bg = pil2np(image_bg)

    # cv2.imshow("get_image_paste", image_bg)
    # cv2.waitKey(0)
    return image_bg


def random_color(dims=3):
    color = [0]*dims

    find_flag = 0
    while not find_flag:
        for dim in range(dims):
            color[dim] = random.randint(0, 255)
            if color[dim] <= 125:
                find_flag = 1

    # RGB
    # color_list = [
    #     [207, 91, 85],
    #     [0, 201, 88],
    #     [117, 74, 57],
    #     [210, 210, 27],
    #     [160, 157, 152],
    #     [181, 210, 210],
    #     [27, 112, 107],
    #     [87, 26, 44],
    #     [115, 19, 20],
    #     [161, 210, 68],
    #     [210, 108, 12],
    #     [112, 9, 142],
    #     [50, 41, 84],
    #     [72, 52, 210],
    #     [210, 177, 89],
    #     [148, 200, 89],
    #     [173, 116, 109],
    #     [185, 185, 210],
    #     [181, 7, 210],
    #     [80, 210, 30],
    #     [65, 72, 98],
    #     [210, 123, 109],
    #     [19, 64, 95],
    #     [128, 21, 210],
    #     [129, 137, 60]
    # ]
    # color = random.sample(color_list, 1)[0]
    return tuple(color)


def create_noise(image_np):
    ic = ImageCaptcha()
    image_pil = np2pil(image_np)
    image_pil = ic.create_noise_curve(image_pil, random_color())
    image_pil = ic.create_noise_curve(image_pil, random_color())
    image_pil = ic.create_noise_dots(image_pil, random_color())
    image_np = pil2np(image_pil)
    return image_np


def get_iou(x1, y1, x2, y2, a1, b1, a2, b2):
    # 相交区域左上角横坐标
    ax = max(x1, a1)
    # 相交区域左上角纵坐标
    ay = max(y1, b1)
    # 相交区域右下角横坐标
    bx = min(x2, a2)
    # 相交区域右下角纵坐标
    by = min(y2, b2)

    area_n = (x2 - x1) * (y2 - y1)
    area_m = (a2 - a1) * (b2 - b1)

    w = max(0, bx - ax)
    h = max(0, by - ay)
    area_x = w * h

    return area_x / (area_n + area_m - area_x)


def preprocess_true_boxes(true_boxes, input_shape, anchors, num_classes):
    """Preprocess true boxes to training input format
    Parameters
    ----------
    true_boxes: array, shape=(m, T, 5)
        Absolute x_min, y_min, x_max, y_max, class_id relative to input_shape.
    input_shape: array-like, hw, multiples of 32
    anchors: array, shape=(N, 2), wh
    num_classes: integer
    Returns
    -------
    y_true: list of array, shape like yolo_outputs, xywh are reletive value
    """
    # print(true_boxes[..., 4])
    # print(num_classes)
    assert (true_boxes[..., 4] < num_classes).all(), 'class id must be less than num_classes'
    # default setting
    num_layers = len(anchors)//3
    anchor_mask = [[6, 7, 8], [3, 4, 5], [0, 1, 2]] if num_layers == 3 else [[3, 4, 5], [1, 2, 3]]

    true_boxes = np.array(true_boxes, dtype='float32')
    input_shape = np.array(input_shape, dtype='int32')
    boxes_xy = (true_boxes[..., 0:2] + true_boxes[..., 2:4]) // 2
    boxes_wh = true_boxes[..., 2:4] - true_boxes[..., 0:2]
    true_boxes[..., 0:2] = boxes_xy/input_shape[::-1]
    true_boxes[..., 2:4] = boxes_wh/input_shape[::-1]

    m = true_boxes.shape[0]
    grid_shapes = [input_shape//{0: 32, 1: 16, 2: 8}[l] for l in range(num_layers)]
    y_true = [np.zeros((m, grid_shapes[l][0], grid_shapes[l][1],len(anchor_mask[l]), 5+num_classes),
                       dtype='float32') for l in range(num_layers)]

    # Expand dim to apply broadcasting.
    anchors = np.expand_dims(anchors, 0)
    anchor_maxes = anchors / 2.
    anchor_mins = -anchor_maxes
    valid_mask = boxes_wh[..., 0] > 0

    for b in range(m):
        # Discard zero rows.
        wh = boxes_wh[b, valid_mask[b]]
        if len(wh) == 0:
            continue
        # Expand dim to apply broadcasting.
        wh = np.expand_dims(wh, -2)
        box_maxes = wh / 2.
        box_mins = -box_maxes

        intersect_mins = np.maximum(box_mins, anchor_mins)
        intersect_maxes = np.minimum(box_maxes, anchor_maxes)
        intersect_wh = np.maximum(intersect_maxes - intersect_mins, 0.)
        intersect_area = intersect_wh[..., 0] * intersect_wh[..., 1]
        box_area = wh[..., 0] * wh[..., 1]
        anchor_area = anchors[..., 0] * anchors[..., 1]
        iou = intersect_area / (box_area + anchor_area - intersect_area)

        # Find best anchor for each true box
        best_anchor = np.argmax(iou, axis=-1)

        for t, n in enumerate(best_anchor):
            for l in range(num_layers):
                if n in anchor_mask[l]:
                    i = np.floor(true_boxes[b,t,0]*grid_shapes[l][1]).astype('int32')
                    j = np.floor(true_boxes[b,t,1]*grid_shapes[l][0]).astype('int32')
                    k = anchor_mask[l].index(n)
                    c = true_boxes[b, t, 4].astype('int32')
                    y_true[l][b, j, i, k, 0:4] = true_boxes[b, t, 0:4]
                    y_true[l][b, j, i, k, 4] = 1
                    y_true[l][b, j, i, k, 5+c] = 1

    return y_true


def get_puzzle(shape=(80, 80)):
    # 创建空图
    image_pil = Image.new('RGBA', (shape[1], shape[0]), (255, 255, 255, 0))
    draw = ImageDraw.Draw(image_pil)

    # 居中创建矩形
    rec_shape = (40, 40)
    left_up_point = [int((shape[0]-rec_shape[0])/2), int((shape[1]-rec_shape[1])/2)]
    right_down_point = [left_up_point[0]+rec_shape[0], left_up_point[1]+rec_shape[1]]

    # 透明度
    flag = random.choice([0, 1])
    if flag:
        alpha = random.randint(100, 150)
    else:
        alpha = random.randint(160, 255)

    # 背景色
    if flag:
        r = random.randint(0, 30)
    else:
        r = random.randint(100, 180)
    # r = random.randint(0, 255)
    fill_color = (r, r, r, alpha)

    # 边缘色
    if random.choice([0, 1, 1]):
        if flag:
            r = random.randint(140, 170)
        else:
            r = random.randint(70, 100)
        outline_color = (r, r, r)
    else:
        outline_color = (fill_color[0], fill_color[1], fill_color[2])

    draw.rectangle((left_up_point[1],
                    left_up_point[0],
                    left_up_point[1]+rec_shape[1],
                    left_up_point[0]+rec_shape[0]),
                   fill=fill_color,
                   outline=outline_color)

    # 拼图的圆或半圆
    radius = random.randint(int(rec_shape[0] / 3 / 2), int(rec_shape[0] / 3 / 1.2))
    center_list = [[left_up_point[1], int((right_down_point[0]+left_up_point[0])/2), 1],
                   [right_down_point[1], int((right_down_point[0]+left_up_point[0])/2), 1],
                   [int((right_down_point[1]+left_up_point[1])/2), left_up_point[0], 0],
                   [int((right_down_point[1]+left_up_point[1])/2), right_down_point[0], 0]
                   ]
    circle_num = random.randint(1, 4)
    # print("circle_num", circle_num)
    center_list = random.sample(center_list, circle_num)

    min_w, min_h = left_up_point[1], left_up_point[0]
    max_w, max_h = right_down_point[1], right_down_point[0]
    for center in center_list:
        w, h = center[:2]
        is_width = center[2]

        # 判断长宽
        into_ratio = random.randint(int(1/2*radius), int(3/4*radius))
        if is_width:
            # 挑选圆是凸还是凹进去
            if random.choice([0, 1]):
                center = (center[0]+into_ratio, center[1])
            else:
                center = (center[0]-into_ratio, center[1])
        else:
            if random.choice([0, 1]):
                center = (center[0], center[1]+into_ratio)
            else:
                center = (center[0], center[1]-into_ratio)

        # 判断透明度
        color = fill_color
        if is_width:
            if left_up_point[1] <= center[0] <= right_down_point[1]:
                color = (0, 0, 0, 0)
        else:
            if left_up_point[0] <= center[1] <= right_down_point[0]:
                color = (0, 0, 0, 0)

        # print("center, color, alpha", center, color, alpha)
        draw.ellipse([(center[0]-radius, center[1]-radius),
                      (center[0]+radius, center[1]+radius)],
                     fill=color,
                     outline=outline_color)

        # 修补内部圆的边缘颜色
        if color[3] == alpha:
            if is_width:
                if center[0] < w:
                    draw.rectangle((w,
                                    h-radius,
                                    center[0]+radius,
                                    center[1]+radius),
                                   fill=fill_color)
                else:
                    draw.rectangle((center[0]-radius,
                                    center[1]-radius,
                                    w,
                                    h+radius),
                                   fill=fill_color)
            else:
                if center[1] < h:
                    draw.rectangle((w-radius,
                                    h,
                                    center[0]+radius,
                                    center[1]+radius),
                                   fill=fill_color)
                else:
                    draw.rectangle((center[0]-radius,
                                    center[1]-radius,
                                    w+radius,
                                    h),
                                   fill=fill_color)
        # 修补外部圆的边缘颜色
        else:
            if is_width:
                if center[0] > w:
                    draw.rectangle((center[0]-radius,
                                    center[1]-radius,
                                    w,
                                    h+radius),
                                   fill=(0, 0, 0, 0))
                else:
                    draw.rectangle((w,
                                    h-radius,
                                    center[0]+radius,
                                    center[1]+radius),
                                   fill=(0, 0, 0, 0))
            else:
                if center[1] > h:
                    draw.rectangle((center[0]-radius,
                                    center[1]-radius,
                                    w+radius,
                                    h),
                                   fill=(0, 0, 0, 0))
                else:
                    draw.rectangle((w-radius,
                                    h,
                                    center[0]+radius,
                                    center[1]+radius),
                                   fill=(0, 0, 0, 0))

        # 新增面积
        if color[3] == alpha:
            if center[0]-radius <= min_w:
                min_w = center[0]-radius
            if center[0]+radius >= max_w:
                max_w = center[0]+radius
            if center[1]-radius <= min_h:
                min_h = center[1]-radius
            if center[1]+radius >= max_h:
                max_h = center[1]+radius

    image_pil = image_pil.crop([min_w, min_h, max_w+1, max_h+1])
    # image_pil.show("2")
    return image_pil


def puzzle_on_image(image_np, puzzle_shape, image_shape):
    position_list = []
    # 获取拼图图片
    puzzle_image_pil = get_puzzle()
    puzzle_image_pil = puzzle_image_pil.resize(puzzle_shape)
    image_np = pil_resize(image_np, image_shape[0], image_shape[1])

    # h, w
    fg_w, fg_h = puzzle_image_pil.size[:2]
    bg_h, bg_w = image_np.shape[:2]
    # 拼图放置的位置
    position_h = random.randint(0, bg_h-fg_h)
    position_w = random.randint(0, bg_w-fg_w)
    position_list.append([position_h, position_w])
    # for p in position_list:
    #     cv2.rectangle(image_np, (p[1], p[0]),
    #                   (p[1]+puzzle_shape[1], p[0]+puzzle_shape[0]),
    #                   (0, 0, 255), 1)
    # 拼图添加到背景图上
    image_np = get_image_paste(image_np, puzzle_image_pil, position_h, position_w)
    # cv2.imshow("puzzle_on_image", image_np)
    # cv2.waitKey(0)
    return image_np, position_list


def distort_image(image_np, hue=.1, sat=1.5, val=1.5):
    """
    图像失真
    :return:
    """
    def rand(a=0, b=1):
        return np.random.rand()*(b-a) + a

    # cv2.imshow("distort_image1", image_np)
    hue = rand(-hue, hue)
    sat = rand(1, sat) if rand() < .5 else 1/rand(1, sat)
    val = rand(1, val) if rand() < .5 else 1/rand(1, val)

    image_np = cv2.cvtColor(image_np, cv2.COLOR_BGR2RGB)
    x = rgb_to_hsv(image_np/255.)
    x[..., 0] += hue
    x[..., 0][x[..., 0] > 1] -= 1
    x[..., 0][x[..., 0] < 0] += 1
    x[..., 1] *= sat
    x[..., 2] *= val
    x[x > 1] = 1
    x[x < 0] = 0
    image_np = hsv_to_rgb(x)
    image_np = cv2.cvtColor(np.uint8(image_np*255), cv2.COLOR_RGB2BGR)

    # cv2.imshow("distort_image2", image_np)
    # cv2.waitKey(0)
    return image_np


def flip_image(image_np):
    # cv2.imshow("flip_image1", image_np)
    if random.choice([0, 1]):
        # 水平翻转
        image_np = cv2.flip(image_np, 1)
    else:
        # 垂直翻转
        image_np = cv2.flip(image_np, 0)
    # cv2.imshow("flip_image2", image_np)
    # cv2.waitKey(0)
    return image_np


def get_drag_image(image_np):
    h, w = image_np.shape[:2]

    # 取一定高度图片
    clip_h = random.randint(int(1/4*h), int(3/4*h))
    image_clip = image_np[:clip_h, ...]

    # 将图片在一定宽度截断，重新拼接
    clip_w = random.randint(int(1/6*w), int(5/6*w))
    image_w1 = image_clip[:, :clip_w, ...]
    image_w2 = image_clip[:, clip_w:, ...]
    image_new = np.concatenate([image_w2, image_w1], axis=1)

    # 分割线
    clip_line = [(image_w2.shape[1], 0), (image_w2.shape[1], clip_h)]

    # show
    # print(clip_line)
    # cv2.line(image_new, clip_line[0], clip_line[1], (0, 0, 255), 2)
    # cv2.imshow("get_drag_image", image_new)
    # cv2.waitKey(0)
    return image_new, clip_line


def get_real_data_puzzle(shape=(160, 256)):
    paths = glob("../data/detect2_real/*")
    i = 10000
    for p in paths:
        image = cv2.imread(p)
        image = pil_resize(image, shape[0], shape[1])
        cv2.imwrite("../data/detect2_real/"+str(i)+".jpg", image)
        i += 1

        image = distort_image(image)
        cv2.imwrite("../data/detect2_real/"+str(i)+".jpg", image)
        i += 1

        image = flip_image(image)
        cv2.imwrite("../data/detect2_real/"+str(i)+".jpg", image)
        i += 1


def read_label_puzzle():
    paths = glob("../data/detect2_real/*.json")
    map_path = "../data/detect2_real/map.txt"
    with open(map_path, "a") as f:
        for p in paths:
            with open(p, "r") as fp:
                _dict = json.loads(fp.read())
            points = _dict.get("shapes")[0].get("points")
            image_path = _dict.get("imagePath")
            ps = [str(int(points[0][0])), str(int(points[0][1])),
                  str(int(points[1][0])), str(int(points[1][1]))]
            p_str = ",".join(ps)
            f.write(image_path + " " + p_str + ",0" + "\n")


def fix_map_txt():
    path = "../data/map.txt"
    with open(path, "r") as f:
        _list = f.readlines()

    with open("../data/map_new.txt", "w") as f:
        new_list = []
        for line in _list:
            ss = line.split(" ")
            ps = ss[-1][:-1].split(",")[:-1]

            if random.choice([0, 1, 1, 1]):
                pix = random.choice([1, 2, 2, 3, 3, 4, 4])
                for i in range(len(ps)):
                    if i < 2:
                        ps[i] = str(int(ps[i]) - pix)
                    else:
                        ps[i] = str(int(ps[i]) + pix)
            new_line = ss[0] + " " + ",".join(ps) + ",0\n"
            new_list.append(new_line)
            print("line", line)
            print("new_line", new_line)
        f.writelines(new_list)


def get_char_map():
    path = "../data/phrase/phrase3.txt"
    with open(path, "r") as f:
        _list = f.readlines()
    path = "../data/phrase/phrase4.txt"
    with open(path, "r") as f:
        _list += f.readlines()
    path = "../data/phrase/phrase5.txt"
    with open(path, "r") as f:
        _list += f.readlines()

    _str = "".join(_list)
    _str = re.sub("\n", "", _str)
    _list = list(set([x+"\n" for x in _str]))
    _list.sort(key=lambda x: x)
    with open("../data/phrase/char.txt", "w") as f:
        f.writelines(_list)


if __name__ == "__main__":
    # from click_captcha.utils import get_classes, get_anchors
    # annotation_path = '../data/detect/map.txt'
    # log_dir = 'yolo_data/logs/000/'
    # classes_path = 'yolo_data/my_classes.txt'
    # anchors_path = 'yolo_data/tiny_yolo_anchors.txt'
    # class_names = get_classes(classes_path)
    # num_classes = len(class_names)
    # anchors = get_anchors(anchors_path)
    #
    # with open(annotation_path) as f:
    #     lines = f.readlines()
    # random.shuffle(lines)
    # input_shape = (160, 256)
    # g = gen_yolo(lines, 10, input_shape, anchors, num_classes)
    # list(g)

    generate_data_phrase()

    # generate_data_yolo_puzzle()
    # gen_yolo_puzzle()
    # _path = "../data/base/0b16d1f1a4e017d4a7ab5779263887f1.jpeg"
    # get_drag_image(cv2.imread(_path))

    # for ii in range(10):
    #     im = get_puzzle()

    # with open("../data/chinese.txt", "r") as f:
    #     _str = f.read()
    #
    # _list = [c for c in _str]
    # _list = list(set(_list))
    # _str = "".join(_list)

    # with open("../data/chinese.txt", "w") as f:
    #     f.write(_str)