# -*- coding:utf-8 -*-
import os
import torch
import torch.nn as nn
from Resnet import resnet18
from torch.autograd import Variable
from torchvision import transforms
from PIL import Image
import math
def load_net():
net = resnet18(pretrained=True)
#channel_in = net.fc.in_features
#net.fc = nn.Linear(channel_in, num_classes)
#model_dict = torch.load(model_path).state_dict()
#net.load_state_dict(model_dict)
print(net)
net=net.to(device)
return net
def change_img_to_tensor(path):
img_size = transforms.Resize((224, 224))
img_to_tensor=transforms.ToTensor()
img_Normalize = transforms.Normalize(mean=[.5, .5, .5], std=[.5, .5, .5])
img=Image.open(path).convert('RGB')
img_sized=img_size(img)
tensor=img_to_tensor(img_sized)
tensor=img_Normalize(tensor)
return tensor
def get_feature(images_paths,net):
all_feature= []
paths = []
for fi in os.listdir(images_paths):
img_paths = images_paths+fi+'/'
for fj in os.listdir(img_paths):
img_path = img_paths+fj
img_tensor = change_img_to_tensor(img_path)
images = img_tensor.to(device)
input_img = torch.rand(1, 3,224,224)
input_img[0,:,:,:]=images
features,_ = net(Variable(input_img))
feature = features.data.cpu().numpy()
all_feature.append(feature)
paths.append(list(img_path))
return all_feature,paths
if __name__ == '__main__':
gpu_ids = [0]
#os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3"
device=torch.device("cuda:0" if torch.cuda.is_available() else 'cpu')
model_path = "./resnet18-5c106cde.pth"
net = load_net()
if len(gpu_ids) > 1:
net = nn.DataParallel(net, device_ids=gpu_ids)
net = net.eval()
#net = net.cuda()
images_paths = 'D:/02lunwen/01/paper_code/WANG/'
all_feature,paths = get_feature(images_paths,net)
