使用预训练的ResNet50提取图片的特征向量

最新推荐文章于 2024-08-20 15:14:09 发布

岳戴

最新推荐文章于 2024-08-20 15:14:09 发布

阅读量1w

点赞数 16

本文链接：https://blog.csdn.net/Jackydyy/article/details/119238657

版权

本文介绍了如何使用预训练的ResNet50模型来提取单张图片及整个文件夹中图片的特征向量，涵盖从单张图片到批量处理的流程。

摘要由CSDN通过智能技术生成

1.提取单张图片的特征向量

import os.path

import torch
import torch.nn as nn
from torchvision import models, transforms
from torch.autograd import Variable
import numpy as np
from PIL import Image

features_dir = './features'    # 存放特征的文件夹路径

img_path = "FullFrame1/images0001.png"     # 图片路径
file_name = img_path.split('/')[-1]        # 图片路径的最后一个/后面的名字
feature_path = os.path.join(features_dir, file_name + '.txt')    # /后面的名字加txt

transform1 = transforms.Compose([      # 串联多个图片变换的操作
    transforms.Scale(256),   # 缩放
    transforms.CenterCrop(224),  # 中心裁剪
    transforms.ToTensor()]     # 转换成Tensor
)

img = Image.open(img_path)   # 打开图片
img1 = transform1(img)       # 对图片进行transform1的各种操作

# resnet18 = models.resnet18(pretrained = True)
resnet50_feature_extractor = models.resnet50(pretrained=True)    # 导入ResNet50的预训练模型
resnet50_feature_extractor.fc = nn.Linear(2048, 2048)