【深度學習】雪碧圖識別（CNN 卷積神經網絡訓練）

目錄

• 雪碧圖識別（CNN 卷積神經網絡訓練）
  • 數據獲取
  • 制作數據集
  • 訓練模型
  • 模型調用
  • 通過爬蟲測試

雪碧圖識別（CNN 卷積神經網絡訓練）

鍍金的天空 是一個互聯網技能認證網站, 都是些爬蟲題目。其中有一道題 爬蟲-雪碧圖-2 需要使用到圖片識別。所以模仿 mnist ，用 CNN 卷積神經網絡訓練一個模型，准確率達到 99.90% 。

Github 項目源碼

# 基於 tensorflow 2.0
pip install -i https://pypi.tuna.tsinghua.edu.cn/simple tensorflow==2.0.0

# 項目構成
├─glidesky
    │  model.h5                      # 模型文件
    │  predict.py                    # 模型調用
    │  train.py                      # 模型訓練
    │
    ├─data_source
    │  │  data.h5                 # 數據集文件
    │  │  make_dataset.py         # 生成數據集
    │  │  spider.py               # 爬蟲
    │  │
    │  └─imgs                     # 存放采集圖片
    │
    ├─logs            # 訓練可視化日志
    │
    ├─test            # 測試圖片

數據獲取

數據獲取，首先找到一頁內容涵蓋 0-9 所有數字，然后用爬蟲將數字圖片采集下來，以供后續作為深度學習的數據集。

1. 因為每次請求都是不一樣的圖，但是數字是固定的，所以只要不斷請求同一頁即可
2. 每次請求只保留 10 張圖片，從而保證樣本數據的均勻分布
3. 采集過程較為耗時無聊，所以原計划采 100 萬張，后面只采集了 45 萬張

import re
import os
import uuid
import base64
import requests

from PIL import Image
from io import BytesIO
from bs4 import BeautifulSoup
from concurrent.futures import ThreadPoolExecutor

Cookie = 'your cookies'
headers = {
    'Accept': 'text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.9',
    'Accept-Encoding': 'gzip, deflate', 'Accept-Language': 'zh-CN,zh;q=0.9', 'Cache-Control': 'max-age=0',
    'Connection': 'keep-alive',
    'Cookie': Cookie,
    'Host': 'www.glidedsky.com',
    'Referer': 'http://www.glidedsky.com/level/web/crawler-basic-2?page=1',
    'Upgrade-Insecure-Requests': '1',
    'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/81.0.4044.92 Safari/537.36'
}


def get_img(text):
    """
    :param text: 獲取圖片模板
    :return:
    """
    img_str = re.findall('base64,(.*?)"', text)[0]
    img_fp = BytesIO(base64.b64decode(img_str.encode('utf-8')))
    img = Image.open(img_fp)
    return img


def crawler(url):
    text = requests.get(url, headers=headers).text
    img = get_img(text)
    rows = BeautifulSoup(text, 'lxml').find_all('div', class_="col-md-1")
    num_labels = list(str(123171140339373274129338158411319368))
    num_imgs = []
    for row in rows:
        for div in row.find_all('div'):
            css_name = div.get('class')[0].split(' ')[0]
            tag_x = re.findall(f'\.{css_name} \{{ background-position-x:(.*?)px \}}', text)
            tag_y = re.findall(f'\.{css_name} \{{ background-position-y:(.*?)px \}}', text)
            width = re.findall(f'\.{css_name} \{{ width:(.*?)px \}}', text)
            height = re.findall(f'\.{css_name} \{{ height:(.*?)px \}}', text)
            tag_x = abs(int(tag_x[0]))
            tag_y = abs(int(tag_y[0]))
            width = int(width[0])
            height = int(height[0])
            box = (tag_x, tag_y, tag_x + width, tag_y + height)
            num_imgs.append(img.crop(box))
    save_list = [str(i) for i in range(10)]
    for num_img, num_label in zip(num_imgs, num_labels):
        if num_label in save_list:
            file_name = f'./imgs/{num_label}_{uuid.uuid1()}.png'
            num_img = num_img.resize((20, 20))
            num_img.save(file_name)
            save_list.remove(num_label)

