用python写一个基于卷积神经网络的猫狗分类代码
时间: 2023-11-16 19:07:16 浏览: 83
import tensorflow as tf
import numpy as np
import os
import cv2
# 设置训练数据路径和标签
train_data_path = 'train/'
train_labels = ['cat', 'dog']
# 设置训练参数
batch_size = 32
epochs = 10
learning_rate = 0.001
# 加载训练数据
train_images = []
train_labels = []
for label in train_labels:
path = os.path.join(train_data_path, label)
for file in os.listdir(path):
image_path = os.path.join(path, file)
image = cv2.imread(image_path)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
image = cv2.resize(image, (224, 224))
train_images.append(image)
train_labels.append(label)
train_images = np.array(train_images)
train_labels = np.array(train_labels)
# 编码标签
train_labels = np.where(train_labels == 'cat', 0, 1)
# 划分训练集和验证集
from sklearn.model_selection import train_test_split
X_train, X_val, y_train, y_val = train_test_split(train_images, train_labels, test_size=0.2, random_state=42)
# 数据增强
from tensorflow.keras.preprocessing.image import ImageDataGenerator
train_datagen = ImageDataGenerator(rotation_range=20,
zoom_range=0.15,
width_shift_range=0.2,
height_shift_range=0.2,
shear_range=0.15,
horizontal_flip=True,
fill_mode="nearest")
# 定义模型
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Conv2D, MaxPooling2D, Flatten, Dense, Dropout
model = Sequential()
model.add(Conv2D(32, (3, 3), activation='relu', input_shape=(224, 224, 3)))
model.add(MaxPooling2D((2, 2)))
model.add(Dropout(0.2))
model.add(Conv2D(64, (3, 3), activation='relu'))
model.add(MaxPooling2D((2, 2)))
model.add(Dropout(0.2))
model.add(Conv2D(128, (3, 3), activation='relu'))
model.add(MaxPooling2D((2, 2)))
model.add(Dropout(0.2))
model.add(Conv2D(256, (3, 3), activation='relu'))
model.add(MaxPooling2D((2, 2)))
model.add(Dropout(0.2))
model.add(Flatten())
model.add(Dense(64, activation='relu'))
model.add(Dropout(0.5))
model.add(Dense(1, activation='sigmoid'))
# 编译模型
model.compile(optimizer=tf.keras.optimizers.Adam(lr=learning_rate),
loss='binary_crossentropy',
metrics=['accuracy'])
# 训练模型
history = model.fit(train_datagen.flow(X_train, y_train, batch_size=batch_size),
steps_per_epoch=len(X_train) // batch_size,
epochs=epochs,
validation_data=(X_val, y_val))
# 测试模型
test_data_path = 'test/'
test_images = []
test_names = []
for file in os.listdir(test_data_path):
image_path = os.path.join(test_data_path, file)
image = cv2.imread(image_path)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
image = cv2.resize(image, (224, 224))
test_images.append(image)
test_names.append(file)
test_images = np.array(test_images)
predictions = model.predict(test_images)
predictions = np.where(predictions > 0.5, 'dog', 'cat')
# 保存预测结果
import pandas as pd
df = pd.DataFrame({'id': test_names, 'label': predictions.ravel()})
df.to_csv('submission.csv', index=False)
print('预测完成!')
阅读全文