甘肃建投土木工程建设有限公司网站后台,江西建设银行分行网站,河北网站建设费用,开放平台api一. 场景说明
我们通常遇到一个问题#xff0c;当很多图片放在一个文件夹中#xff0c;要把这些文件夹中的图片按规律分为几类。当图片比较少时#xff0c;我们可以手动完成#xff0c;但是当图片的数量是几千甚至几万时#xff0c;手动挑选图片的工作量就太大了。 因此当很多图片放在一个文件夹中要把这些文件夹中的图片按规律分为几类。当图片比较少时我们可以手动完成但是当图片的数量是几千甚至几万时手动挑选图片的工作量就太大了。 因此需要一个算法对图片进行自动化分类并保存也就是本文的《基于K-Means的图片聚类算法实战》。
二. 算法步骤
基于深度学习模型提取图片特征基于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 KMeansclass FeatureExtract(object):def __init__(self):# 加载预训练的ResNet18模型self.resnet models.resnet18(pretrainedTrue)# 移除最后一层全连接层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_clusters2):# 定义聚类器, 设置迭代次数和终止条件kmeans KMeans(n_clusters2, max_iter1000, tol1e-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-MeansK-均值聚类算法
