当前位置：首页 > article >正文

基础：用卷积神经网络（CNN）进行猫狗图像分类

article 2024/11/18 22:25:16

在本篇教程中，我们将通过卷积神经网络（CNN）实现一个简单的猫狗图像分类器。我们将介绍如何处理数据、构建CNN模型、训练模型并在测试集上进行预测。最终，你将能够用这个模型对未知图像进行猫狗分类。
在这里插入图片描述

1. 环境准备

首先，确保你已经安装了以下库：

tensorflow（用于深度学习模型）
opencv-python（用于图像处理）
numpy（用于数值计算）
matplotlib（用于数据可视化）
tqdm（用于显示进度条）

pip install tensorflow opencv-python numpy matplotlib tqdm

2. 数据预处理

2.1 数据集结构

假设你已经准备好了猫狗分类的数据集，其中包含两个文件夹：train/ 和 test/。每个文件夹下包含多个 .jpg 格式的图像。训练集中的图像每个都对应一个标签，标签通过文件名中的cat或dog来标识。例如，cat.0.jpg代表一只猫，dog.1.jpg代表一只狗。

2.2 标签编码

为了将标签转换为机器学习模型可以处理的格式，我们使用独热编码（One-Hot Encoding）。具体来说，如果图像是猫，则标签为[1, 0]；如果图像是狗，则标签为[0, 1]。

def label_img(img):
    word_label = img.split('.')[-3]
    if word_label == 'cat':
        return [1, 0]
    elif word_label == 'dog':
        return [0, 1]

2.3 读取训练数据

我们将读取训练数据，将每个图像调整为固定大小，并转换为灰度图像。然后，将图像和标签组合成一个训练数据集，并进行随机打乱。

def create_train_data():
    training_data = []
    for img in tqdm(os.listdir(train_dir)):
        if not img.endswith('.jpg'):
            continue
        label = label_img(img)
        path = os.path.join(train_dir, img)
        img = cv2.imread(path, cv2.IMREAD_GRAYSCALE)  # 读取灰度图
        img = cv2.resize(img, (img_size, img_size))  # 调整图像大小
        training_data.append([np.array(img), np.array(label)])
    shuffle(training_data)
    return training_data

2.4 读取测试数据

测试数据与训练数据相似，但是没有标签。我们仅将测试数据中的图像加载并调整为固定大小。

def process_test_data():
    testing_data = []
    for img in tqdm(os.listdir(test_dir)):
        if not img.endswith('.jpg'):
            continue
        path = os.path.join(test_dir, img)
        img_num = img.split('.')[0]
        img = cv2.imread(path, cv2.IMREAD_GRAYSCALE)
        img = cv2.resize(img, (img_size, img_size))
        testing_data.append([np.array(img), img_num])
    shuffle(testing_data)
    return testing_data

3. 构建卷积神经网络（CNN）

接下来，我们将构建一个卷积神经网络来进行图像分类。该网络由多个卷积层、池化层和全连接层组成。

3.1 定义模型结构

model = Sequential()

# 输入层和卷积层
model.add(Conv2D(32, (5, 5), activation='relu', input_shape=(img_size, img_size, 1), padding='same'))
model.add(MaxPooling2D(pool_size=(2, 2)))

# 添加更多卷积层和池化层
model.add(Conv2D(64, (5, 5), activation='relu', padding='same'))
model.add(MaxPooling2D(pool_size=(2, 2)))

model.add(Conv2D(128, (5, 5), activation='relu', padding='same'))
model.add(MaxPooling2D(pool_size=(2, 2)))

model.add(Conv2D(64, (5, 5), activation='relu', padding='same'))
model.add(MaxPooling2D(pool_size=(2, 2)))

model.add(Conv2D(32, (5, 5), activation='relu', padding='same'))
model.add(MaxPooling2D(pool_size=(2, 2)))

# 全连接层
model.add(Flatten())
model.add(Dense(1024, activation='relu'))
model.add(Dropout(0.8))

# 输出层
model.add(Dense(2, activation='softmax'))

3.2 编译模型

使用Adam优化器，损失函数采用categorical_crossentropy，因为这是一个多类分类问题。

model.compile(optimizer=Adam(lr), loss='categorical_crossentropy', metrics=['accuracy'])

4. 准备数据并训练模型

4.1 划分训练和验证集

我们将训练数据分为训练集和验证集。

train = train_data[:-500]
test = train_data[-500:]

X = np.array([i[0] for i in train], dtype=np.float64).reshape(-1, img_size, img_size, 1)
y = np.array([i[1] for i in train], dtype=np.float64)
Xtest = np.array([i[0] for i in test], dtype=np.float64).reshape(-1, img_size, img_size, 1)
ytest = np.array([i[1] for i in test], dtype=np.float64)

4.2 训练模型

我们使用fit方法训练模型，并将验证数据传入以监控验证集上的性能。

model.fit(X, y, epochs=3, validation_data=(Xtest, ytest), batch_size=32, verbose=1)

5. 在测试集上进行预测

5.1 加载测试数据并进行预测

我们将加载测试数据并对每个图像进行分类预测。

test_data = process_test_data()

# 可视化预测结果
fig = plt.figure()
for num, data in enumerate(test_data[:16]):
    img_num = data[1]
    img_data = data[0]
    
    y = fig.add_subplot(4, 4, num + 1)
    orig = img_data
    data = img_data.reshape(1, img_size, img_size, 1)
    
    model_out = model.predict(data)[0]
    
    if np.argmax(model_out) == 1:
        label = 'Dog'
    else:
        label = 'Cat'
    
    y.imshow(orig, cmap='gray')
    plt.title(label)
    y.axes.get_xaxis().set_visible(False)
    y.axes.get_yaxis().set_visible(False)

plt.tight_layout()
plt.show()