os.makedirs('./imgs', exist_ok=True)
urls = []
for _ in range(90000):
    url = f'http://www.glidedsky.com/level/web/crawler-sprite-image-2?page=999'
    urls.append(url)

pool = ThreadPoolExecutor(max_workers=20)
for result in pool.map(crawler, urls):
    ...

制作數據集

將所有圖片統一尺寸為 20*20 后，轉為灰度值；對應的標簽轉為獨熱編碼，通過 sklearn 隨機切分訓練集和測試集數據，最后保存為 h5 數據集文件。

1. 測試集最好不要和訓練集重疊，這樣才能評估模型的泛化能力
2. 保存數據集時，不事先進行預處理的原因：直接保存，數據文件大小為 190 M; 歸一化后再保存，則為 1.9 G
3. h5 層次數據格式第5代的版本（Hierarchical Data Format，HDF5），它是用於存儲科學數據的一種文件格式和庫文件
4. 獨熱編碼即一位有效編碼，比如 0-9 共十個數字，可以用一個長度為 10 的 list 表示。比如 2 是 [0,0,1,0,0,0,0,0,0,0]，9 是 [0,0,0,0,0,0,0,0,0,1]，以此類推；值可以通過 np.argmax() 獲取。

import os
import h5py
import numpy as np
from PIL import Image
from sklearn.model_selection import train_test_split

images = []
labels = []

for path in os.listdir('./imgs'):
    label = int(path.split('_')[0])
    label_one_hot = [0 if i != label else 1 for i in range(10)]
    labels.append(label_one_hot)

    img = Image.open('./imgs/' + path).resize((20, 20)).convert('L')
    img_arr = np.reshape(img, 20 * 20)
    images.append(img_arr)

# 拆分訓練集、測試集
train_images, test_images, train_labels, test_labels = train_test_split(images, labels, test_size=0.1, random_state=0)

with h5py.File('./data.h5', 'w') as f:
    f.create_dataset('train_images', data=np.array(train_images))
    f.create_dataset('train_labels', data=np.array(train_labels))
    f.create_dataset('test_images', data=np.array(test_images))
    f.create_dataset('test_labels', data=np.array(test_labels))

訓練模型

構建卷積神經網絡模型，喂數據（40萬的訓練集，4萬的測試集），對模型進行訓練。

1. 數據集是由白底黑字的灰度圖轉成矩陣（20*20）構成的，每個數字是在 0-255 之間，黑色 0，白色 255。預處理將其轉成黑底白字后，除以 255.0 即完成歸一化。數據歸一化后，有助於提高模型的准確度。為什么要歸一化
2. epochs，訓練集的數據全部被訓練一次，即為一個 epoch ; epochs 設置多次次合適，目前沒有萬能公式，需要不斷嘗試
3. 模型編譯的時候需要指定 optimizer 優化器、loss 損失函數、metrics 衡量指標等參數
4. 模型訓練的過程中，可以指定回調函數，比如保存模型、記錄日志等等
5. 訓練過的模型，可以加載后繼續訓練

經過訓練后，模型准確率達到了 99.91%

在測試集的准確率達到了 99.90%

import os
import h5py
import tensorflow as tf
from tensorflow.keras import layers, models


