import os
import numpy as np
from sklearn.cluster import KMeans
import cv2
from imutils import build_montages
import torch.nn as nn
import torchvision.models as models
from PIL import Image
from torchvision import transforms
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
resnet50 = models.resnet50(pretrained=True)
self.resnet = nn.Sequential(resnet50.conv1,
resnet50.bn1,
resnet50.relu,
resnet50.maxpool,
resnet50.layer1,
resnet50.layer2,
resnet50.layer3,
resnet50.layer4)
def forward(self, x):
x = self.resnet(x)
return x
net = Net().eval()
image_path = []
all_images = []
images = os.listdir('./images_lianghua_v3')
for image_name in images:
image_path.append('./images_lianghua_v3/' + image_name)
for path in image_path:
image = Image.open(path).convert('RGB')
image = transforms.Resize([224,224])(image)
image = transforms.ToTensor()(image)
image = image.unsqueeze(0)
image = net(image)
print(image.shape)
image = image.reshape(-1, )
all_images.append(image.detach().numpy())
clt = KMeans(n_clusters=15)
clt.fit(all_images)
labelIDs = np.unique(clt.labels_)
print(labelIDs)
all_labels = 0
for labelID in labelIDs:
idxs = np.where(clt.labels_ == labelID)[0]
all_labels += len(idxs)
print(len(idxs))
idxs = np.random.choice(idxs, size=min(25, len(idxs)),replace=False)
show_box = []
for i in idxs:
image = cv2.imread(image_path[i])
image = cv2.resize(image, (96, 96))
show_box.append(image)
montage = build_montages(show_box, (96, 96), (5, 5))[0]
title = "Type{}".format(labelID)
cv2.imwrite(title + ".jpg",montage)
# cv2.imshow(title, montage)
# cv2.waitKey(0)
print("所有图片:",all_labels)
特征提取+聚类
最新推荐文章于 2024-03-04 11:14:45 发布
该代码示例首先导入必要的库,包括PyTorch和OpenCV等。它定义了一个基于ResNet50的网络结构,然后加载一批图像,对图像进行预处理,通过ResNet提取特征。接着,利用KMeans算法对这些特征进行聚类,最后展示每个类别的样本图像。
摘要由CSDN通过智能技术生成