写python代码,读取文件夹中图片,用神经网络对图片进行特征提取,对提取到的特征进行k-means聚类,显示聚类结果图,并将聚类后的结果保存到新的文件夹中
时间: 2024-05-12 19:13:02 浏览: 84
以下是一个示例代码,需要根据具体情况进行修改:
```python
import os
import cv2
import numpy as np
from sklearn.cluster import KMeans
# 读取文件夹中的图片,获取特征向量
def get_features(folder_path):
features = []
for file_name in os.listdir(folder_path):
if file_name.endswith(".jpg"):
file_path = os.path.join(folder_path, file_name)
img = cv2.imread(file_path)
img = cv2.resize(img, (224, 224)) # 调整图片大小
feature = extract_feature(img) # 提取特征
features.append(feature)
return np.array(features)
# 提取图片特征,这里使用了一个预训练的神经网络模型
def extract_feature(img):
# TODO: 替换为你自己的神经网络模型
model = None
feature = model.predict(img)
return feature.flatten()
# 对特征向量进行k-means聚类
def kmeans_cluster(features, n_clusters):
kmeans = KMeans(n_clusters=n_clusters)
kmeans.fit(features)
return kmeans.labels_
# 显示聚类结果图
def show_cluster_result(folder_path, labels):
for i, file_name in enumerate(os.listdir(folder_path)):
if file_name.endswith(".jpg"):
file_path = os.path.join(folder_path, file_name)
img = cv2.imread(file_path)
img = cv2.resize(img, (224, 224)) # 调整图片大小
label = labels[i]
cv2.putText(img, str(label), (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1.0, (0, 255, 0), 2)
cv2.imshow("Cluster Result", img)
cv2.waitKey(0)
# 将聚类后的结果保存到新的文件夹中
def save_cluster_result(folder_path, labels):
for i, file_name in enumerate(os.listdir(folder_path)):
if file_name.endswith(".jpg"):
file_path = os.path.join(folder_path, file_name)
label = labels[i]
new_folder_path = os.path.join(os.path.dirname(folder_path), "cluster_" + str(label))
if not os.path.exists(new_folder_path):
os.makedirs(new_folder_path)
new_file_path = os.path.join(new_folder_path, file_name)
os.rename(file_path, new_file_path)
# 主函数
if __name__ == "__main__":
folder_path = "path/to/folder" # 文件夹路径
n_clusters = 5 # 聚类数
features = get_features(folder_path)
labels = kmeans_cluster(features, n_clusters)
show_cluster_result(folder_path, labels)
save_cluster_result(folder_path, labels)
```
其中 `extract_feature` 函数需要替换为你自己的神经网络模型,可以使用预训练的模型或者自己训练的模型。另外,聚类结果可能会因为随机初始化而有所不同。
阅读全文