def __getitem__(self, index):
"""
Args:
index (int): Index
Returns:
image (PIL): Loaded image
"""
path = self.images_fn[index]
img = self.loader(path)
imfullsize = max(img.size)
if self.bbxs is not None:
img = img.crop(self.bbxs[index])
if self.imsize is not None:
if self.bbxs is not None:
img = imresize(img, self.imsize * max(img.size) / imfullsize)
else:
img = imresize(img, self.imsize)
if self.transform is not None:
img = self.transform(img)
return img, path
def __getitem__(self, index):
"""
Args:
index (int): Index
Returns:
images tuple (q,p,n1,...,nN): Loaded train/val tuple at index of self.qidxs
"""
if self.__len__() == 0:
raise (RuntimeError(
"List qidxs is empty. Run ``dataset.create_epoch_tuples(net)`` method to create subset for train/val!"
))
# query image
output = (self.loader(self.images[self.qpool[index]]))
if self.imsize is not None:
if self.random_size:
w, h = output.size
imsize = get_random_size(self.imsize, w, h)
output = imresize(output, imsize)
else:
output = imresize(output, self.imsize)
if self.transform is not None:
output = self.transform(output)
return output
def __getitem__(self, index):
# output = self.loader(self.images[self.qidxs[index]])
output = self.loaded_imgs[index]
if self.imsize is not None:
if self.random:
w, h = output.size
imsize = get_random_size(self.imsize, w, h)
output = imresize(output, imsize)
else:
output = imresize(output, self.imsize)
if self.transform is not None:
output = self.transform(output)
return output
def __getitem__(self, index):
"""
Args:
index (int): Index
Returns:
images tuple (q,p,n1,...,nN): Loaded train/val tuple at index of self.qidxs
"""
if self.__len__() == 0:
raise (RuntimeError(
"List qidxs is empty. Run ``dataset.create_epoch_tuples(net)`` method to create subset for train/val!"
))
output = []
# query image
output.append(self.loader(self.images[self.qidxs[index]]))
# positive image
output.append(self.loader(self.images[self.pidxs[index]]))
# negative images
for i in range(len(self.nidxs[index])):
output.append(self.loader(self.images[self.nidxs[index][i]]))
if self.imsize is not None:
output = [imresize(img, self.imsize) for img in output]
if self.transform is not None:
output = [
self.transform(output[i]).unsqueeze_(0)
for i in range(len(output))
]
target = torch.Tensor([-1, 1] + [0] * len(self.nidxs[index]))
return output, target
def __getitem__(self, index):
"""
Args:
index (int): Index
Returns:
image (PIL): Loaded image
"""
path = self.images_fn[index]
img = self.loader(path)
if isinstance(img, Exception):
sys.stderr.write("Warning: Image '%s' was not found\n" % path)
if self.ignore_errors:
return {}
else:
raise img
if self.bbxs and self.bbxs[index]:
img = img.crop(self.bbxs[index])
if self.imsize is not None:
img = imresize(img, self.imsize)
if self.transform is not None:
img = self.transform(img)
return img
def img_loader(image_name, imsize, bbx=None):
img = Image.open(image_name)
imsize_ = np.max((img.height, img.width))
if bbx: img = img.crop(bbx)
img = imresize(
img, np.floor(1.0 * np.max(
(img.height, img.width)) * imsize / imsize_))
return Variable(TF.to_tensor(img)).unsqueeze(0)
def get_descriptor_from_image(net: ImageRetrievalNet, img: Image,
img_resolution: int, device) -> np.array:
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=net.meta['mean'], std=net.meta['std'])
])
imgr: torch.Tensor = transform(imresize(img, img_resolution))
return _multi_scale_extract(
net, imgr.unsqueeze(0),
[1, 1 / np.sqrt(2), np.sqrt(2)], device)
def inference(self, img):
try:
input = Image.open(img).convert("RGB")
input = imresize(input, 224)
input = transforms(input).unsqueeze()
with torch.no_grad():
vect = net(input)
return vect
except:
print('cannot indentify error')
def __getitem__(self, index):
"""
Args:
index (int): Index
Returns:
image (PIL): Loaded image
"""
path = self.images_fn[index]
img = self.loader(path)
if self.bbxs:
img = img.crop(self.bbxs[index])
if self.imsize is not None:
if self.random_size:
w, h = img.size
imsize = get_random_size(self.imsize, w, h)
img = imresize(img, imsize)
else:
img = imresize(img, self.imsize)
if self.transform is not None:
img = self.transform(img)
return img
def __init__(
self,
imsize=None,
nnum=5,
qsize=2000,
poolsize=20000,
transform=None,
loader=default_loader,
filename=None,
q_percent=1,
):
# setting up paths
data_root = get_data_root()
name = "retrieval-SfM-120k"
db_root = os.path.join(data_root, "train", name)
ims_root = os.path.join(db_root, "ims")
# loading db
db_fn = os.path.join(db_root, "{}.pkl".format(name))
with open(db_fn, "rb") as f:
db = pickle.load(f)["val"]
# initializing tuples dataset
self.imsize = imsize
self.images = [
cid2filename(db["cids"][i], ims_root)
for i in range(len(db["cids"]))
]
self.clusters = db["cluster"]
self.qpool = db["qidxs"]
# self.ppool = db['pidxs']
# size of training subset for an epoch
self.nnum = nnum
self.qsize = min(qsize, len(self.qpool))
self.poolsize = min(poolsize, len(self.images))
self.qidxs = self.qpool
self.index = np.arange(len(self.qidxs))
if q_percent < 1:
number = int(len(self.qidxs) * q_percent)
self.index = np.random.permutation(self.index)
self.index = self.index[:number]
self.pidxs = []
self.nidxs = []
self.poolvecs = None
self.transform = transform
self.loader = loader
self.filename = filename
self.phase = 1
self.ranks = torch.load(f"{filename}/ranks_362")
if os.path.isfile(f"{filename}/pool_vecs"):
self.pool_vecs = pickle.load(open(f"{filename}/pool_vecs", "rb"))
print(len(self.images))
self.loaded_images = []
if os.path.exists("./images"):
self.loaded_images = pickle.load(open("./images", "rb"))
else:
for i in range(len(self.images)):
try:
img = self.loader(self.images[i])
if self.imsize is not None:
img = imresize(img, self.imsize)
if self.transform is not None:
img_tensor = self.transform(img).unsqueeze(0)
img.close()
self.loaded_images.append(img_tensor)
except:
self.loaded_images.append(None)
pickle.dump(self.loaded_images, open("./images", "wb"))
def main():
# setting up the visible GPU
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
input_resol = 512
# resolution of input image, will resize to that if larger
# input_resol = 1024; # resolution of input image, will resize to that if larger
scales = [1, 1 / np.sqrt(2),
1 / 2] # re-scaling factors for multi-scale extraction
# sample image
img_file = 'sanjuan.jpg'
if not path.exists(img_file):
os.system(
'wget https://raw.githubusercontent.com/gtolias/tma/master/data/input/'
+ img_file)
img = default_loader(img_file)
print("use network trained with gem pooling and FC layer")
state = load_url(TRAINED['rSfM120k-tl-resnet101-gem-w'],
model_dir=os.path.join(get_data_root(), 'networks'))
net = init_network({
'architecture': state['meta']['architecture'],
'pooling': state['meta']['pooling'],
'whitening': state['meta'].get('whitening')
})
net.load_state_dict(state['state_dict'])
net.eval()
net.cuda()
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=state['meta']['mean'],
std=state['meta']['std'])
])
# single-scale extraction
vec = extract_ss(net,
transform(imresize(img, input_resol)).unsqueeze(0).cuda())
vec = vec.data.cpu().numpy()
print(vec)
# multi-scale extraction
vec = extract_ms(net,
transform(imresize(img, input_resol)).unsqueeze(0).cuda(),
ms=scales,
msp=1.0)
vec = vec.data.cpu().numpy()
print(vec)
print("\n")
print(
"use network trained with gem pooling, and apply the learned whitening transformation"
)
state = load_url(TRAINED['retrievalSfM120k-resnet101-gem'],
model_dir=os.path.join(get_data_root(), 'networks'))
net = init_network({
'architecture': state['meta']['architecture'],
'pooling': state['meta']['pooling'],
'whitening': state['meta'].get('whitening')
})
net.load_state_dict(state['state_dict'])
net.eval()
net.cuda()
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=state['meta']['mean'],
std=state['meta']['std'])
])
# single-scale extraction
vec = extract_ss(net,
transform(imresize(img, input_resol)).unsqueeze(0).cuda())
vec = vec.data.cpu().numpy()
print(vec)
whiten_ss = state['meta']['Lw']['retrieval-SfM-120k']['ss']
vec = whitenapply(vec.reshape(-1, 1), whiten_ss['m'],
whiten_ss['P']).reshape(-1)
print(vec)
# multi-scale extraction
vec = extract_ms(net,
transform(imresize(img, input_resol)).unsqueeze(0).cuda(),
ms=scales,
msp=net.pool.p.item())
vec = vec.data.cpu().numpy()
print(vec)
whiten_ms = state['meta']['Lw']['retrieval-SfM-120k']['ms']
vec = whitenapply(vec.reshape(-1, 1), whiten_ms['m'],
whiten_ms['P']).reshape(-1)
print(vec)
print("\n")
print("use pre-trained (on ImageNet) network with appended mac pooling")
net = init_network({
'architecture': 'resnet101',
'pooling': 'mac',
'pretrained': True
})
net.eval()
net.cuda()
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=net.meta['mean'], std=net.meta['std'])
])
# single-scale extraction
vec = extract_ss(net,
transform(imresize(img, input_resol)).unsqueeze(0).cuda())
vec = vec.data.cpu().numpy()
print(vec)
# multi-scale extraction
vec = extract_ms(net,
transform(imresize(img, input_resol)).unsqueeze(0).cuda(),
ms=scales,
msp=1.0)
vec = vec.data.cpu().numpy()
print(vec)
print("\n")
def loader(image_name, im_size):
return Variable(TF.to_tensor(imresize(Image.open(image_name),
im_size))).unsqueeze(0)
