一. 场景说明

我们通常遇到一个问题，当很多图片放在一个文件夹中，要把这些文件夹中的图片按规律分为几类。当图片比较少时，我们可以手动完成，但是当图片的数量是几千甚至几万时，手动挑选图片的工作量就太大了。
因此，需要一个算法对图片进行自动化分类并保存，也就是本文的《基于K-Means的图片聚类算法实战》。

二. 算法步骤

1. 基于深度学习模型提取图片特征；
2. 基于K-Means聚类算法对图片特征进行聚类；
3. 根据不同的类别，把相应的图片保存在不同的文件夹中。

下面是一个例子，./imgs中保存了很多猫和狗的图片，运行脚本后，就把猫和狗的图片分开，分别保存在./0 和 ./1 两个文件夹中。
代码如下：

import torch
import torchvision.models as models
import torchvision.transforms as transforms
from PIL import Image
import torch.nn.functional as F
import glob
import numpy as np
from tqdm import tqdm
from sklearn.cluster import KMeansclass FeatureExtract(object):def __init__(self):# 加载预训练的ResNet18模型self.resnet = models.resnet18(pretrained=True)# 移除最后一层全连接层self.resnet = torch.nn.Sequential(*list(self.resnet.children())[:-1])# 设置模型为评估模式self.resnet.eval()self.preprocess = transforms.Compose([transforms.Resize(224),transforms.ToTensor(),transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),])def feature_extract(self, image_path):# 加载和预处理图像image = Image.open(image_path)input_tensor = self.preprocess(image)input_batch = input_tensor.unsqueeze(0)# 使用模型提取特征with torch.no_grad():features = self.resnet(input_batch)# 输出特征向量return features.squeeze()def k_means(data, n_clusters=2):# 定义聚类器, 设置迭代次数和终止条件kmeans = KMeans(n_clusters=2, max_iter=1000, tol=1e-4)kmeans.fit(data)# 获取每个数据点的簇标签labels = kmeans.labels_return labelsif __name__ == "__main__":import shutilimport osextract = FeatureExtract()features = []images = glob.glob("./imgs/*.jpg")print("开始抽取图片特征...")for img_path in tqdm(images):feature = extract.feature_extract(img_path)features.append(feature.tolist())data = np.array(features)print("开始特征聚类...")n_clusters = 2labels = k_means(data, n_clusters)print("开始保存图片")for i in range(n_clusters):os.mkdir(str(i))for label, img_path in zip(labels, images):shutil.copy(img_path, os.path.join(str(label), os.path.basename(img_path)))

理论基础可学习：K-Means（K-均值）聚类算法