【机器学习】使用VGG16与K-Means对大量图片进行自动分类

在数字时代，图片已经成为我们生活中不可或缺的一部分。从社交媒体上的自拍到电商平台的商品展示，再到个人电脑中存储的旅行照片，我们每天都在生成和接触大量的图片数据。然而，随着图片数量的增加，如何高效地组织和分类这些图片成为了一个令人头疼的问题。手动分类方法不仅耗时费力，还容易出错，尤其是在面对成千上万张图片时几乎不可行。

通过结合深度学习和无监督学习技术，我们可以自动对大量图片进行分类。本文将详细介绍如何使用VGG16模型（卷积神经网络）和K-Means聚类算法来实现这一目标。无论你是技术爱好者还是希望整理自己图片收藏的用户，这篇文章都将为你提供实用的指导和灵感。

下面将从背景知识开始，逐步深入到代码实现，最后探讨实际应用场景和优化建议。

一、背景知识

在正式开始代码实践之前，我们先了解两个核心技术：VGG16模型和K-Means聚类算法。

1、什么是VGG16？

VGG16是由牛津大学Visual Geometry Group（VGG）提出的一种深度卷积神经网络（CNN）模型。它在2014年的ImageNet大规模视觉识别挑战赛（ILSVRC）中表现出色，因其简单而有效的架构而广受欢迎。VGG16的名字来源于其包含16个带权重的层（13个卷积层和3个全连接层）。

VGG16结构简介：

输入层：224×224像素的RGB图像。

卷积块：包含13个卷积层和5个最大池化层，用于提取图像特征。

全连接层：3个全连接层，最后一层通常用于分类（在ImageNet上为1000类）。

在本项目中，我们利用VGG16前部的卷积层来提取图像特征，而不是直接进行分类。

2、什么是K-Means聚类？

K-Means是一种常用的无监督学习算法，算法目标是通过迭代过程找到数据集的簇划分，使得每个簇内的样本与簇内均值的平方误差最小化。

算法基本步骤如下：

1）随机初始化K个簇中心。

2）将每个数据点分配到最近的簇中心。

3）更新簇中心位置为簇内数据点的均值。

4）重复步骤2和3，直到簇中心稳定。

具体理论知识可参照前文：【机器学习】九、聚类：（一）K-means聚（https://www.guyuehome.com/detail?id=1825504051827896322）

下面将使用VGG16提取的图像特征作为K-Means算法的输入，从而将图片根据相似性自动分组。

二、实现步骤与代码详解

1. 导入库

import osimport shutilimport numpy as npfrom PIL import Imagefrom sklearn.cluster import KMeansfrom torchvision import models, transformsimport torch

2. 设置路径

image_dir = 'image_10k'target_dir = 'clustered_images_10k_100'

3. 加载VGG16模型

model = models.vgg16(pretrained=False)model.load_state_dict(torch.load('./vgg16-397923af.pth'))model.eval()

4. 图像预处理

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

5. 提取图片特征

def extract_features(image_path):    image = Image.open(image_path).convert('RGB')    image = preprocess(image).unsqueeze(0)    with torch.no_grad():        features = model(image)    return features.numpy().flatten()

6. 收集图片路径

image_paths = [os.path.join(image_dir, f) for f in os.listdir(image_dir)               if f.lower().endswith(('.jpg', '.jpeg', '.png'))]

7. 提取所有图片的特征

features = []for path in image_paths:    try:        features.append(extract_features(path))    except Exception as e:        print(f"Error processing 


    
{path}: {e}")
features = np.array(features)

8. K-Means聚类

kmeans = KMeans(n_clusters=100, random_state=42)labels = kmeans.fit_predict(features)

9. 创建分类文件夹并复制图片

os.makedirs(target_dir, exist_ok=True)for i in range(100):    os.makedirs(os.path.join(target_dir, f'cluster_{i}'), exist_ok=True)
for path, label in zip(image_paths, labels):    shutil.copy(path, os.path.join(target_dir, f'cluster_{label}', os.path.basename(path)))

三、关键技术深度解析

1、实际应用场景

图像检索：快速检索相似图片。

照片整理：按场景自动分组个人照片。

数据预处理：为监督学习任务提供基础分类！！！

2、性能优化建议

GPU加速：

device = torch.device("cuda"if torch.cuda.is_available() else"cpu")model = model.to(device)

并行处理图片：

from multiprocessing import Pool
with Pool() as p:    features = p.map(extract_features, image_paths)

批量处理图片：

batch_images = torch.stack([preprocess(Image.open(p).convert('RGB')) for p in image_paths[:32]])features = model(batch_images).detach().numpy()

3、聚类算法选型对比

from sklearn.metrics import silhouette_scorefrom hdbscan import HDBSCAN
# 不同算法效果对比algorithms = {    "K-Means": KMeans(n_clusters=100),    "DBSCAN": DBSCAN(eps=0.5, min_samples=5),    "HDBSCAN": HDBSCAN(min_cluster_size=50)}
for name, algo in algorithms.items():    labels = algo.fit_predict(features)    if hasattr(algo, 'labels_'):        labels = algo.labels_    score = silhouette_score(features, labels) if len(set(labels))>1else-1    print(f"{name}轮廓系数：{score:.2f}")

4、常见问题

图片格式处理：代码已支持.jpg、.jpeg和.png格式。对于其他格式，可以扩展if条件或使用PIL转换为支持的格式。

选择簇数量：使用肘部法则绘制聚类效果曲线选择。

特征提取速度：可以尝试更轻量模型（如MobileNet）或降低图片分辨率。

在这篇文章中，我们详细探讨了如何使用VGG16模型和K-Means聚类算法对大量图片进行自动分类。从背景知识到代码实现，再到实际应用和优化建议，我们提供了一个完整的解决方案。通过这种方法，你可以轻松地将杂乱无章的图片整理成有序的类别，无论是用于个人收藏还是商业项目~

点击“阅读原文”查看详情