class Train:
    def __init__(self):
        # 最終模型存放路徑
        self.modelpath = './model.h5'

        # 定義模型
        if os.path.exists(self.modelpath):
            self.model = tf.keras.models.load_model(self.modelpath)
            print(f"{self.model} 模型加載成功，繼續訓練...")
        else:
            self.model = models.Sequential([
                # 第1層卷積，卷積核大小為3*3，32個，28*28為待訓練圖片的大小
                layers.Conv2D(32, (3, 3), activation='relu', input_shape=(20, 20, 1)),
                layers.MaxPooling2D((2, 2)),
                # 第2層卷積，卷積核大小為3*3，64個
                layers.Conv2D(64, (3, 3), activation='relu'),
                layers.MaxPooling2D((2, 2)),
                # 第3層卷積，卷積核大小為3*3，64個
                layers.Conv2D(64, (3, 3), activation='relu'),
                layers.Flatten(),
                layers.Dense(64, activation='relu'),
                layers.Dense(10, activation='softmax'),
            ])
        self.model.summary()

        # 讀取數據
        with h5py.File('./data_source/data.h5', 'r') as f:
            self.train_images = f['train_images'][()]
            self.train_labels = f['train_labels'][()]
            self.test_images = f['test_images'][()]
            self.test_labels = f['test_labels'][()]

        train_count, test_count = 400000, 40000
        self.train_images = self.train_images[:train_count].reshape((train_count, 20, 20, 1))
        self.train_labels = self.train_labels[:train_count]
        self.test_images = self.test_images[:test_count].reshape((test_count, 20, 20, 1))
        self.test_labels = self.test_labels[:test_count]

        # 數據處理 歸一化
        self.train_images = 1 - self.train_images / 255.0
        self.test_images = 1 - self.test_images / 255.0

    def train(self):
        # 可視化 tensorboard --logdir=D:\GitHub\antman\glidedsky\logs
        TensorBoardcallback = tf.keras.callbacks.TensorBoard(
            log_dir='logs',
            histogram_freq=1,
            write_graph=True,
            write_images=True,
            update_freq=10000
        )
        self.model.compile(optimizer='Adam',
                           loss='categorical_crossentropy',
                           metrics=['accuracy'])
        self.model.fit(self.train_images, self.train_labels, epochs=10, callbacks=[TensorBoardcallback])
        self.model.save(self.modelpath)

    def test(self):
        self.model = tf.keras.models.load_model(self.modelpath)
        test_loss, test_acc = self.model.evaluate(self.test_images, self.test_labels)
        print("准確率: %.4f，共測試了%d張圖片 " % (test_acc, len(self.test_labels)))


if __name__ == "__main__":
    app = Train()
    app.train()
    app.test()

模型調用

模型的調用輸入：歸一化的三維矩陣（尺寸 20*20，需要轉成黑底白字） 構成的列表；輸出： 標簽獨熱編碼 構成的列表。

1. 模型的輸入應該與訓練時的數據使用相同的處理方式
2. 獨熱編碼取最大值的下標，即代表的標簽數字

最后，測試准備的這 5 張圖，模型都能正確識別

import numpy as np
import tensorflow as tf
from PIL import Image

class Predict(object):
    def __init__(self):
        self.cnn = tf.keras.models.load_model('./model.h5')

    def predict(self, image_path):
        # 以黑白方式讀取圖片
        img = Image.open(image_path).resize((20, 20)).convert('L')
        img_arr = 1 - np.reshape(img, (20, 20, 1)) / 255.0
        x = np.array([img_arr])

        # API refer: https://keras.io/models/model/
        y = self.cnn.predict(x)

        # 因為x只傳入了一張圖片，取y[0]即可
        # np.argmax()取得最大值的下標，即代表的數字
        print(image_path)
        print(y[0])
        print('        -> Predict digit', np.argmax(y[0]))

if __name__ == "__main__":
    app = Predict()
    app.predict('./test/0.png')
    app.predict('./test/3.png')
    app.predict('./test/4.png')
    app.predict('./test/7.png')
    app.predict('./test/9.png')

通過爬蟲測試

直接在爬蟲中調用模型，因為存在概率問題，所以多跑幾次，就可以解決這道爬蟲題目。有興趣的話，可以參考完整的代碼 glidedsky 通關筆記