AaronJny/captcha_detection

# -*- coding: utf-8 -*-
import time

import requests
from tqdm import tqdm

for i in tqdm(range(300)):
    response = requests.get('http://xxxxxxxx/ValidateImage')
    with open('./images/{}.png'.format(i), 'wb') as f:
        f.write(response.content)
    time.sleep(2)

# -*- coding: utf-8 -*-
# @Time    : 2020/11/12 21:46
# @Author  : AaronJny
# @File    : convert.py
# @Desc    :
# 引入转换脚本
from xyolo import voc2xyolo

# voc格式的标注数据路径的正则表达式
input_path = './labels/*.xml'
# classes是我们要检测的所有有效类别名称构成的txt文件，每个类别一行
classes_path = './classes.txt'
# 转换后的xyolo数据集存放路径
output_path = './xyolo_label.txt'
# 开始转换
voc2xyolo(input_path=input_path, classes_path=classes_path, output_path=output_path)

# -*- coding: utf-8 -*-
# @Time    : 2020/11/12 22:18
# @Author  : AaronJny
# @File    : config.py
# @Desc    :
from xyolo import DefaultYolo3Config


# 创建一个DefaultYolo3Config的子类，在子类里覆盖默认的配置
class XYoloConfig(DefaultYolo3Config):
    def __init__(self):
        super(XYoloConfig, self).__init__()
        # 数据集路径，推荐使用绝对路径
        self._dataset_path = '/Users/aaron/code/captcha_detection/xyolo_label.txt'
        # 类别名称文件路径，推荐使用绝对路径
        self._classes_path = '/Users/aaron/code/captcha_detection/classes.txt'
        # 模型保存路径，默认是保存在当前路径下的xyolo_data下的，也可以进行更改
        # 推荐使用绝对路径
        self._output_model_path = 'detect_model.h5'
        

# -*- coding: utf-8 -*-
# @Time    : 2020/11/12 22:08
# @Author  : AaronJny
# @File    : train_detect.py
# @Desc    :
# 导入包
from xyolo import YOLO
from xyolo import init_yolo_v3

from config import XYoloConfig

# 使用修改后的配置创建yolo对象
config = XYoloConfig()
init_yolo_v3(config)
# 如果是训练，在创建yolo对象时要传递参数train=True
yolo = YOLO(config, train=True)
# 开始训练，训练完成后会自动保存
yolo.fit()

# -*- coding: utf-8 -*-
# @Time    : 2020/11/12 22:28
# @Author  : AaronJny
# @File    : test_detect.py
# @Desc    :
from xyolo import YOLO

from config import XYoloConfig

yolo = YOLO(XYoloConfig())
img = yolo.detect_and_draw_image('./images/286.png')
img.show()

# -*- coding: utf-8 -*-
# @Time    : 2020/11/12 22:54
# @Author  : AaronJny
# @File    : extract.py
# @Desc    :
import os
from glob import glob

import cv2
import numpy as np
from PIL import Image
from tqdm import tqdm


def extract_correct_word():
    save_dir = './correct_words'
    if not os.path.exists(save_dir):
        os.mkdir(save_dir)
    images = glob('images/*.png')
    for image_path in tqdm(images):
        # 获取图片编号
        num = image_path.split('/')[-1].split('.')[0]
        # 读取图片
        im = np.asarray(Image.open(image_path).convert("RGB"))
        # 切割第一个字
        word_im = im[1:34, 186:216]
        cv2.imwrite(os.path.join(save_dir, 'x-{}-{}.png'.format(num, 0)), word_im)
        # 切割第二个字
        word_im = im[1:34, 216:246]
        cv2.imwrite(os.path.join(save_dir, 'x-{}-{}.png'.format(num, 1)), word_im)



if __name__ == '__main__':
    extract_correct_word()

# -*- coding: utf-8 -*-
# @Time    : 2020/11/12 22:54
# @Author  : AaronJny
# @File    : extract.py
# @Desc    :
import os
from glob import glob

import cv2
import numpy as np
from PIL import Image
from tqdm import tqdm


def extract_correct_word():
    save_dir = './correct_words'
    if not os.path.exists(save_dir):
        os.mkdir(save_dir)
    images = glob('images/*.png')
    for image_path in tqdm(images):
        # 获取图片编号
        num = image_path.split('/')[-1].split('.')[0]
        # 读取图片
        im = np.asarray(Image.open(image_path).convert("RGB"))
        # 切割第一个字
        word_im = im[1:34, 186:216]
        cv2.imwrite(os.path.join(save_dir, 'x-{}-{}.png'.format(num, 0)), word_im)
        # 切割第二个字
        word_im = im[1:34, 216:246]
        cv2.imwrite(os.path.join(save_dir, 'x-{}-{}.png'.format(num, 1)), word_im)


