基于卷積神經網絡的圖像二分類檢測模型訓練與推理實現教程

前言

對于本教程，說白了，就是期望能通過一個程序判斷一張圖片是否為某個物體，或者說判斷一張圖片是否為某個缺陷。因為本教程是針對二分類問題，因此主要處理 是與不是 的問題，比如我的模型是判斷一張圖片是否為蘋果，那么拿一張圖片給模型去推理，他會得出這張圖是蘋果的概率，如果概率大于0.5(這個概率在0~1之間)，那么就判斷為是蘋果。

教程內容

使用了Python的 TensorFlow 和 Keras 庫構建卷積神經網絡來完成二分類模型訓練，以及使用模型完成對一張圖片的推理。原文鏈接：基于卷積神經網絡的圖像二分類檢測模型訓練與推理實現教程 | 幽絡源

大致步驟

1.確定環境與庫

2.準備數據集并且劃分

3.數據集的命名問題注意事項

4.編寫訓練代碼完成模型訓練

5.編寫推理代碼

6.測試二分類檢測結果

7.根據結果優化數據集

步驟1.確定環境與庫

Python環境是必備的，我這里所使用的Python版本為3.12.3

其次還需要以下庫，依次執行如下命令即可

pip install tensorflow
pip install pillow
pip install scipy

如圖

步驟2.準備數據集并且劃分

我這里以判斷圖片是否為沖溝缺陷來準備數據集，首先創建數據集的目錄結構，結構如下

data/train/true_sample/ false_sample/  val/true_sample/false_sample/

QQ_1734065732662

目錄解釋：

data：作為數據集的根目錄

train和val分別為訓練集、驗證集目錄

true_sample：正類樣本，也就是我這里需要把含有沖溝缺陷的圖放到這個目錄

false_sample：負類樣本，也就是這里需要將不含有沖溝缺陷的圖片放進這個目錄

如圖，我向train和val的true_sample目錄加入了一些含有沖溝缺陷的圖片

對于負類樣本，也不是無腦的只要不是沖溝就往里面放，而是放置你認為訓練出的模型可能會將什么識別為正類樣本。比如滑坡和沖溝其實是有聯系的，但不完全等同于，所以我需要將滑坡相關的，但是沒有沖溝情況的圖片放入false_sample中，期望模型不要誤判。再比如一個蘋果，你可能需要把紅色氣球作為父類樣本，防止模型將紅氣球判斷為是蘋果，如圖是我的負類樣本

步驟3.數據集的命名問題注意事項

關于數據集的命名，這里其實有一個坑，但是先說避免坑的做法：就像步驟2一樣，你的正類樣本所放置的目錄命名為true_sample、負類樣本所放置的目錄命名為false_sample就行了。（如果看不懂下面的解釋，按照這里做法做就是了）

然后我來解釋下是什么坑，對于這個二分類模型訓練，訓練出來的模型，無非是識別是與不是的問題，但是模型怎么區分我的哪個目錄放置的為是，哪個目錄放置的為不是呢，步驟4會給出訓練代碼，訓練代碼中的加載數據集時有一行如下代碼

class_mode='binary'  # 二分類（沖溝缺陷 vs. 非沖溝缺陷）

這表示我們要做二分類模型訓練，加上這行代碼，在加載數據集時，Keras 會自動將這些文件夾的名稱作為標簽，分別命名為1 和 0，如果被命名為標簽1 的目錄，則在推理時，概率越接近于1，則越表示是標為1的目錄的樣本，反之概率越接近于0，則越表示是標為0的目錄的樣本。而keras自動命名標簽1和0時是根據目錄名首字母的順序來的字，字母靠前的標為0，后者為1，true_sample的首字母為t，false_sample的首字母為f，因此false_sample標為0，true_sample標為1，這是符合我們的正常預期的。

反面例子：

如果我把正類樣本放置于名為defect的目錄，負類樣本放置于no_defect目錄會怎樣呢，按照如上解釋，defect目錄會被標為0，no_defect目錄會被標為1，這就和我們預期相反了，什么意思呢。我把正類樣本放置defect目錄中，其推理結果將會是越接近0，則越表示為正類了，因此這里特別需要注意（如果你要自定義目錄名的話）。

步驟4.編寫訓練代碼完成模型訓練

先直接上訓練代碼

