get_feature.py

# -*- 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)