一. 场景说明
我们通常遇到一个问题,当很多图片放在一个文件夹中,要把这些文件夹中的图片按规律分为几类。当图片比较少时,我们可以手动完成,但是当图片的数量是几千甚至几万时,手动挑选图片的工作量就太大了。
因此,需要一个算法对图片进行自动化分类并保存,也就是本文的《基于K-Means的图片聚类算法实战》。
二. 算法步骤
下面是一个例子,./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)))