from tensorflow.keras import layers, models
from tensorflow.keras.preprocessing.image import ImageDataGenerator
import tensorflow as tftrain_dir='data/train'
val_dir='data/val'# 設置圖像的尺寸和批量大小，不用改，保持150是最平衡的
IMG_HEIGHT = 150
IMG_WIDTH = 150
BATCH_SIZE = 12# 數據預處理與增強
train_datagen = ImageDataGenerator(rescale=1./255,  # 將像素值歸一化到 [0, 1] 區間shear_range=0.2,zoom_range=0.2,horizontal_flip=True
)validation_datagen = ImageDataGenerator(rescale=1./255)# 加載訓練和驗證數據
train_generator = train_datagen.flow_from_directory(train_dir,  # 訓練數據目錄target_size=(IMG_HEIGHT, IMG_WIDTH),  # 圖像尺寸batch_size=BATCH_SIZE,class_mode='binary'  # 二分類（沖溝缺陷 vs. 非沖溝缺陷）
)train_class_labels = train_generator.class_indices
print("訓練集自動標簽映射關系為："+str(train_class_labels))validation_generator = validation_datagen.flow_from_directory(val_dir,  # 驗證數據目錄target_size=(IMG_HEIGHT, IMG_WIDTH),batch_size=BATCH_SIZE,class_mode='binary'
)val_class_labels = validation_generator.class_indices
print("測試集自動標簽映射關系為："+str(val_class_labels))# 將數據生成器轉換為 tf.data.Dataset 并應用 repeat() 方法
train_dataset = tf.data.Dataset.from_generator(lambda: train_generator,output_signature=(tf.TensorSpec(shape=(None, IMG_HEIGHT, IMG_WIDTH, 3), dtype=tf.float32),tf.TensorSpec(shape=(None,), dtype=tf.int32))
)
train_dataset = train_dataset.repeat()  # 確保數據重復validation_dataset = tf.data.Dataset.from_generator(lambda: validation_generator,output_signature=(tf.TensorSpec(shape=(None, IMG_HEIGHT, IMG_WIDTH, 3), dtype=tf.float32),tf.TensorSpec(shape=(None,), dtype=tf.int32))
)
validation_dataset = validation_dataset.repeat()  # 確保數據重復# 構建模型
model = models.Sequential([layers.InputLayer(shape=(IMG_HEIGHT, IMG_WIDTH, 3)),  # 添加 Input 層layers.Conv2D(32, (3, 3), activation='relu'),layers.MaxPooling2D((2, 2)),layers.Conv2D(64, (3, 3), activation='relu'),layers.MaxPooling2D((2, 2)),layers.Conv2D(128, (3, 3), activation='relu'),layers.MaxPooling2D((2, 2)),layers.Flatten(),layers.Dense(128, activation='relu'),layers.Dense(1, activation='sigmoid')  # 輸出層，二分類問題
])# 編譯模型
model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])# 訓練模型
model.fit(train_dataset,steps_per_epoch=train_generator.samples // BATCH_SIZE,epochs=30,validation_data=validation_dataset,validation_steps=validation_generator.samples // BATCH_SIZE
)# 保存模型
model.save('defect_detector_model.keras')  # 使用 .keras 格式保存模型

使用這段代碼訓練數據集你唯一需要注意的是保持代碼文件于數據集文件在同一目錄，或者使用絕對路徑，如圖

QQ_1734070943655

我們啟動訓練代碼，可以看到控制臺在按照規定的輪次30在訓練中，而且可以看到我在訓練代碼中加入了輸出標簽映射關系來確保正類與負類的映射關系正確，如圖

QQ_1734071390702

訓練后，你會得到一個名為defect_detector_nodel.keras的文件，推理時會使用該模型進行推理

步驟5.編寫推理代碼

代碼如下：

import os
from tensorflow.keras.models import load_model
import numpy as np
from tensorflow.keras.preprocessing import image# 加載訓練好的模型
model = load_model('defect_detector_model.keras')  # 注意加載的是 .keras 格式# 設置輸入圖像的目標尺寸（與訓練時相同）
IMG_HEIGHT = 150
IMG_WIDTH = 150# 定義函數來加載并預測圖像
def predict_image(img_path):# 加載圖像并進行預處理img = image.load_img(img_path, target_size=(IMG_HEIGHT, IMG_WIDTH))img_array = image.img_to_array(img)  # 將圖像轉換為數組img_array = np.expand_dims(img_array, axis=0)  # 擴展維度，成為一個 batchimg_array = img_array / 255.0  # 歸一化處理（與訓練時一致）# 預測圖像類別prediction = model.predict(img_array)  # 返回的是一個包含概率的數組return prediction[0][0]  # 提取預測的概率值picPath=r"測試圖.jpg"
confidence = predict_image(picPath)
print("有沖溝缺陷的概率為："+str(confidence))

這段推理代碼中，我們加載了剛才訓練出的模型，然后使用了一張名為測試圖.jpg的圖片來進行推理，然后輸出他有缺陷的概率

步驟6.測試二分類檢測結果

我這里就不用一張圖片來測試了，我這里指定一個目錄，進行整個目錄來測試里面的圖片，還是附上我這個推理代碼吧

import os
from tensorflow.keras.models import load_model
import numpy as np
from tensorflow.keras.preprocessing import image# 加載訓練好的模型
model = load_model('defect_detector_model.keras')  # 注意加載的是 .keras 格式# 設置輸入圖像的目標尺寸（與訓練時相同）
IMG_HEIGHT = 150
IMG_WIDTH = 150# 定義函數來加載并預測圖像
def predict_image(img_path):# 加載圖像并進行預處理img = image.load_img(img_path, target_size=(IMG_HEIGHT, IMG_WIDTH))img_array = image.img_to_array(img)  # 將圖像轉換為數組img_array = np.expand_dims(img_array, axis=0)  # 擴展維度，成為一個 batchimg_array = img_array / 255.0  # 歸一化處理（與訓練時一致）# 預測圖像類別prediction = model.predict(img_array)  # 返回的是一個包含概率的數組return prediction[0][0]  # 提取預測的概率值# 測試目錄，包含要進行推理的圖像
testDir = r"D:\virtualTemp\pythonProject\CNN分類檢測\data\train\true_sample"
pics = os.listdir(testDir)
# 遍歷目錄中的所有圖片并進行預測
for pic in pics:picPath = os.path.join(testDir, pic)  # 獲取圖片的完整路徑# 獲取預測結果的置信度confidence = predict_image(picPath)# 輸出圖像的置信度和類別print(f"{pic} 置信度: {confidence:.4f}, 預測結果: {'有缺陷' if confidence >= 0.5 else '無缺陷'}")

我先使用正類樣本來測試，先看看拿訓練的數據如何，然后再用另外的圖片來測試

結果如下圖，正類樣本中只有一張圖判定為了無沖溝，但是我正類樣本中其實都應當是沖溝，而我有101張圖，因此這里正確率為99.009%

QQ_1734071615033