import visdom
import time
import datetime
import os
import torch as t
import numpy as np
import matplotlib.pyplot as plt
from util.Data import Data
import torch
import models
from util.utils import extract_feature_test
from util.utils import sort_img,imshow
from util.data.market1501_1 import Market1501
from train_trihard import args
query_index = 777
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
Baseline = models.init_model(name='resnet50', num_classes=751, loss={'softmax', 'metric'})
# Adf.load_state_dict(t.load('F:/My_ReID/experiment_data/0401/weights/model_444.pt'))
checkpoint = torch.load('E:/gyx/reid_tutorial/log/checkpoint_ep200.pth.tar')
Baseline.load_state_dict(checkpoint['state_dict'])
market_train = Market1501(dtype='train',root='E:/gyx/')
market_test = Market1501(dtype='test',root='E:/gyx/')
market_query = Market1501(dtype='query',root='E:/gyx/')
# trainset,num_train_pids,num_train_imgs = market_train._process_dir()
testset,num_test_pids,test_camid,num_test_imgs = market_test._process_dir( dir_path='E:/gyx/market1501/bounding_box_test/', relabel=False)
queryset,num_query_pids,query_camid,num_query_imgs = market_query._process_dir( dir_path='E:/gyx/market1501/query/', relabel=False)
data = Data(args)
test_loader = data.test_loader
query_loader = data.query_loader
# print('queryset',queryset)
# print('testset',testset)
Baseline.eval()
dataset_query,query_label,query_camid,query_path = queryset,num_query_pids,query_camid,num_query_imgs
dataset_test,gallery_label,test_camid,gallery_path = testset,num_test_pids,test_camid,num_test_imgs
# print('query_label',query_label)
# print('gallery_label',gallery_label)
# print('query_path',query_path)
# print('gallery_path',gallery_path)
# Extract feature
query_feature = extract_feature_test(Baseline, query_loader)
gallery_feature = extract_feature_test(Baseline, test_loader)
query_feature = query_feature.cuda()
gallery_feature = gallery_feature.cuda()
i = query_index
index = sort_img(query_feature[i],query_label[i],query_camid[i],gallery_feature,gallery_label,test_camid)
query_path = query_path[i]
query_label = query_label[i]
print(query_path)
print('Top 10 images are as follow:')
try: # Visualize Ranking Result
# Graphical User Interface is needed
fig = plt.figure(figsize=(16,4))
ax = plt.subplot(1,11,1)
ax.axis('off')
imshow(query_path,'query')
for i in range(10):
ax = plt.subplot(1,11,i+2)
ax.axis('off')
img_path = gallery_path[index[i]]
label = gallery_label[index[i]]
imshow(img_path)
if label == query_label:
ax.set_title('%d'%(i+1), color='green')
else:
ax.set_title('%d'%(i+1), color='red')
print(img_path)
except RuntimeError:
for i in range(10):
# log_path = "./show" + '/Log %d.txt' % query_index
# if not os.path.exists(log_path):
# os.system(r"touch {}".format(log_path))
img_path = gallery_path.imgs[index[i]]
print(img_path[0])
# f = open(log_path, 'a')
# f.write(img_path + '\n')
print('If you want to see the visualization of the ranking result, graphical user interface is needed.')
fig.savefig("./show/show %d.png"% query_index)
from importlib import import_module
from torchvision import transforms
from torch.utils.data import dataloader
from util import data_manager
from util import transforms as T
from util.samplers import RandomIdentitySampler
from torch.utils.data import DataLoader
from util.dataset_loader import ImageDataset
class Data:
def __init__(self, args):
dataset = data_manager.init_img_dataset(root=args.root, name=args.dataset, split_id=args.split_id,
cuhk03_labeled=args.cuhk03_labeled,
cuhk03_classic_split=args.cuhk03_classic_split, )
self.dataset = dataset
train_list = T.Compose([
T.Random2DTranslation(args.height, args.width),
T.RandomHorizontalFlip(),
# T.RandomErasing(),
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),])
test_list = T.Compose([
T.Resize((args.height, args.width)),
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),])
# module = import_module('data.' + args.data_train.lower()+'_1')
self.train_loader = DataLoader(
ImageDataset(dataset.train, transform=train_list),
sampler=RandomIdentitySampler(dataset.train, num_instances=args.num_instances),
batch_size=args.train_batch, num_workers=args.workers, pin_memory=True, drop_last=True, )
self.test_loader = DataLoader(
ImageDataset(dataset.gallery, transform=test_list),
batch_size=args.test_batch, shuffle=False, num_workers=args.workers,
pin_memory=True, drop_last=False,
)
self.query_loader = DataLoader(
ImageDataset(dataset.query, transform=test_list),
batch_size=args.test_batch, shuffle=False, num_workers=args.workers,
pin_memory=True, drop_last=False,
)
from __future__ import absolute_import
import os
import sys
import errno
import shutil
import json
import os.path as osp
from PIL import Image
import matplotlib.pyplot as plt
import cv2
import numpy as np
from numpy import array,argmin
from torch.autograd import Variable
import torch
import torch.nn as nn
#####################################################################
#Show result
def imshow(path, title=None):
"""Imshow for Tensor."""
im = plt.imread(path)
plt.imshow(im)
if title is not None:
plt.title(title)
plt.pause(0.001) # pause a bit so that plots are updated
######################################################################
# Extract feature
# ----------------------
# Extract feature from a trained model.
#'--ms',default='1', type=str,help='multiple_scale: e.g. 1 1,1.1 1,1.1,1.2'
ms = 1, 1
PCB = False
def fliplr(img):
'''flip horizontal'''
inv_idx = torch.arange(img.size(3)-1,-1,-1).long() # N x C x H x W
inv_idx= inv_idx.cuda()
img_flip = img.index_select(3,inv_idx)
return img_flip
def extract_feature(model,dataloaders):
features = torch.FloatTensor()
count = 0
for data in dataloaders:
img, label = data
# print('img',img)
# print('label',label)
n, c, h, w = img.size()
count += n
# print(count)
ff = torch.FloatTensor(n, 512).zero_().cuda()
# ff = torch.FloatTensor(n,2048).zero_().cuda()
if PCB:
ff = torch.FloatTensor(n,2048,6).zero_().cuda() # we have six parts
for i in range(2):
if(i==1):
img = fliplr(img)
input_img = Variable(img.cuda())
for scale in ms:
if scale != 1:
# bicubic is only available in pytorch>= 1.1
input_img = nn.functional.interpolate(input_img, scale_factor=scale, mode='bicubic', align_corners=False)
# print('input_img.shape',input_img[0].shape) #torch.Size([1, 128, 64, 3])
# input_img = input_img[0].cpu().numpy()
# input_img = np.transpose(input_img, (2, 0, 1))
# input_img = torch.FloatTensor(input_img)
outputs = model(input_img)
# print('outputs.shape',outputs.shape) torch.Size([1, 2048])
# print('ff.shape', ff.shape)
ff += outputs
# norm feature
if PCB:
# feature size (n,2048,6)
# 1. To treat every part equally, I calculate the norm for every 2048-dim part feature.
# 2. To keep the cosine score==1, sqrt(6) is added to norm the whole feature (2048*6).
fnorm = torch.norm(ff, p=2, dim=1, keepdim=True) * np.sqrt(6)
ff = ff.div(fnorm.expand_as(ff))
ff = ff.view(ff.size(0), -1)
else:
fnorm = torch.norm(ff, p=2, dim=1, keepdim=True)
ff = ff.div(fnorm.expand_as(ff))
features = torch.cat((features,ff.data.cpu()), 0)
return features
def extract_feature_test(model,dataloaders):
features = torch.FloatTensor()
count = 0
# print(len(dataloaders)) #995
for data in dataloaders:
img, label, camids = data
# print('img',img)
# print('label',label)
n, c, h, w = img.size()
count += n
# print(count)
# ff = torch.FloatTensor(n,2048).zero_().cuda()
ff = torch.FloatTensor(n, 2048).zero_().cuda()
# ff = torch.FloatTensor(n, 512).zero_().cuda()
if PCB:
ff = torch.FloatTensor(n,2048,6).zero_().cuda() # we have six parts
for i in range(2):
if(i==1):
img = torch.FloatTensor(img).cuda()
img = fliplr(img)
# img = Variable(img.cuda())
# input_img = Variable(img)
# input_img = img.type(torch.FloatTensor)
input_img = img.cuda()
for scale in ms:
if scale != 1:
# bicubic is only available in pytorch>= 1.1
input_img = nn.functional.interpolate(input_img, scale_factor=scale, mode='bicubic', align_corners=False)
# print('input_img.shape',input_img[0].shape) #torch.Size([1, 128, 64, 3])
# input_img = input_img[0].cpu().numpy()
# input_img = np.transpose(input_img, (2, 0, 1))
# input_img = torch.FloatTensor(input_img)
# print(input_img.shape)#torch.Size([16, 3, 256, 128])
model.cuda()
outputs = model(input_img)
# print('outputs.shape',outputs.shape) #torch.Size([1, 2048])
# print('ff.shape', ff.shape)
ff += outputs
# norm feature
if PCB:
# feature size (n,2048,6)
# 1. To treat every part equally, I calculate the norm for every 2048-dim part feature.
# 2. To keep the cosine score==1, sqrt(6) is added to norm the whole feature (2048*6).
fnorm = torch.norm(ff, p=2, dim=1, keepdim=True) * np.sqrt(6)
ff = ff.div(fnorm.expand_as(ff))
ff = ff.view(ff.size(0), -1)
else:
fnorm = torch.norm(ff, p=2, dim=1, keepdim=True)
ff = ff.div(fnorm.expand_as(ff))
features = torch.cat((features,ff.data.cpu()), 0)
return features
#######################################################################
# sort the images
def sort_img(qf, ql, qc, gf, gl, gc):
query = qf.view(-1, 1)
# print(query.shape)
score = torch.mm(gf, query)
score = score.squeeze(1).cpu()
score = score.numpy()
# predict index
index = np.argsort(score) # from small to large
index = index[::-1]
# index = index[0:2000]
# good index
query_index = np.argwhere(gl == ql)
# same camera
camera_index = np.argwhere(gc == qc)
# good_index = np.setdiff1d(query_index, camera_index, assume_unique=True)
junk_index1 = np.argwhere(gl == -1)
junk_index2 = np.intersect1d(query_index, camera_index)
junk_index = np.append(junk_index2, junk_index1)
mask = np.in1d(index, junk_index, invert=True)
index = index[mask]
return index
"""Image ReID"""
import os.path as osp
import glob
import re
class Market1501(object):
"""
Market1501
Reference:
Zheng et al. Scalable Person Re-identification: A Benchmark. ICCV 2015.
URL: http://www.liangzheng.org/Project/project_reid.html
Dataset statistics:
# identities: 1501 (+1 for background)
# images: 12936 (train) + 3368 (query) + 15913 (gallery)
"""
dataset_dir = 'market1501'
def __init__(self,dtype,root='E:/gyx/',**kwargs):
self.dataset_dir = osp.join(root, self.dataset_dir)
if dtype == 'train':
self.dir_path = osp.join(self.dataset_dir, 'bounding_box_train')
self.relabel=True
elif dtype == 'test':
self.dir_path = osp.join(self.dataset_dir, 'bounding_box_test')
self.relabel=False
else:
self.dir_path = osp.join(self.dataset_dir, 'query')
self.relabel=False
def _process_dir(self, dir_path, relabel=False):
img_paths = glob.glob(osp.join(dir_path, '*.jpg'))
pattern = re.compile(r'([-\d]+)_c(\d)')
# pattern = re.compile(r'([-\d]+)_c(\d)s(\d)_([-\d]+)_([-\d]+)')
pid_container = set()
for img_path in img_paths:
pid, camid = map(int, pattern.search(img_path).groups())
if pid == -1: continue # junk images are just ignored
pid_container.add(pid)
pid2label = {pid: label for label, pid in enumerate(pid_container)}
dataset = []
pid_list = []
cam_list = []
img_list= []
for img_path in img_paths:
pid, camid = map(int, pattern.search(img_path).groups())
if pid == -1: continue # junk images are just ignored
assert 0 <= pid <= 1501 # pid == 0 means background
assert 1 <= camid <= 6
camid -= 1 # index starts from 0
if relabel: pid = pid2label[pid]
dataset.append((img_path, pid, camid))
pid_list.append(pid)
cam_list.append(camid)
img_list.append(img_path)
# num_pids = len(pid_container)
# num_imgs = len(dataset)
return dataset, pid_list,cam_list, img_list
from __future__ import absolute_import
from .ResNet import *
__factory = {
'resnet50': ResNet50,
}
def get_names():
return __factory.keys()
def init_model(name, *args, **kwargs):
if name not in __factory.keys():
raise KeyError("Unknown model: {}".format(name))
return __factory[name](*args, **kwargs)
from __future__ import absolute_import
import torch
import torchvision
from torch import nn
from torch.nn import functional as F
from IPython import embed
import torch.nn as nn
class ResNet50(nn.Module):
def __init__(self, num_classes, loss={'softmax', 'metric'}, aligned=False, **kwargs):
super(ResNet50, self).__init__()
self.loss = loss
resnet50 = torchvision.models.resnet50(pretrained=True)
# 可以用来提取出model每一层的网络结构(去除后两层):
self.base = nn.Sequential(*list(resnet50.children())[:-2])
self.classifier = nn.Linear(2048, num_classes)
self.feat_dim = 2048 # feature dimension
self.aligned = aligned
if self.aligned:
self.bn = nn.BatchNorm2d(2048)
self.relu = nn.ReLU(inplace=True)
self.conv1 = nn.Conv2d(2048, 128, kernel_size=1, stride=1, padding=0, bias=True)
def forward(self, x):
# print('x.shape',x.shape)# x.shape torch.Size([16, 3, 256, 128])
x = self.base(x)
if not self.training:
lf = self.horizon_pool(x)
if self.aligned and self.training:
lf = self.bn(x)
lf = self.relu(lf)
lf = self.horizon_pool(lf)
lf = self.conv1(lf)
if self.aligned or not self.training:
# print('lf.shape', lf.shape) #lf.shape torch.Size([16, 128, 8, 1])
# print('lf.size()[0:3]',lf.size()[0:3]) #lf.size()[0:3] torch.Size([16, 128, 8])
# print('lf.size',lf.size) #lf.size
lf = lf.view(lf.size()[0:3])
# print('lf.size', lf.size) #lf.size
# print('lf.shape', lf.shape) #lf.shape torch.Size([16, 128, 8])
lf = lf / torch.pow(lf,2).sum(dim=1, keepdim=True).clamp(min=1e-12).sqrt()
# print('lf.shape', lf.shape) #lf.shape torch.Size([16, 128, 8])
x = F.avg_pool2d(x, x.size()[2:])
f = x.view(x.size(0), -1)
#f = 1. * f / (torch.norm(f, 2, dim=-1, keepdim=True).expand_as(f) + 1e-12)
if not self.training:
return f
# return f, lf # if using local feature
y = self.classifier(f)
if self.loss == {'softmax'}:
return y
elif self.loss == {'metric'}:
if self.aligned: return f, lf
return f
elif self.loss == {'softmax', 'metric'}:
if self.aligned: return y, f, lf
return y, f
else:
raise KeyError("Unsupported loss: {}".format(self.loss))
if __name__ == '__main__':
model = ResNet50(num_classes=751, loss={'softmax', 'metric'},aligned=True)
imgs = torch.Tensor(32, 3, 256, 128)
y, f, lf = model(imgs)
embed()
query_index:777
query_index:3202
query_index:888
query_index:2