def extract_gen_word():
    # 创建目录
    save_dir = './gen_words'
    if not os.path.exists(save_dir):
        os.mkdir(save_dir)
    # 读取数据标注结果
    with open('xyolo_label.txt', 'r') as f:
        lines = f.readlines()
    # 对于每一张图片
    for line in lines:
        n_line = line.strip()
        image_path, *pos = n_line.split()
        num = image_path.split('/')[-1].split('.')[0]
        im = np.asarray(Image.open(image_path).convert("RGB"))
        # 对于每一个框框
        for index, _pos in enumerate(pos):
            x1, y1, x2, y2, cls = map(int, _pos.split(','))
            n_im = im.copy()
            cv2.rectangle(n_im, (x1, y1), (x2, y2), (0, 255, 0), 2)
            cv2.imshow('text', n_im)
            cv2.waitKey(1000)
            word = input('请输入当前选中汉字：')
            cv2.imwrite(os.path.join('gen_word/{}-{}-{}.png'.format(word, num, index)), im[y1:y2, x1:x2])
            del n_im
            cv2.destroyAllWindows()


if __name__ == '__main__':
    # extract_correct_word()
    extract_gen_word()

# -*- coding: utf-8 -*-
# @Time    : 2020/11/13 08:54
# @Author  : AaronJny
# @File    : create_classify_dataset.py
# @Desc    : 创建二分类器的数据集
import json
import random
from glob import glob
from os.path import abspath


def create_dataset():
    correct_word_images = glob('correct_words/*.png')
    # 建立原始字（即标签部分的字）->图片路径列表的映射
    correct_word_map = {}
    for image in correct_word_images:
        word = image.split('/')[-1].split('-')[0]
        # 过滤掉认不出来的字
        if word == '1':
            continue
        correct_word_map.setdefault(word, []).append(image)
    # 建立生成字（即验证部分的字）->图片路径列表的映射
    gen_word_images = glob('gen_words/*.png')
    gen_word_images_set = set(gen_word_images)
    gen_word_map = {}
    for image in gen_word_images:
        word = image.split('/')[-1].split('-')[0]
        # 过滤掉认不出来的字
        if word == '1':
            continue
        gen_word_map.setdefault(word, []).append(image)
    # 通过负采样生成数据集
    records = []
    for word, correct_images in correct_word_map.items():
        for correct_image in correct_images:
            gen_images = gen_word_map.get(word, [])
            tmp_records = []
            # 如果有生成的图片
            if gen_images:
                # 负采样数和正样本数相同
                sample_num = len(gen_images)
                # 添加正样本
                for gen_image in gen_images:
                    tmp_records.append((1, abspath(correct_image), abspath(gen_image)))
                # 随机选择负样本
                random_images = gen_word_images_set - set(gen_images)
                for image in random.sample(list(random_images), sample_num):
                    tmp_records.append((0, abspath(correct_image), abspath(image)))
            records.append(tmp_records)
    with open('data.json', 'w') as f:
        json.dump(records, f, ensure_ascii=False, indent=2)


if __name__ == '__main__':
    create_dataset()

class ClassifyConfig:
    # 数据集路径
    DATA_PATH = './data.json'
    # 验证集比例
    VAL_SPLIT = 0.2
    # 测试集比例
    TEST_SPLIT = 0
    # batch_size
    BATCH_SIZE = 32
    # Dropout
    DROPOUT_RATE = 0.2
    # imagenet数据集均值
    IMAGE_MEAN = [0.485, 0.456, 0.406]
    # imagenet数据集标准差
    IMAGE_STD = [0.299, 0.224, 0.225]
    # 学习率
    LR = 1e-3
    # 图片大小
    IMAGE_SIZE = (32, 32)
    # 图像信道
    IMAGE_CHANNELS = 1
    # 模型保存地址
    MODEL_PATH = './word_classify_best_weights.h5'

# -*- coding: utf-8 -*-
# @Time    : 2020/11/13
# @Author  : AaronJny
# @File    : process_classify_dataset.py
# @Desc    :
import json
import math

import numpy as np
from PIL import Image

from config import ClassifyConfig


def preprocess_image(image: Image.Image):
    """
    对图片进行预处理，并返回处理后的图片矩阵
    """
    if isinstance(image, str):
        image = Image.open(image)
    x = image.resize(ClassifyConfig.IMAGE_SIZE)
    convert_flag = {
        1: 'L',
        3: 'RGB'
    }
    x = np.asarray(x.convert(convert_flag[ClassifyConfig.IMAGE_CHANNELS]), dtype=np.float32)
    x = x / 255.
    x = x.reshape((*ClassifyConfig.IMAGE_SIZE, ClassifyConfig.IMAGE_CHANNELS))
    return x


