KNN算法

import os
import numpy as np
from PIL import Image
from sklearn.model_selection import KFold
from sklearn.neighbors import KNeighborsClassifier
from sklearn.metrics import accuracy_score
from torchvision import transforms
from torch.utils.data import Dataset, DataLoader

class ImageDataset(Dataset):
    def __init__(self, image_dir, transform=None):
        self.image_dir = image_dir
        self.transform = transform
        self.image_paths = []
        self.labels = []
        
        folder_to_label = {'1': 0, '2': 1}
        
        for label_folder, label in folder_to_label.items():
            class_dir = os.path.join(image_dir, label_folder)
            if os.path.isdir(class_dir):
                for file_name in os.listdir(class_dir):
                    file_path = os.path.join(class_dir, file_name)
                    if file_path.endswith(('.png', '.jpg', '.jpeg')):
                        self.image_paths.append(file_path)
                        self.labels.append(label)
    
    def __len__(self):
        return len(self.image_paths)
    
    def __getitem__(self, idx):
        image = Image.open(self.image_paths[idx]).convert('RGB')
        label = self.labels[idx]
        if self.transform:
            image = self.transform(image)
        return image, label

transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

dataset = ImageDataset(image_dir='/path/to/images', transform=transform)

# 将数据集中的所有图像转换为特征向量并存储
data = []
labels = []

for i in range(len(dataset)):
    image, label = dataset[i]
    data.append(image.numpy().flatten())  # 将图像转换为一维特征向量
    labels.append(label)

data = np.array(data)
labels = np.array(labels)

# KFold交叉验证
kf = KFold(n_splits=5, shuffle=True, random_state=42)
accuracy_list = []

for train_index, test_index in kf.split(data):
    X_train, X_test = data[train_index], data[test_index]
    y_train, y_test = labels[train_index], labels[test_index]
    
    # 定义KNN分类器
    knn = KNeighborsClassifier(n_neighbors=5)  # 使用k=5的KNN分类器
    knn.fit(X_train, y_train)  # 训练KNN分类器
    
    # 预测并计算精度
    y_pred = knn.predict(X_test)
    accuracy = accuracy_score(y_test, y_pred)
    accuracy_list.append(accuracy)

# 输出平均准确率
mean_accuracy = np.mean(accuracy_list)
print(f'平均准确率: {mean_accuracy:.4f}')