加载模型
import torch
model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet18', pretrained=True)
# or any of these variants
# model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet34', pretrained=True)
# model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet50', pretrained=True)
# model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet101', pretrained=True)
# model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet152', pretrained=True)
model.eval()
可供选择的模型有很多,并且不断有新的模型加入,可以参考torchvision官方文档。分类常用的如:[AlexNet, ConvNeXt, DenseNet, EfficientNet, EfficientNetV2, GoogLeNet, Inception V3, MNASNet, MobileNet V2, MobileNet V3, RegNet, ResNet, ResNeXt, ShuffleNet V2, SqueezeNet, SwinTransformer, VGG, VisionTransformer, Wide ResNet]。这里以resnet为例。
处理数据
# sample execution (requires torchvision)
from PIL import Image
from torchvision import transforms
import torchvision
preprocess = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
path_to_data = '/path/to/data/'
BATCH_SIZE = 256
test_data = torchvision.datasets.ImageFolder(path_to_data, preprocess)
data_loader = torch.utils.data.DataLoader(test_data, batch_size=BATCH_SIZE)
该模型的输入需要设置 batchsize(N),3 通道的 RGB 图像(3 x H x W),H 和 W 最小为224,图像归一化到 [0, 1] 上,均值 mean = [0.485, 0.456, 0.406] 方差 std = [0.229, 0.224, 0.225]。本例中使用 torchvision 的 dataloader 组织数据。由于我们测试时没有标签,图像文件保存在 /path/to/data/ 下的一个文件夹内即可:
/path/to/data/
└── class_x
├── xxx.jpg
├── xxy.jpg
└── ...
测试
# load filename of test images
image_names = test_data.samples
# move the input and model to GPU for speed if available
model.to('cuda')
count = 0
result = [] # save the filename if it is predicted to be class 703 in imagenet labels.
with torch.no_grad():
for (x, y) in tqdm(data_loader, desc="Evaluating", leave=False):
x = x.to('cuda')
y = y.to('cuda')
output = model(x)
# Tensor of shape 1000, with confidence scores over Imagenet's 1000 classes
# print(output[0])
for index in range(output.shape[0]):
if torch.argmax(output[index,...])==703:
result.append(image_names[count*BATCH_SIZE+index][0])
count += 1
# The output has unnormalized scores. To get probabilities, you can run a softmax on it.
# print(torch.nn.functional.softmax(output[-1], dim=0))
model 的输出是图像的 logits,1000 维向量,通过查找其中的最大值可以得到其类别,imagenet 对应的类别标注参考 imagenet 或者IMAGENET 1000 Class List。要想得到 [0,1] 归一化后的概率表示,需要再进行一次 softmax。
全部代码如下
我们可以将需要查找的类别图像保存到一个目录(targetdir)下,也可以输出所有图像的路径。
import torch
#model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet18', pretrained=True)
# or any of these variants
# model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet34', pretrained=True)
# model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet50', pretrained=True)
model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet101', pretrained=True)
# model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet152', pretrained=True)
model.eval()
# sample execution (requires torchvision)
from PIL import Image
from torchvision import transforms
import torchvision
from tqdm import tqdm
preprocess = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
path_to_data = '/home/llj/ld/data/OpenEdit/train_latest/images/input_image/'
BATCH_SIZE = 256
nThreads = 4
target_dir = 'targetdir'
test_data = torchvision.datasets.ImageFolder(path_to_data, preprocess)
image_names = test_data.samples
data_loader = torch.utils.data.DataLoader(test_data, batch_size=BATCH_SIZE,)
model.to('cuda')
probability = []
count = 0
result = []
with torch.no_grad():
for (x, y) in tqdm(data_loader, desc="Evaluating", leave=False):
x = x.to('cuda')
y = y.to('cuda')
output = model(x)
for index in range(output.shape[0]):
if output[index,703]>0.5:
result.append(image_names[count*BATCH_SIZE+index][0])
count += 1
# Tensor of shape 1000, with confidence scores over Imagenet's 1000 classes
# print(output[-1])
# The output has unnormalized scores. To get probabilities, you can run a softmax on it.
# print(torch.nn.functional.softmax(output[-1], dim=0))
from shutil import copyfile
import os
os.makedirs(target_dir, exist_ok=True)
for r in result:
copyfile(r, target_dir+'/'+r.split('/')[-1])
for r in result:
print(r)
扩展应用
想通过缩略图找原图?汇报时使用压缩过的图像现在想找原图?如何从大量图像文件中快速找到与目标图像相似的那个?使用pytorch中的resnet预训练模型进行特征提取,以及查找相似图像
参考:
ResNet: Deep residual networks pre-trained on ImageNet
PyTorch Image Classification