基于K-Means的图片聚类算法实战

一. 场景说明

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


class 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 labels


if __name__ == "__main__":
    import shutil
    import os
    extract = 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 = 2
    labels = 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-均值）聚类算法