Captcha solver

train.py

from tensorflow.keras import Input, Model
from tensorflow.keras.models import load_model
from tensorflow.keras.layers import (
    Conv2D,
    Activation,
    MaxPooling2D,
    Dropout,
    Flatten,
    Dense,
)
from tensorflow.keras.applications import MobileNetV2
from tensorflow.keras.optimizers import RMSprop
from tensorflow.keras.metrics import top_k_categorical_accuracy
from PIL import Image
from tensorflow.keras.preprocessing.image import img_to_array
import numpy as np
from tensorflow.keras.callbacks import (
    ModelCheckpoint,
    TensorBoard,
    LambdaCallback,
    TerminateOnNaN,
)
from os.path import isfile
from subprocess import Popen, PIPE, DEVNULL
from urllib.request import urlopen
from urllib.error import URLError
from time import sleep
from base64 import b64decode
from json import loads
from io import BytesIO
import backoff
from hashlib import blake2b
import tensorflow as tf
import matplotlib.pyplot as plt

alphabet = "abcdefghijklmnopqrstuvwxyz0123456789".upper()
char_size = len(alphabet)
categories = 6
model_path = "./checkpoints/model.hdf5"


def text2vec(label):
    vecs = []
    for i, char in enumerate(label):
        vector = np.zeros(char_size)
        vector[alphabet.index(char)] = 1
        vecs.append(vector)
    return np.array(vecs)


@backoff.on_exception(backoff.expo, URLError)
def get():
    return urlopen(
        "http://0.0.0.0:3000/botdetect.php?get=image&c=DemoCaptcha&t=239c90fd000cfe78d241c9d2f7743188"
    )


def generator(port=3000):
    while True:
        x = []
        y = []

        proc = Popen(
            ["php", "-S", f"0.0.0.0:{port}"],
            cwd="./bdc-php-free/v4.2.5/examples/t_api-captcha~demo",
            stderr=PIPE,
            stdout=DEVNULL,
        )
        sleep(0.05)

        while len(x) < 32:
            r = get()
            resp = loads(r.read())
            img = BytesIO(b64decode(resp["image"]))
            img = img_to_array(Image.open(img).convert("L")).reshape((50, 250, 1))
            img /= 127.5
            img -= 1.0
            code = resp["code"]
            x.append(img)
            y.append(text2vec(code))

        proc.kill()
        x = np.array(x)
        y = np.array(y).reshape((32, -1))
        if np.any(np.isnan(x)) or np.any(np.isnan(y)):
            print("OPOO got NAN")
            continue
        yield (x, y)


def vec2text(label):
    arr = [l.argmax().tolist() for l in label.reshape((categories, char_size))]
    ret = [alphabet[l] for l in arr]
    return ret


class CaptchaModel:
    # ref: https://github.com/yeguixin/captcha_solver/tree/master/src/models
    def __init__(self):
        self.filewriter = tf.summary.create_file_writer("./logs/")
        if isfile(model_path):
            self.model = load_model(model_path)
            return
        inp = Input(shape=(50, 250, 1))
        x = Conv2D(32, (3, 3), strides=(1, 1), padding="same")(inp)
        x = Activation("relu")(x)
        x = MaxPooling2D(pool_size=(2, 2), padding="same")(x)
        x = Dropout(0.5)(x)
        x = Conv2D(64, (3, 3), strides=(1, 1), padding="same")(x)
        x = Activation("relu")(x)
        x = MaxPooling2D(pool_size=(2, 2), padding="same")(x)
        x = Dropout(0.5)(x)
        x = Conv2D(128, (3, 3), strides=(1, 1), padding="same")(x)
        x = Activation("relu")(x)
        x = MaxPooling2D(pool_size=(2, 2), padding="same")(x)
        x = Dropout(0.5)(x)
        x = Conv2D(256, (3, 3), strides=(1, 1), padding="same")(x)
        x = Activation("relu")(x)
        x = MaxPooling2D(pool_size=(2, 2), padding="same")(x)
        x = Dropout(0.5)(x)
        x = Conv2D(512, (3, 3), strides=(1, 1), padding="same")(x)
        x = Activation("relu")(x)
        x = MaxPooling2D(pool_size=(2, 2), padding="same")(x)
        x = Dropout(0.5)(x)
        x = Flatten()(x)
        x = Dense(1024)(x)
        x = Activation("relu")(x)
        x = Dropout(0.5)(x)
        x = Dense(char_size * categories, activation="softmax")(x)
        model = Model(inp, x)
        model.compile(
            optimizer=RMSprop(learning_rate=0.0001, clipvalue=0.5),
            loss=["categorical_crossentropy"],
            metrics=[self.accuracy],
        )
        self.model = model

    def accuracy(self, y_true, y_pred):
        return top_k_categorical_accuracy(y_true, y_pred, k=categories)

    def test(self, batch, logs=None):
        print("")
        (images, labels) = next(generator(port="9000"))
        for i, image in enumerate(images[:5]):
            label = "".join(vec2text(labels[i].reshape((1, -1))))
            predicted = self.model.predict(image.reshape((1, 50, 250, 1)))
            predicted_label = "".join(vec2text(predicted))
            print(
                blake2b(image.tobytes()).hexdigest()[:5],
                predicted[0][:5],
                vec2text(predicted),
                vec2text(labels[i].reshape((1, -1))),
            )
            with self.filewriter.as_default():
                tf.summary.image(
                    f"{label} {predicted_label}", image.reshape((1, 50, 250, 1)), step=0
                )

    def train(self):
        checkpoint = ModelCheckpoint(
            model_path, save_best_only=True, monitor="accuracy", mode="max"
        )
        board = TensorBoard(log_dir="./logs", write_images=True)
        log = LambdaCallback(on_epoch_end=self.test)
        term = TerminateOnNaN()
        self.model.fit_generator(
            generator(),
            steps_per_epoch=20,
            epochs=100000,
            callbacks=[checkpoint, board, log, term],
        )


if __name__ == "__main__":
    model = CaptchaModel()
    model.train()