class DataGenerator:
    """
    封装了数据集的生成器
    """

    def __init__(self, data, batch_size=ClassifyConfig.BATCH_SIZE, shuffle=True):
        self.data = data
        self.batch_size = batch_size
        self.shuffle = shuffle
        self.steps = int(math.ceil(len(self.data) / self.batch_size))

    @classmethod
    def load_image(cls, image_path):
        im = Image.open(image_path)
        image = np.asarray(im.convert("RGB"))
        del im
        return image

    @classmethod
    def preprocess_record(cls, record):
        """
        对一条给定数据进行预处理
        """
        label, image_path1, image_path2 = record
        image1 = preprocess_image(image_path1)
        image2 = preprocess_image(image_path2)
        return image1, image2, label

    def __iter__(self):
        if self.shuffle:
            np.random.shuffle(self.data)
        total = len(self.data)
        for start in range(0, total, self.batch_size):
            end = min(total, start + self.batch_size)
            batch_images1 = []
            batch_images2 = []
            batch_labels = []
            for record in self.data[start:end]:
                image1, image2, label = self.preprocess_record(record)
                batch_images1.append(image1)
                batch_images2.append(image2)
                batch_labels.append(label)
            np_batch_images1 = np.asarray(batch_images1)
            np_batch_images2 = np.asarray(batch_images2)
            np_batch_labels = np.asarray(batch_labels)
            np_batch_labels.reshape((-1, 1))
            del batch_images1, batch_images2, batch_labels
            yield [np_batch_images1, np_batch_images2], np_batch_labels
            del np_batch_images1, np_batch_images2, np_batch_labels

    def for_fit(self):
        while True:
            yield from self.__iter__()


def load_dataset():
    """
    加载并划分、封装数据集，返回(训练数据集,验证数据集,测试数据集)
    """
    with open(ClassifyConfig.DATA_PATH, 'r') as f:
        records = json.load(f)
    np.random.shuffle(records)
    total = len(records)
    val_num = int(total * ClassifyConfig.VAL_SPLIT)
    test_num = int(total * ClassifyConfig.TEST_SPLIT)
    val_data = DataGenerator([x for record in records[:val_num] for x in record])
    test_data = DataGenerator([x for record in records[val_num:val_num + test_num] for x in record])
    train_data = DataGenerator([x for record in records[val_num + test_num:] for x in record])
    return train_data, val_data, test_data

# -*- coding: utf-8 -*-
# @Time    : 2020/11/13
# @Author  : AaronJny
# @File    : classify_model.py
# @Desc    :
import tensorflow as tf
from tensorflow.keras import layers

from config import ClassifyConfig


class SiameseNetwork(tf.keras.Model):

    def __init__(self):
        super(SiameseNetwork, self).__init__()
        self.net_stage_1 = tf.keras.Sequential([
            layers.Input(shape=(*ClassifyConfig.IMAGE_SIZE, ClassifyConfig.IMAGE_CHANNELS)),
            layers.Conv2D(6, (3, 3), padding='same'),
            layers.MaxPooling2D((2, 2), 2),
            layers.Dropout(ClassifyConfig.DROPOUT_RATE),
            layers.ReLU(),
            layers.Conv2D(16, (5, 5)),
            layers.MaxPooling2D((2, 2), 2),
            layers.Dropout(ClassifyConfig.DROPOUT_RATE),
            layers.ReLU()
        ])
        self.net_stage_2 = tf.keras.Sequential([
            layers.Conv2D(6, (3, 3)),
            layers.MaxPooling2D((2, 2), 2),
            layers.Dropout(ClassifyConfig.DROPOUT_RATE),
            layers.ReLU(),
            layers.Flatten(),
            layers.Dense(84),
            # layers.Dropout(ClassifyConfig.DROPOUT_RATE),
            layers.ReLU(),
            layers.Dense(1, activation='sigmoid')
        ])

    @tf.function
    def call(self, inputs, training=None, mask=None):
        outs = []
        for x in inputs:
            out = self.net_stage_1(x)
            outs.append(out)
        out = tf.concat(outs, axis=-1)
        out = self.net_stage_2(out)
        return out


def load_classify_model():
    model = SiameseNetwork()
    model.build(
        [(None, *ClassifyConfig.IMAGE_SIZE, ClassifyConfig.IMAGE_CHANNELS),
         (None, *ClassifyConfig.IMAGE_SIZE, ClassifyConfig.IMAGE_CHANNELS)])
    model.load_weights(ClassifyConfig.MODEL_PATH)
    return model

# -*- coding: utf-8 -*-
# @Time    : 2020/11/13
# @Author  : AaronJny
# @File    : train_classify.py
# @Desc    :
from tensorflow.keras import losses, optimizers, callbacks

