captcha.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import tensorflow as tf\n",
    "from tensorflow import keras\n",
    "import numpy as np\n",
    "import PIL\n",
    "from functools import reduce"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "IMG_HEIGHT = 128\n",
    "IMG_WIDTH = 128\n",
    "def path2img(path):\n",
    "    img = PIL.Image.open(path)\n",
    "    return img\n",
    "def path2imgarray(path):\n",
    "    import cv2\n",
    "    img = PIL.Image.open(path).resize((IMG_WIDTH, IMG_HEIGHT))\n",
    "    img = np.array(img)\n",
    "    def binaryzation(img):\n",
    "        cv_img = cv2.cvtColor(img,cv2.COLOR_RGB2GRAY)\n",
    "        return cv_img\n",
    "\n",
    "    img = binaryzation(img) if len(img.shape) == 3 else img / 255.0\n",
    "    img = img.reshape(IMG_HEIGHT, IMG_WIDTH, 1)\n",
    "    return img\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 识别验证码模型\n",
    "import functools\n",
    "class CaptchaIdentifier(keras.Model):\n",
    "    \n",
    "    labelLen = 4\n",
    "    \n",
    "    charSet = [chr(ord('A') + i) for i in range(26)]\n",
    "    charSetLen = len(charSet)\n",
    "    def text2vector(self, text):\n",
    "        vectors = np.zeros([self.labelLen, self.charSetLen], dtype=np.float32)\n",
    "        for i, c in enumerate(text):\n",
    "            vectors[i, self.charSet.index(c)] = 1.0\n",
    "        return vectors\n",
    "    def vector2text(self, vectors):\n",
    "        return ''.join(map(lambda vector: chr(ord('A') + vector.index(np.max(vector))),vectors))\n",
    "    def __init__(self):\n",
    "        layers = keras.layers\n",
    "        super().__init__()\n",
    "        self.model = tf.keras.Sequential([\n",
    "            layers.Input(shape=(IMG_HEIGHT, IMG_WIDTH, 1)), # (128, 128, 1)\n",
    "\n",
    "            layers.Conv2D(16, (7,7), activation='relu'),    # (128, 128, 16)\n",
    "            layers.MaxPooling2D((2, 2)),                    # (64, 64, 16)\n",
    "            layers.BatchNormalization(),\n",
    "            \n",
    "            layers.Conv2D(128, (3, 3), activation='relu'),  # (64, 64, 128)\n",
    "            layers.MaxPooling2D((2, 2),),                   # (32, 32, 128)\n",
    "            layers.BatchNormalization(),\n",
    "            \n",
    "            layers.Flatten(),                             \n",
    "            layers.Dense(1024, activation='relu', ),        # (1024)\n",
    "            layers.Dropout(0.2),\n",
    "            \n",
    "            layers.Dense(CaptchaIdentifier.labelLen * CaptchaIdentifier.charSetLen), # (144)\n",
    "            layers.Reshape([CaptchaIdentifier.labelLen, CaptchaIdentifier.charSetLen]), # 144 -> (4, 36)\n",
    "\n",
    "            layers.Softmax()                              # (4, 36) -> (4, 36)\n",
    "        ])\n",
    "        \n",
    "                \n",
    "    def call(self, x):\n",
    "        return self.model(x)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create an instance of the model\n",
    "identifier = CaptchaIdentifier()\n",
    "\n",
    "trainLoss = tf.keras.metrics.Mean(name='trainLoss')\n",
    "valLoss = tf.keras.metrics.Mean(name='valLoss')\n",
    "trainAccuracy = tf.keras.metrics.CategoricalAccuracy(name='trainAccuracy')\n",
    "valAccuracy = tf.keras.metrics.CategoricalAccuracy(name='valAccuracy')\n",
    "calLoss = tf.keras.losses.CategoricalCrossentropy()\n",
    "\n",
    "identifier.build(input_shape=(None, IMG_HEIGHT, IMG_WIDTH, 1))\n",
    "identifier.summary()\n",
    "# 查看每一层\n",
    "print('model')\n",
    "\n",
    "for layer in identifier.model.layers:\n",
    "    print(layer.name, layer.output_shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def getTrainStap(model):\n",
    "    optimizer = tf.keras.optimizers.Adam()\n",
    "    @tf.function\n",
    "    def trainStep(images, labels):\n",
    "        with tf.GradientTape() as tape:\n",
    "            predictions = model(images)\n",
    "            # print(f'predict shape {predictions.shape}')\n",
    "            loss = calLoss(labels, predictions)\n",
    "            # print('calLoss end')\n",
    "        gradients = tape.gradient(loss, model.trainable_variables)\n",
    "        # print(1)\n",
    "        optimizer.apply_gradients(zip(gradients, model.trainable_variables))\n",
    "        # print(1)\n",
    "        trainLoss(loss)\n",
    "        trainAccuracy(labels, predictions)\n",
    "        # print(1)\n",
    "    return trainStep\n",
    "\n",
    "def getcValStep(model):\n",
    "    @tf.function\n",
    "    def valStep(images, labels):\n",
    "        predictions = model(images, training=False)\n",
    "        loss = calLoss(labels, predictions)\n",
    "\n",
    "        valLoss(loss)\n",
    "        valAccuracy(labels, predictions)\n",
    "    return valStep"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "def basenameWithoutExt(path):\n",
    "    return os.path.basename(path).split('.')[0]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "BATCH_SIZE = 50\n",
    "\n",
    "    \n",
    "trainDir = './dataset1/train/'\n",
    "testDir = './dataset1/test/'\n",
    "# noise 和文件名是一一对应的\n",
    "dataset = tf.data.Dataset.from_tensor_slices(\n",
    "    ([path2imgarray(path) for path in tf.io.gfile.glob(trainDir + '*.jpg')], \n",
    "    [identifier.text2vector(basenameWithoutExt(path).replace('_resample', '')) for path in tf.io.gfile.glob(trainDir + '*.jpg')])\n",
    ")\n",
    "dataset = dataset.shuffle(buffer_size=1000)\n",
    "dataset = dataset.batch(BATCH_SIZE)\n",
    "dataset = dataset.prefetch(buffer_size=tf.data.experimental.AUTOTUNE)\n",
    "\n",
    "valDataset = tf.data.Dataset.from_tensor_slices(\n",
    "    ([path2imgarray(path) for path in tf.io.gfile.glob(testDir + '*.jpg')],\n",
    "    [identifier.text2vector(basenameWithoutExt(path).replace('_resample', '')) for path in tf.io.gfile.glob(testDir + '*.jpg')])\n",
    ")\n",
    "\n",
    "valDataset = valDataset.shuffle(buffer_size=1000)\n",
    "valDataset = valDataset.batch(BATCH_SIZE)\n",
    "valDataset = valDataset.prefetch(buffer_size=tf.data.experimental.AUTOTUNE)\n",
    "\n",
    "print(len(dataset))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "text = [*dataset.take(1)][0]\n",
    "print(text[1].shape)\n",
    "# print(text[1][0])\n",
    "print(identifier.vector2text(np.array(text[1][0]).tolist()))\n",
    "\n",
    "print(text[0].shape)\n",
    "plt.imshow(text[0][0])\n",
    "\n",
    "# print(text[0][0])\n",
    "# identifier.text2vector(text.numpy().decode('utf-8'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "# train face1\n",
    "trainStep = getTrainStap(identifier)\n",
    "valStep = getcValStep(identifier)\n",
    "EPOCHS = 500\n",
    "\n",
    "trainLossHistory = []\n",
    "valLossHistory = []\n",
    "trainAccHistory = []\n",
    "valAccHistory = []\n",
    "epochsRange = []\n",
    "\n",
    "\n",
    "\n",
    "start = time.time()\n",
    "for epoch in range(EPOCHS):\n",
    "    trainLoss.reset_states()\n",
    "    valLoss.reset_states()\n",
    "    trainAccuracy.reset_states()\n",
    "    valAccuracy.reset_states()\n",
    "        \n",
    "    for images, labels in dataset:\n",
    "        trainStep(images, labels)\n",
    "    \n",
    "    # random validate\n",
    "    for images, labels in valDataset.take(1):\n",
    "        valStep(images, labels)\n",
    "\n",
    "    \n",
    "    print(\n",
    "        f'Epoch {epoch + 1}, '\n",
    "        f'\\nTrain      Loss: {trainLoss.result()}, '\n",
    "        f'Train      Acc: {trainAccuracy.result()} '\n",
    "        f'\\nValidation Loss: {valLoss.result()}'\n",
    "        f'Validation  Acc: {valAccuracy.result()}'\n",
    "        f' - {int((time.time() - start) / (epoch + 1) * (EPOCHS - epoch)) / 60} minutes left,\\n '\n",
    "    )\n",
    "    trainLossHistory.append(trainLoss.result())\n",
    "    valLossHistory.append(valLoss.result())\n",
    "    trainAccHistory.append(trainAccuracy.result())\n",
    "    valAccHistory.append(valAccuracy.result())\n",
    "    epochsRange.append(epoch)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "minloss = np.min(valLossHistory)\n",
    "maxacc = np.max(valAccHistory)\n",
    "epochln = len(epochsRange)\n",
    "\n",
    "\n",
    "for i in range(4):\n",
    "  # dpi = 150\n",
    "  plt.figure(figsize=(5, 5), dpi=80)\n",
    "  # 画虚线 y = minloss\n",
    "  plt.plot([0, epochln], [minloss,minloss],linestyle='--',color='black')\n",
    "  plt.text(0, minloss, '%.5f'%minloss, fontsize=10) \n",
    "  plt.plot(epochsRange[5 * i:], trainLossHistory[5 * i:], label='Training Loss')\n",
    "  plt.plot(epochsRange[5 * i:], valLossHistory[5 * i:], label='Validation Loss')\n",
    "  # 图例\n",
    "  plt.legend(loc='upper right')\n",
    "  plt.show()\n",
    "\n",
    "for i in range(4):\n",
    "  plt.figure(figsize=(5, 5), dpi=80)\n",
    "  # 画虚线 y = maxacc\n",
    "  plt.plot([0, epochln], [maxacc,maxacc],linestyle='--',color='black')\n",
    "  plt.text(0, maxacc, '%.5f'%maxacc, fontsize=10)\n",
    "  plt.plot(epochsRange[5 * i:], trainAccHistory[5 * i:], label='TrainingAcc')\n",
    "  plt.plot(epochsRange[5 * i:], valAccHistory[5 * i:], label='ValidationAcc')\n",
    "  # 图例\n",
    "  plt.legend(loc='upper right')\n",
    "  plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "text = [*valDataset.take(1)][0]\n",
    "print(text[1][0])\n",
    "print(identifier(text[0][:1]))\n",
    "print(identifier.vector2text(np.array(text[1][0]).tolist()))\n",
    "print(identifier.vector2text(np.array(identifier(text[0][:1],training=False)[0]).tolist()))\n",
    "change = tf.keras.Sequential([\n",
    "              # keras.layers.RandomFlip(\"horizontal\"),\n",
    "            keras.layers.RandomRotation(0.01),\n",
    "            keras.layers.RandomZoom(0.1),\n",
    "])\n",
    "plt.imshow(change(text[0][0]))\n"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "2536a8b7348cda3df24488cb4703f2095d9cd957069c8076928c7127913566d2"
  },
  "kernelspec": {
   "display_name": "Python 3.9.10 ('base')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.10"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}