def __init__(self, num_classes):
super(ResNet_VLAD, self).__init__()
modelPreTrain = resnet.resnet18(pretrained=True)
self.model = modelPreTrain
modelPreTrain2 = resnet.resnet18(pretrained=True)
self.model2 = modelPreTrain2
self.num_clusters = 32
self.conv1 = nn.Conv2d(8, 64, kernel_size=7, stride=1, padding=3, bias=False)
self.conv1.weight.data[:,0:3,:,:] = self.model.conv1.weight.data
self.conv1.weight.data[:,3:6,:,:] = self.model.conv1.weight.data
self.conv1.weight.data[:,6:8,:,:] = self.model.conv1.weight.data[:,0:2,:,:]
self.conv2 = nn.Conv2d(10, 64, kernel_size=7, stride=1, padding=3, bias=False)
self.conv2.weight.data[:,0:3,:,:] = self.model.conv1.weight.data
self.conv2.weight.data[:,3:6,:,:] = self.model.conv1.weight.data
self.conv2.weight.data[:,6:9,:,:] = self.model.conv1.weight.data
self.conv2.weight.data[:,9,:,:] = self.model.conv1.weight.data[:,0,:,:]
self.vlad1 = netvlad.NetVLAD(num_clusters=self.num_clusters, dim=512, alpha=1.0)
self.vlad2 = netvlad.NetVLAD(num_clusters=self.num_clusters, dim=512, alpha=1.0)
#self.avg_pool = nn.AdaptiveAvgPool2d((1,1))
self.bn1 = nn.BatchNorm1d(512*self.num_clusters)
self.bn2 = nn.BatchNorm1d(512*self.num_clusters)
self.classifier1 = nn.Linear(512*self.num_clusters, num_classes)
self.classifier2 = nn.Linear(512*self.num_clusters, num_classes)
for l in layers[:-5]:
for p in l.parameters():
p.requires_grad = False
if opt.mode.lower() == 'cluster' and not opt.vladv2:
layers.append(L2Norm())
encoder = nn.Sequential(
*layers) # to create small model(taken from the layers)
model = nn.Module()
model.add_module('encoder', encoder) ## add VGG (n-2) layers to the model
if opt.mode.lower() != 'cluster':
if opt.pooling.lower() == 'netvlad':
net_vlad = netvlad.NetVLAD(num_clusters=opt.num_clusters,
dim=encoder_dim,
vladv2=opt.vladv2)
if not opt.resume:
if opt.mode.lower() == 'train':
initcache = join(
opt.dataPath, 'centroids',
opt.arch + '_' + train_set.dataset + '_' +
str(opt.num_clusters) + '_desc_cen.hdf5')
else:
initcache = join(
opt.dataPath, 'centroids',
opt.arch + '_' + whole_test_set.dataset + '_' +
str(opt.num_clusters) + '_desc_cen.hdf5')
if not exists(initcache):
raise FileNotFoundError(
else:
raise ValueError('Unknown dataset split: ' + opt.split)
print('====> Query count:', whole_test_set.dbStruct.numQ)
elif opt.mode.lower() == 'cluster':
whole_train_set = dataset.get_whole_training_set(onlyDB=True)
print('===> Building model')
pretrained = not opt.fromscratch
if opt.arch.lower() == 'self_define':
print('using customized network')
if opt.dataset in ['mapillary','mapillary40k']:
encoder_dim = 256
if opt.casa:
net_vlad = netvlad.NetVLAD(num_clusters=opt.num_clusters, dim=encoder_dim, vladv2=False).cuda()
print('using Channel attention and spatial attention model')
model = network.AttentAlex(net_vlad,opt.num_clusters,encoder_dim,opt.num_PCA,opt.bn).cuda()
elif opt.atten:
print('using Spatial Attention')
if opt.rect_atten:
print('using rectified attention')
#print(opt.pooling)
if opt.pooling in ['atten_vlad']:
# print('using attention vlad!')
pca_module=None
if opt.num_PCA >0:
print('using pca module')
pca_module = layers.WPCA(opt.num_clusters*encoder_dim,opt.num_PCA)
pca_module.init_params(opt.dataset)
for l in layers[:-2]:
for p in l.parameters():
p.requires_grad = False
if opt.mode.lower() == 'cluster':
layers.append(L2Norm())
encoder = nn.Sequential(*layers)
model = nn.Module()
model.add_module('encoder', encoder)
# 初始化model中的pooling模块
if opt.mode.lower() != 'cluster':
if opt.pooling.lower() == 'netvlad':
net_vlad = netvlad.NetVLAD(num_clusters=opt.num_clusters,
dim=encoder_dim,
vladv2=False,
normalize_output=True)
if not opt.resume:
if opt.mode.lower() == 'train':
initcache = join(
opt.dataPath, 'centroids',
opt.arch + '_' + train_set.dataset + '_' +
str(opt.num_clusters) + '_desc_cen.hdf5')
else:
initcache = join(
opt.dataPath, 'centroids',
opt.arch + '_' + whole_test_set.dataset + '_' +
str(opt.num_clusters) + '_desc_cen.hdf5')
if not exists(initcache):
raise FileNotFoundError(
'Could not find clusters, please run with --mode=cluster before proceeding'
# Test scene place recognition network (NetVLAD), which is applicable to unwrapped panoramic images
from __future__ import print_function
import argparse
import random
from os.path import join, exists
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
import h5py
import faiss
import numpy as np
import netvlad
import cv2
import SceneModel
import warnings
parser = argparse.ArgumentParser(description='ScenePlaceRecognitionTest')
parser.add_argument('--cacheBatchSize', type=int, default=8, help='Batch size for caching and testing')
parser.add_argument('--nGPU', type=int, default=1, help='number of GPU to use.')
parser.add_argument('--nocuda', action='store_true', help='Dont use cuda')
parser.add_argument('--threads', type=int, default=8, help='Number of threads for each data loader to use')
parser.add_argument('--seed', type=int, default=123, help='Random seed to use.')
parser.add_argument('--dataPath', type=str, default='data/', help='Path for centroid data.')
parser.add_argument('--cachePath', type=str, default='/tmp/', help='Path to save cache to.')
parser.add_argument('--resume', type=str, default='checkpoints_res',
help='Path to load checkpoint from, for resuming training or testing.')
parser.add_argument('--ckpt', type=str, default='best',
help='Resume from latest or best checkpoint.', choices=['latest', 'best'])
parser.add_argument('--dataset', type=str, default='Yuquan',
help='DataSet to use', choices=['MOLP', 'Yuquan', 'Highway'])
