一个欲儿的博客

一个欲儿的博客

Python sklearn 之猫狗识别
2024-12-28

import os
from PIL import Image
import numpy as np
from sklearn.svm import SVC
from sklearn.metrics import accuracy_score, confusion_matrix, ConfusionMatrixDisplay
import matplotlib.pyplot as plt
import joblib

# 加载数据和标签
def load_images_and_labels(folder_path, target_size=(64, 64)):
    images = []
    labels = []
    for file in os.listdir(folder_path):
        file_path = os.path.join(folder_path, file)
        try:
            # 提取标签:文件名包含 'cat' 是猫(0),包含 'dog' 是狗(1)
            label = 0 if "cat" in file.lower() else 1
            # 加载并预处理图片
            img = Image.open(file_path).convert('RGB')  # 转为RGB
            img = img.resize(target_size)  # 调整图像大小
            images.append(np.array(img))
            labels.append(label)
        except Exception as e:
            print(f"Error loading image {file_path}: {e}")
    return np.array(images), np.array(labels)

# 加载训练集和测试集
print("Loading train dataset...")
train_images, train_labels = load_images_and_labels("dataset/train")
print("Loading test dataset...")
test_images, test_labels = load_images_and_labels("dataset/test")
#print("Unique labels in training data:", np.unique(train_labels))

# 扁平化图像(将 64x64x3 转为一维向量)
train_features = train_images.reshape(len(train_images), -1)
test_features = test_images.reshape(len(test_images), -1)

# 创建并训练 SVM 模型
print("Training SVM model...")
model = SVC(kernel='linear', verbose=True)  # 使用线性核
#model.fit(train_features, train_labels)

batch_size = 100  # 设置每批次的大小
n_batches = len(train_features) // batch_size  # 计算总批次数

# 手动训练过程,打印每个批次的进度
for i in range(n_batches // 2):
    start_idx = i * batch_size
    end_idx = (i + 1) * batch_size
    
    # 获取当前批次的标签
    batch_labels = train_labels[start_idx:end_idx]
    
    # 获取倒数第 i 批次的标签
    reverse_start_idx = (n_batches - 1 - i) * batch_size
    reverse_end_idx = (n_batches - i) * batch_size if (n_batches - i) * batch_size <= len(train_features) else len(train_features)
    
    # 获取倒数批次的数据
    reverse_batch_labels = train_labels[reverse_start_idx:reverse_end_idx]
    
    # 合并当前批次和倒数批次的标签与特征
    combined_labels = np.concatenate([batch_labels, reverse_batch_labels])
    combined_features = np.vstack([train_features[start_idx:end_idx], train_features[reverse_start_idx:reverse_end_idx]])

    # 检查合并后的批次是否包含两个标签
    if len(np.unique(combined_labels)) < 2:
        print(f"Skipping batch pair {i+1}/{n_batches//2} because it contains only one class")
        continue  # 如果合并后的批次只有一个标签,跳过该批次

    # 训练模型
    model.fit(combined_features, combined_labels)

    # 打印当前进度
    print(f"Batch pair {i+1}/{n_batches//2} processed")

joblib.dump(model, 'CatorDog.pkl')    
    
# 测试模型
print("Evaluating model...")
# 设置每批次的大小
batch_size = 100
n_batches = len(test_features) // batch_size  # 计算总批次数

# 存储所有批次的预测标签
all_predictions = []

# 手动分批测试过程
print("Evaluating model in batches...")
for i in range(n_batches + 1):
    start_idx = i * batch_size
    end_idx = (i + 1) * batch_size
    if end_idx > len(test_features):
        end_idx = len(test_features)
    
    # 获取当前批次的特征数据
    batch_features = test_features[start_idx:end_idx]
    
    # 进行预测
    batch_predictions = model.predict(batch_features)
    
    # 将当前批次的预测结果添加到列表中
    all_predictions.extend(batch_predictions)
    
    # 打印当前批次的进度
    print(f"Batch {i+1}/{n_batches+1} processed")

# 计算整体的准确率
accuracy = accuracy_score(test_labels, all_predictions)
print(f"Test Accuracy: {accuracy * 100:.2f}%")

数据集链接如下

通过网盘分享的文件:dataset.zip
链接: https://pan.baidu.com/s/1a9uXIGtJ-Pwohc-CaimOPg?pwd=yuer 提取码: yuer


发表评论: