# test.py: eval query image
import torch
from math import floor
import torch.nn.functional as F
from torchvision.transforms import ToTensor, Normalize, Resize, RandomHorizontalFlip, RandomRotation, Compose
import utils.transforms as trf
from PIL import Image
def extract_query(net, dataset, q_idx, scale=800, crop=True, flip=None, rotate=None):
# Load query image
img = Image.open(dataset.get_query_filename(q_idx))
# Crop the query ROI
if crop:
img = img.crop(tuple(dataset.get_query_roi(q_idx)))
# Apply transformations
img = trf.resize_image(img, scale)
# Flip
if flip: img = trf.flip_image(img, flip)
# Rotation
if rotate: img = trf.rotate(img, rotate)
# Convert to Pytorch's tensor and normalize
I = trf.to_tensor(img)
I = trf.normalize(I, dict(rgb_means=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]))
I = I.unsqueeze(0)#.to(device)
# Forward pass to extract the features
with torch.no_grad():
print ('Extracting features using input={} pixels...'.format(scale))
q_feat = net(I).cpu().numpy()
return q_feat, img
def q_eval(net, dataset, q_idx, flip=False, rotate=False, scale=1):
# load query image
q_im = dataset.get_image(q_idx)
q_size = q_im.size
# list of transformation lists
trfs_chains = [[]]
if rotate:
eps = 1e-6
trfs_chains[0] += [RandomRotation((rotate-eps,rotate+eps))]
if flip:
trfs_chains[0] += [RandomHorizontalFlip(1)]
if scale == 0: # AlexNet asks for resized images of 224x224
edge_list = [224]
resize_list = [Resize((edge,edge)) for edge in edge_list]
elif scale == 1:
edge_list = [800]
resize_list = [lambda im: imresize(im, edge) for edge in edge_list]
elif scale == 1.5:
edge_list = [1200]
resize_list = [lambda im: imresize(im, edge) for edge in edge_list]
elif scale == 2: # multiscale
edge_list = [600,800,1000,1200]
resize_list = [lambda im: imresize(im, edge) for edge in edge_list]
else:
raise ValueError()
if len(resize_list) == 1:
trfs_chains[0] += resize_list
else:
add_trf(trfs_chains, resize_list )
# default transformations
mean = [0.485, 0.456, 0.406]
std = [0.229, 0.224, 0.225]
for chain in trfs_chains:
chain += [ToTensor(), Normalize(mean, std)]
net = net.eval()
q_feat = torch.zeros( (len(trfs_chains), net.out_features) )
print ('Computing the forward pass and extracting the image representation...')
for i in range(len(trfs_chains)):
q_tensor = Compose(trfs_chains[i])(q_im)
import pdb; pdb.set_trace() # XXX BREAKPOINT
q_feat[i] = net.forward(q_tensor.view(1,q_tensor.shape[0],q_tensor.shape[1],q_tensor.shape[2]))
return F.normalize(q_feat.mean(dim=0), dim=0).detach().numpy()
def add_trf(trfs_chains, new_trfs):
n_chains = len(trfs_chains)
for trf in new_trfs:
for i in range(n_chains):
trfs_chains.append(trfs_chains[i] + [trf])
return
def imresize(im, maxedge):
''' creates image in a different size, where the aspect ratio as maintainted max max(height,width)=maxedge
'''
h,w = im.size
if h<w:
return im.resize((int(h/w*maxedge), maxedge))
else:
return im.resize((maxedge, int(h/w*maxedge)))