from classify_model import SiameseNetwork
from config import ClassifyConfig
from process_classify_dataset import load_dataset

model = SiameseNetwork()
model.build(
    [(None, *ClassifyConfig.IMAGE_SIZE, ClassifyConfig.IMAGE_CHANNELS),
     (None, *ClassifyConfig.IMAGE_SIZE, ClassifyConfig.IMAGE_CHANNELS)])
model.summary()
model.compile(optimizer=optimizers.Adam(ClassifyConfig.LR), loss=losses.binary_crossentropy, metrics=['accuracy'])

train_data, val_data, test_data = load_dataset()

checkpoint = callbacks.ModelCheckpoint(ClassifyConfig.MODEL_PATH,
                                       monitor='val_accuracy', save_best_only=True, save_weights_only=True)
# 这里稍微注意一下，实际上因为数据集太小了，模型非常容易过拟合，因此early_stopping非常容易被触发，所以这里我创建了但并没有使用，当有更大的数据集的时候可以考虑启用它
early_stopping = callbacks.EarlyStopping(monitor='val_accuracy', patience=20, verbose=1)
model.fit(train_data.for_fit(), epochs=300, validation_data=val_data.for_fit(), steps_per_epoch=train_data.steps,
          validation_steps=val_data.steps, callbacks=[checkpoint])

if ClassifyConfig.TEST_SPLIT:
    print(model.evaluate(test_data.for_fit(), steps=test_data.steps))

# -*- coding: utf-8 -*-
# @Time    : 2020/11/14
# @Author  : AaronJny
# @File    : solution.py
# @Desc    : 点选式验证码解决方案
import random
from glob import glob

import cv2
import numpy as np
from PIL import Image
from xyolo import YOLO

from classify_model import load_classify_model
from config import XYoloConfig
from process_classify_dataset import preprocess_image

config = XYoloConfig()
yolo = YOLO(config)
siamese_network = load_classify_model()


def detect(image_path):
    # 接收图片
    img = Image.open(image_path)
    img_array = np.asarray(img.convert('RGB'))
    # 提取两个标签
    word1_array = img_array[1:34, 186:216]
    word1_image = Image.fromarray(word1_array)
    word2_array = img_array[1:34, 216:246]
    word2_image = Image.fromarray(word2_array)
    # 检测文字
    results = yolo.detect_image(img)
    word_images = [word1_image, word2_image]
    word_arrays = [preprocess_image(x) for x in word_images]
    word_arrays = [x.reshape((1, *x.shape)) for x in word_arrays]
    correct_pos = []
    scores = []
    for word_array in word_arrays:
        sub_scores = []
        for index, (_, _, _, x1, y1, x2, y2) in enumerate(results):
            tmp_img_array = img_array[y1:y2, x1:x2]
            tmp_img_array = preprocess_image(Image.fromarray(tmp_img_array))
            tmp_img_array = tmp_img_array.reshape((1, *tmp_img_array.shape))
            score = siamese_network.predict([word_array, tmp_img_array])[0]
            sub_scores.append((score, index))
            print(score)
        sub_scores.sort(key=lambda x: -x[0])
        scores.append(sub_scores)
    score1, index1 = scores[0][0]
    score2, index2 = scores[1][0]
    if index1 == index2:
        if score1 > score2:
            index2 = scores[1][1][1]
        else:
            index1 = scores[0][1][1]
    for index in (index1, index2):
        correct_pos.append(results[index][3:])
    colors = [(0, 255, 0), (0, 0, 255)]
    for (x1, y1, x2, y2), color in zip(correct_pos, colors):
        cv2.rectangle(img_array, (x1, y1), (x2, y2), color, 2)
    cv2.imshow('result', img_array)
    for pos in correct_pos:
        print(pos)
    key = cv2.waitKey(0)
    cv2.destroyAllWindows()
    return key


def run():
    images = glob('./images/2??.png')
    correct_cnt = 0
    half_correct_cnt = 0
    total = 30
    for image_path in random.sample(images, total):
        key = detect(image_path)
        # 如果正确选中了标签中的两个字，就按数字1
        # 注意，所谓的正确是指顺序也要正确。
        # 第一个字用绿框圈起来，第二个字用红框
        if key == ord('1'):  # 数字1
            correct_cnt += 1
            half_correct_cnt += 1
        # 如果只正确选中了其中一个字，就按数字2
        # 注意，这里的正确也包括了顺序，顺序不对不认为是正确的
        elif key == ord('2'):  # 数字2
            half_correct_cnt += 1
        # 都不是就按任意键
        else:
            pass
    print('识别通过率: {}'.format(correct_cnt / total))
    print('至少识别出一个的概率: {}'.format(half_correct_cnt / total))


if __name__ == '__main__':
    run()