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

使用KNN预测图像

news 来源：原创 2025/4/29 7:06:31

题目：使用KNN算法实现机器学习给我一个水果的图片我能预测出这个是什么水果

import cv2
import numpy as np
import os
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
from sklearn.neighbors import KNeighborsClassifier
from sklearn.metrics import classification_report, accuracy_score# 数据集路径和水果标签
dataset_path = './fruits'  # 你的数据集在这个文件夹里
fruit_labels = ['apple', 'banana', 'grape', 'orange', 'pear']  # 定义水果标签
fruit_dict = {fruit: idx for idx, fruit in enumerate(fruit_labels)}  # 构建标签字典（字符串到数字的映射）# 特征提取函数：读取图像，调整大小并展平
def extract_features(image_path):# 读取图像img = cv2.imread(image_path)if img is None:print(f"Warning: Could not read image at {image_path}")return None  # 如果图像读取失败，返回 None# 将图像调整为统一的大小（100x100）img = cv2.resize(img, (100, 100))# 展平图像为一维数组，方便用于训练img = img.flatten()return img# 加载数据集：遍历数据集目录，提取每张图像的特征并保存标签
def load_data(dataset_path):data = []  # 存储所有图像的特征labels = []  # 存储每张图像对应的标签for fruit in fruit_labels:  # 遍历每个水果类别fruit_folder = os.path.join(dataset_path, fruit)  # 获取每个类别的文件夹路径for image_name in os.listdir(fruit_folder):  # 遍历文件夹中的每张图片image_path = os.path.join(fruit_folder, image_name)  # 获取图像的完整路径features = extract_features(image_path)  # 提取图像的特征if features is not None:data.append(features)  # 添加特征到数据列表labels.append(fruit_dict[fruit])  # 添加对应的标签（数字标签）到标签列表return np.array(data), np.array(labels)  # 返回特征和标签的 numpy 数组# 加载数据并进行训练集和测试集划分
data, labels = load_data(dataset_path)  # 加载数据和标签
X_train, X_test, y_train, y_test = train_test_split(data, labels, test_size=0.2, random_state=42)  # 按 80% 训练、20% 测试划分数据集# 使用 KNN 分类器进行训练
knn = KNeighborsClassifier(n_neighbors=3)  # 初始化 KNN 分类器，选择 k=3
knn.fit(X_train, y_train)  # 使用训练集进行模型训练# 模型评估：预测测试集，并输出评估结果
y_pred = knn.predict(X_test)  # 对测试集进行预测
print(f"Accuracy: {accuracy_score(y_test, y_pred)}")  # 输出模型准确度
print(classification_report(y_test, y_pred, target_names=fruit_labels))  # 输出分类报告，包括精度、召回率、F1 分数等# 现在进行对传入图片的预测
def predict_image(image_path):# 读取输入图片img = cv2.imread(image_path)if img is not None:# 提取输入图像的特征img_resized = cv2.resize(img, (100, 100))  # 将图像调整为统一的大小img_flattened = img_resized.flatten()  # 将图像展平为一维特征向量# 使用训练好的 KNN 模型进行预测prediction = knn.predict([img_flattened])  # 注意这里是 [img_flattened]，KNN 需要二维数组# 输出预测结果predicted_label = fruit_labels[prediction[0]]  # 将预测的数字标签转换为水果标签print(f"Predicted label: {predicted_label}")  # 打印预测标签# 可视化输入图片并显示预测结果plt.imshow(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))  # 将图像从 BGR 转为 RGB 并显示plt.title(f"Predicted: {predicted_label}")  # 显示标题（预测结果）plt.show()  # 显示图像else:print(f"Error: Image at {image_path} could not be read.")  # 如果无法读取图片，输出错误信息# 调用 predict_image 函数来预测一个新的图片
# 传入你希望预测的图片路径
image_path = './fruits/pear/1.jpg'  # 替换为你的图像路径，确保图像路径正确
predict_image(image_path)  # 进行预测并显示结果

运行结果：