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

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

news 来源：原创 2025/5/1 9:52:09

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

通过结合深度学习和无监督学习技术，我们可以自动对大量图片进行分类。本文将详细介绍如何使用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是一种常用的无监督学习算法，算法目标是通过迭代过程找到数据集的簇划分，使得每个簇内的样本与簇内均值的平方误差最小化。

算法基本步骤如下：

随机初始化K个簇中心。
将每个数据点分配到最近的簇中心。
更新簇中心位置为簇内数据点的均值。
重复步骤2和3，直到簇中心稳定。

具体理论知识可参照前文：【机器学习】九、聚类：（一）K-means聚类

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

二、实现步骤与代码详解

1. 导入库

import os
import shutil
import numpy as np
from PIL import Image
from sklearn.cluster import KMeans
from torchvision import models, transforms
import 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)))