def init_model(trainId):
tf_config = tf.ConfigProto()
tf_config.gpu_options.allow_growth = True
sess = tf.Session(config=tf_config)
with sess.graph.as_default():
with sess.as_default():
tripletNet = TripletNet()
sess.run(tf.local_variables_initializer())
sess.run(tf.global_variables_initializer())
var_list = [
var for var in tf.global_variables() if "moving" in var.name
]
var_list += [
var for var in tf.global_variables()
if "global_step" in var.name
]
var_list += tf.trainable_variables()
saver = tf.train.Saver(var_list=var_list, max_to_keep=20)
last_file = tf.train.latest_checkpoint("file/" + trainId +
"/models/")
if last_file:
print('Restoring model from {}'.format(last_file))
saver.restore(sess, last_file)
writer = tf.summary.FileWriter("file/" + trainId + "/logs/train",
sess.graph)
return sess, saver, tripletNet, writer
def getModel(weights_file="./static/weights/oxbuild_final.pth"):
"""
Function that returns the model (saved during deploy stage to redce load time)
Args:
weights_file: path of trained weights file
Returns:
model based on weights_file
"""
use_cuda = torch.cuda.is_available()
device = torch.device("cuda" if use_cuda else "cpu")
resnet_model = create_embedding_net()
model = TripletNet(resnet_model)
model.load_state_dict(torch.load(weights_file))
model.to(device)
model.eval()
return model
def run(train_loader, test_pairs, options):
# Construct model
cuda = options["device_id"] != -1
if options["use_pair_feature"]:
point_net = PointNet(options)
pair_net = PairNetWithPairFeatures(options, point_net)
model = TripletNetWithPairFeatures(pair_net)
else:
point_net = PointNet(options)
pair_net = PairNet(options, point_net)
model = TripletNet(pair_net)
if cuda:
model.cuda()
print(model)
optimizer = torch.optim.Adam(model.parameters(), lr=options["lr"], weight_decay=options["weight_decay"])
loss_fn = TripletLoss(options["margin"])
best_overall_metric = 0 # a single value
best_metrics = None # a dictionary
best_epoch = 0
save_model(model, options["save_dir"], 'best', 0) # save the initial first model
for epoch in range(options["epochs"]):
epoch_loss, non_zero_triplet_ratio = train_triplet_epoch(train_loader, model, loss_fn, optimizer, cuda,
use_pair_feature=options["use_pair_feature"])
print("Epoch: {}, train-loss: {:06.4f}, non_zero_triplet_ratio: {:06.4f}, ".format(epoch, epoch_loss,
non_zero_triplet_ratio))
if epoch % options["eval_epoch"] == 0 and epoch != 0:
if options["use_pair_feature"]:
prediction_score = pair_prediction_with_pair_feature(model.pair_net, test_pairs, pair_features, cuda, options, batch_size=10000)
else:
prediction_score = pair_prediction(model.pair_net, test_pairs, cuda, options, batch_size=10000)
test_triplets = []
for term_pair, score in zip(test_pairs, prediction_score):
test_triplets.append((term_pair[0], term_pair[1], -1.0 * score))
metrics = evaluation_main(test_triplets)
if metrics["all"] >= best_overall_metric:
best_overall_metric =metrics["all"]
best_epoch = epoch
best_metrics = metrics
save_model(model, options["save_dir"], 'best', epoch) # save the initial first model
return best_overall_metric, best_epoch, best_metrics
def inference_on_single_labelled_image_pca(
query_img_file,
labels_dir,
img_dir,
img_fts_dir,
weights_file,
top_k=1000,
plot=True,
):
"""
Function that returns the average precision for a given query image and also plots the top 20 results
Args:
query_img_file : path of query image file
labels_dir : Directory for ground truth labels
img_dir : Directory holding the images
img_fts_dir : Directory holding the pca reduced features generated through create_db.py script
weights_file: path of trained weights file
top_k : top_k values used to calculate the average precison
plot : if True, top 20 results are plotted
Returns:
Average precision for the query image file
"""
# Create cuda parameters
use_cuda = torch.cuda.is_available()
np.random.seed(2019)
torch.manual_seed(2019)
device = torch.device("cuda" if use_cuda else "cpu")
print("Available device = ", device)
# Create embedding network
resnet_model = create_embedding_net()
model = TripletNet(resnet_model)
model.load_state_dict(torch.load(weights_file))
model.to(device)
model.eval()
# Get query name
query_img_name = query_img_file.split("/")[-1]
query_img_path = os.path.join(img_dir, query_img_name)
# Create Query extractor object
QUERY_EXTRACTOR = QueryExtractor(labels_dir, img_dir, subset="inference")
# Create query ground truth dictionary
query_gt_dict = QUERY_EXTRACTOR.get_query_map()[query_img_name]
# Creat image database
QUERY_IMAGES_FTS = [
os.path.join(img_fts_dir, file)
for file in sorted(os.listdir(img_fts_dir))
]
QUERY_IMAGES = [
os.path.join(img_fts_dir, file) for file in sorted(os.listdir(img_dir))
]
# Query fts
query_fts = get_query_embedding(model, device,
query_img_file).detach().cpu().numpy()
query_fts = perform_pca_on_single_vector(query_fts)
# Create similarity list
similarity = []
for file in tqdm(QUERY_IMAGES_FTS):
file_fts = np.squeeze(np.load(file))
cos_sim = np.dot(query_fts, file_fts) / (np.linalg.norm(query_fts) *
np.linalg.norm(file_fts))
similarity.append(cos_sim)
# Get best matches using similarity
similarity = np.asarray(similarity)
indexes = (-similarity).argsort()[:top_k]
best_matches = [QUERY_IMAGES[index] for index in indexes]
# Get preds
if plot:
preds = get_preds_and_visualize(best_matches, query_gt_dict, img_dir,
20)
else:
preds = get_preds(best_matches, query_gt_dict)
# Get average precision
ap = ap_per_query(best_matches, query_gt_dict)
return ap
def create_embeddings_db_pca(model_weights_path, img_dir, fts_dir):
"""
Given a model weights path, this function creates a triplet network, loads the parameters and generates the dimension
reduced (using pca) vectors and save it in the provided feature directory.
Args:
model_weights_path : path of trained weights
img_dir : directory that holds the images
fts_dir : directory to store the embeddings
Returns:
None
Eg run:
create_embeddings_db_pca("./weights/oxbuild-exp-3.pth", img_dir="./data/oxbuild/images/", fts_dir="./fts_pca/oxbuild/")
"""
# Create cuda parameters
use_cuda = torch.cuda.is_available()
np.random.seed(2019)
torch.manual_seed(2019)
device = torch.device("cuda" if use_cuda else "cpu")
print("Available device = ", device)
# Create transforms
mean = [0.485, 0.456, 0.406]
std = [0.229, 0.224, 0.225]
transforms_test = transforms.Compose([
transforms.Resize(460),
transforms.FiveCrop(448),
transforms.Lambda(lambda crops: torch.stack(
[transforms.ToTensor()(crop) for crop in crops])),
transforms.Lambda(lambda crops: torch.stack(
[transforms.Normalize(mean=mean, std=std)(crop)
for crop in crops])),
])
# Creat image database
if "paris" in img_dir:
print("> Blacklisted images must be removed")
blacklist = [
"paris_louvre_000136.jpg",
"paris_louvre_000146.jpg",
"paris_moulinrouge_000422.jpg",
"paris_museedorsay_001059.jpg",
"paris_notredame_000188.jpg",
"paris_pantheon_000284.jpg",
"paris_pantheon_000960.jpg",
"paris_pantheon_000974.jpg",
"paris_pompidou_000195.jpg",
"paris_pompidou_000196.jpg",
"paris_pompidou_000201.jpg",
"paris_pompidou_000467.jpg",
"paris_pompidou_000640.jpg",
"paris_sacrecoeur_000299.jpg",
"paris_sacrecoeur_000330.jpg",
"paris_sacrecoeur_000353.jpg",
"paris_triomphe_000662.jpg",
"paris_triomphe_000833.jpg",
"paris_triomphe_000863.jpg",
"paris_triomphe_000867.jpg",
]
files = os.listdir(img_dir)
for blacklisted_file in blacklist:
files.remove(blacklisted_file)
QUERY_IMAGES = [os.path.join(img_dir, file) for file in sorted(files)]
else:
QUERY_IMAGES = [
os.path.join(img_dir, file) for file in sorted(os.listdir(img_dir))
]
# Create dataset
eval_dataset = EmbeddingDataset(img_dir,
QUERY_IMAGES,
transforms=transforms_test)
eval_loader = DataLoader(eval_dataset,
batch_size=1,
num_workers=0,
shuffle=False)
# Create embedding network
resnet_model = create_embedding_net()
model = TripletNet(resnet_model)
model.load_state_dict(torch.load(model_weights_path))
model.to(device)
model.eval()
# Create features
with torch.no_grad():
for idx, image in enumerate(tqdm(eval_loader)):
# Move image to device and get crops
image = image.to(device)
bs, ncrops, c, h, w = image.size()
# Get output
output = model.get_embedding(image.view(-1, c, h, w))
output = output.view(bs, ncrops, -1).mean(1).cpu().numpy()
# Perform pca
output = perform_pca_on_single_vector(output)
# Save fts
img_name = (QUERY_IMAGES[idx].split("/")[-1]).replace(".jpg", "")
save_path = os.path.join(fts_dir, img_name)
np.save(save_path, output.flatten())
del output, image
gc.collect()
# if __name__ == '__main__':
# create_embeddings_db_pca("./weights/oxbuild-exp-3.pth", img_dir="./data/oxbuild/images/", fts_dir="./fts_pca/oxbuild/")
def run_predict():
out_dir = RESULTS_DIR + '/xx10'
initial_checkpoint = None
#RESULTS_DIR + '/xx10/checkpoint/00002200_model.pth'
#None
#start experiments here!
os.makedirs(out_dir + '/backup', exist_ok=True)
backup_project_as_zip(PROJECT_PATH,
out_dir + '/backup/code.%s.zip' % IDENTIFIER)
log = Logger()
log.open(out_dir + '/log.evaluate.txt', mode='a')
log.write('\n--- [START %s] %s\n\n' % (IDENTIFIER, '-' * 64))
log.write('** some experiment setting **\n')
log.write('\tSEED = %u\n' % SEED)
log.write('\tPROJECT_PATH = %s\n' % PROJECT_PATH)
log.write('\tout_dir = %s\n' % out_dir)
log.write('\n')
fig1 = plt.figure(figsize=(5, 5))
ax1 = fig1.add_subplot(111, projection='3d')
fig2 = plt.figure(figsize=(5, 5))
ax2 = fig2.add_subplot(111, projection='3d')
## net ------------------------------
net = Net().cuda()
if initial_checkpoint is not None:
log.write('\tinitial_checkpoint = %s\n' % initial_checkpoint)
net.load_state_dict(
torch.load(initial_checkpoint,
map_location=lambda storage, loc: storage))
log.write('%s\n\n' % (type(net)))
log.write('\n')
## dataset ----------------------------------------
log.write('** dataset setting **\n')
dataset = DummyDataset(
'samples_valid', #'samples_train',
mode='<not_used>',
transform=None,
)
log.write('\n')
for n in range(len(dataset)):
data, combination, label, index = dataset[n]
num_combination = len(combination)
estimate = np.zeros(num_combination, np.float32)
if 1:
ax1.clear()
ax2.clear()
x = data.x.values
y = data.y.values
z = data.z.values
ax1.plot(x, y, z, '.', color=[0.75, 0.75, 0.75], markersize=3)
ax2.plot(x, y, z, '.', color=[0.75, 0.75, 0.75], markersize=3)
for j in range(0, num_combination, 512):
batch_size = min(j + 512, num_combination) - j
truth = label[j:j + batch_size]
tracklets = combination[j:j + batch_size, :, 0:3].reshape(-1, 9)
tracklets = np.vstack(tracklets)
tracklets = torch.from_numpy(tracklets).float()
net.set_mode('test')
with torch.no_grad():
tracklets = tracklets.cuda()
logits = net.forward(tracklets)
tracklet = tracklets.data.cpu().numpy()
prob = np_sigmoid(logits.data.cpu().numpy())
estimate[j:j + batch_size] = prob
#draw results -------------------------------------------------------------
if 1:
threshold = 0.5
for i in range(batch_size):
#print(j,i)
t = tracklet[i].reshape(-1, 3)
if prob[i] > threshold and truth[i] > 0.5: #hit
print('*', end='', flush=True)
color = np.random.uniform(0, 1, (3))
ax1.plot(t[:, 0],
t[:, 1],
t[:, 2],
'.-',
color=color,
markersize=6)
if prob[i] > threshold and truth[i] < 0.5: #fp
ax2.plot(t[:, 0],
t[:, 1],
t[:, 2],
'.-',
color=[0, 0, 0],
markersize=6)
if prob[i] < threshold and truth[i] > 0.5: #miss
ax2.plot(t[:, 0],
t[:, 1],
t[:, 2],
'.-',
color=[1, 0, 0],
markersize=6)
set_figure(ax1,
title='hit @sample%d' % index,
x_limit=(0, -100),
y_limit=(-20, 20),
z_limit=(500, 1000))
set_figure(ax2,
title='error @sample%d' % index,
x_limit=(0, -100),
y_limit=(-20, 20),
z_limit=(500, 1000))
#plt.pause(0.01)
estimate = (estimate > 0.5)
label = (label > 0.5)
print('\n@sample%d :' % index)
print('\t total = %d ' % (len(label)))
print('\t hit = %d (%0.2f)' %
((estimate * label).sum(),
(estimate * label).sum() / label.sum()))
print('\t miss = %d (%0.2f)' %
((~estimate * label).sum(),
(~estimate * label).sum() / label.sum()))
print('\t fp = %d (%0.2f)' % ((estimate * (~label)).sum(),
(estimate *
(~label)).sum() / estimate.sum()))
#plt.show()
plt.pause(5)
plt.show()
pass
#assert(test_num == len(test_loader.sampler))
log.write('-------------\n')
log.write('initial_checkpoint = %s\n' % (initial_checkpoint))
#log.write('tag=%s\n'%tag)
log.write('\n')
def run_train():
out_dir = RESULTS_DIR + '/xx10'
initial_checkpoint = None
#RESULTS_DIR + '/xx10/checkpoint/00002200_model.pth'
pretrain_file = None
skip = []
## setup -----------------
os.makedirs(out_dir + '/checkpoint', exist_ok=True)
os.makedirs(out_dir + '/train', exist_ok=True)
os.makedirs(out_dir + '/backup', exist_ok=True)
backup_project_as_zip(PROJECT_PATH,
out_dir + '/backup/code.train.%s.zip' % IDENTIFIER)
log = Logger()
log.open(out_dir + '/log.train.txt', mode='a')
log.write('\n--- [START %s] %s\n\n' % (IDENTIFIER, '-' * 64))
log.write('\tSEED = %u\n' % SEED)
log.write('\tPROJECT_PATH = %s\n' % PROJECT_PATH)
log.write('\tout_dir = %s\n' % out_dir)
log.write('\n')
# fig = plt.figure(figsize=(5,5))
# ax = fig.add_subplot(111, projection='3d')
fig1 = plt.figure(figsize=(5, 5))
ax1 = fig1.add_subplot(111, projection='3d')
fig2 = plt.figure(figsize=(5, 5))
ax2 = fig2.add_subplot(111, projection='3d')
## net ----------------------
log.write('** net setting **\n')
net = Net().cuda()
if initial_checkpoint is not None:
log.write('\tinitial_checkpoint = %s\n' % initial_checkpoint)
net.load_state_dict(
torch.load(initial_checkpoint,
map_location=lambda storage, loc: storage))
# cfg = load_pickle_file(out_dir +'/checkpoint/configuration.pkl')
if pretrain_file is not None:
log.write('\tpretrain_file = %s\n' % pretrain_file)
net.load_pretrain(pretrain_file, skip)
log.write('%s\n\n' % (type(net)))
log.write('\n')
## optimiser ----------------------------------
iter_accum = 1
batch_size = 2
valid_size = 2
num_iters = 100
iter_smooth = 20
iter_log = 50
iter_valid = 100
iter_save = [0, num_iters-1]\
+ list(range(0,num_iters,200))#1*1000
LR = None #LR = StepLR([ (0, 0.01), (200, 0.001), (300, -1)])
optimizer = optim.SGD(filter(lambda p: p.requires_grad, net.parameters()),
lr=0.001 / iter_accum,
momentum=0.9,
weight_decay=0.0001)
start_iter = 0
start_epoch = 0.
if initial_checkpoint is not None:
checkpoint = torch.load(
initial_checkpoint.replace('_model.pth', '_optimizer.pth'))
start_iter = checkpoint['iter']
start_epoch = checkpoint['epoch']
rate = get_learning_rate(optimizer) #load all except learning rate
optimizer.load_state_dict(checkpoint['optimizer'])
adjust_learning_rate(optimizer, rate)
pass
## dataset ----------------------------------------
log.write('** dataset setting **\n')
train_dataset = DummyDataset('samples_train',
mode='<not_used>',
transform=train_augment)
train_loader = DataLoader(
train_dataset,
sampler=RandomSampler(train_dataset),
#sampler = SequentialSampler(train_dataset),
batch_size=batch_size,
drop_last=True,
num_workers=4,
pin_memory=True,
collate_fn=train_collate)
valid_dataset = DummyDataset('samples_valid',
mode='<not_used>',
transform=train_augment)
valid_loader = DataLoader(
train_dataset,
sampler=RandomSampler(valid_dataset),
#sampler = SequentialSampler(train_dataset),
batch_size=valid_size,
drop_last=True,
num_workers=4,
pin_memory=True,
collate_fn=train_collate)
# log.write('\ttrain_dataset.split = %s\n'%(train_dataset.split))
# log.write('\tvalid_dataset.split = %s\n'%(valid_dataset.split))
log.write('\tlen(train_dataset) = %d\n' % (len(train_dataset)))
log.write('\tlen(valid_dataset) = %d\n' % (len(valid_dataset)))
# log.write('\tlen(train_loader) = %d\n'%(len(train_loader)))
# log.write('\tlen(valid_loader) = %d\n'%(len(valid_loader)))
log.write('\tbatch_size = %d\n' % (batch_size))
log.write('\titer_accum = %d\n' % (iter_accum))
log.write('\tbatch_size*iter_accum = %d\n' % (batch_size * iter_accum))
log.write('\n')
#<debug>========================================================================================
if 0:
#fig = plt.figure(figsize=(5,5))
#ax = fig.add_subplot(111, projection='3d')
for tracklets, truths, datas, labels, lengths, indices in train_loader:
batch_size = len(indices)
print('batch_size=%d' % batch_size)
tracklets = tracklets.data.cpu().numpy()
truths = truths.data.cpu().numpy()
split = np.cumsum(lengths)
tracklets = np.split(tracklets, split)
truths = np.split(truths, split)
for b in range(batch_size):
ax.clear()
data = datas[b]
x = data.x.values
y = data.y.values
z = data.z.values
ax.plot(x, y, z, '.', color=[0.75, 0.75, 0.75], markersize=3)
tracklet = tracklets[b].reshape(-1, 3, 3)
truth = truths[b]
pos = np.where(truth == 1)[0]
for i in pos:
t = tracklet[i]
color = np.random.uniform(0, 1, (3))
ax.plot(t[:, 0],
t[:, 1],
t[:, 2],
'.-',
color=color,
markersize=6)
set_figure(ax,
x_limit=(0, -100),
y_limit=(-20, 20),
z_limit=(500, 1000))
plt.pause(0.01)
plt.show()
#<debug>========================================================================================
## start training here! ##############################################
log.write('** start training here! **\n')
log.write(' optimizer=%s\n' % str(optimizer))
log.write(' momentum=%f\n' % optimizer.param_groups[0]['momentum'])
log.write(' LR=%s\n\n' % str(LR))
log.write(' images_per_epoch = %d\n\n' % len(train_dataset))
log.write(
' rate iter epoch num | valid_loss | train_loss | batch_loss | time \n'
)
log.write(
'----------------------------------------------------------------------------------------------------\n'
)
train_loss = np.zeros(6, np.float32)
valid_loss = np.zeros(6, np.float32)
batch_loss = np.zeros(6, np.float32)
rate = 0
start = timer()
j = 0
i = 0
while i < num_iters: # loop over the dataset multiple times
sum_train_loss = np.zeros(6, np.float32)
sum = 0
net.set_mode('train')
optimizer.zero_grad()
for tracklets, truths, datas, labels, lengths, indices in train_loader:
batch_size = len(indices)
i = j / iter_accum + start_iter
epoch = (i - start_iter) * batch_size * iter_accum / len(
train_dataset) + start_epoch
num_products = epoch * len(train_dataset)
if i % iter_valid == 0:
net.set_mode('valid')
valid_loss = evaluate(net, valid_loader)
#print(valid_loss)
net.set_mode('train')
print('\r', end='', flush=True)
log.write('%0.4f %5.1f k %6.1f %4.1f m | %0.3f | %0.3f | %0.3f | %s\n' % (\
rate, i/1000, epoch, num_products/1000000,
valid_loss[0], #valid_loss[1], valid_loss[2], valid_loss[3], #valid_loss[4], valid_loss[5],#valid_acc,
train_loss[0], #train_loss[1], train_loss[2], train_loss[3], #train_loss[4], train_loss[5],#train_acc,
batch_loss[0], #batch_loss[1], batch_loss[2], batch_loss[3], #batch_loss[4], batch_loss[5],#batch_acc,
time_to_str((timer() - start)/60)))
time.sleep(0.01)
#if 1:
if i in iter_save:
torch.save(net.state_dict(),
out_dir + '/checkpoint/%08d_model.pth' % (i))
torch.save(
{
'optimizer': optimizer.state_dict(),
'iter': i,
'epoch': epoch,
}, out_dir + '/checkpoint/%08d_optimizer.pth' % (i))
# learning rate schduler -------------
if LR is not None:
lr = LR.get_rate(i)
if lr < 0: break
adjust_learning_rate(optimizer, lr / iter_accum)
rate = get_learning_rate(optimizer) * iter_accum
# one iteration update -------------
tracklets = tracklets.cuda()
truths = truths.cuda()
logits = net.forward(tracklets)
loss = F.binary_cross_entropy_with_logits(logits, truths)
# accumulated update
loss.backward()
if j % iter_accum == 0:
#torch.nn.utils.clip_grad_norm(net.parameters(), 1)
optimizer.step()
optimizer.zero_grad()
# print statistics ------------
batch_loss = np.array((
loss.cpu().data.numpy(),
0,
0,
0,
0,
0,
))
sum_train_loss += batch_loss
sum += 1
if i % iter_smooth == 0:
train_loss = sum_train_loss / sum
sum_train_loss = np.zeros(6, np.float32)
sum = 0
print('\r%0.4f %5.1f k %6.1f %4.1f m | %0.3f | %0.3f | %0.3f | %s %d,%d,%s' % (\
rate, i/1000, epoch, num_products/1000000,
valid_loss[0], #valid_loss[1], valid_loss[2], valid_loss[3], #valid_loss[4], valid_loss[5],#valid_acc,
train_loss[0], #train_loss[1], train_loss[2], train_loss[3], #train_loss[4], train_loss[5],#train_acc,
batch_loss[0], #batch_loss[1], batch_loss[2], batch_loss[3], #batch_loss[4], batch_loss[5],#batch_acc,
time_to_str((timer() - start)/60) ,i,j, ''), end='',flush=True)#str(inputs.size()))
j = j + 1
#<debug> ===================================================================
if 1:
#if i%200==0:
net.set_mode('test')
with torch.no_grad():
logits = net.forward(tracklets)
tracklets = tracklets.data.cpu().numpy()
probs = np_sigmoid(logits.data.cpu().numpy())
truths = truths.data.cpu().numpy()
batch_size = len(indices)
split = np.cumsum(lengths)
tracklets = np.split(tracklets, split)
probs = np.split(probs, split)
truths = np.split(truths, split)
for b in range(batch_size):
ax1.clear()
ax2.clear()
data = datas[b]
x = data.x.values
y = data.y.values
z = data.z.values
ax1.plot(x,
y,
z,
'.',
color=[0.75, 0.75, 0.75],
markersize=3)
ax2.plot(x,
y,
z,
'.',
color=[0.75, 0.75, 0.75],
markersize=3)
tracklet = tracklets[b]
prob = probs[b]
truth = truths[b]
#idx = np.where(prob>0.5)[0]
#for i in idx:
threshold = 0.5
for i in range(len(truth)):
t = tracklet[i].reshape(-1, 3)
if prob[i] > threshold and truth[i] > 0.5: #hit
color = np.random.uniform(0, 1, (3))
ax1.plot(t[:, 0],
t[:, 1],
t[:, 2],
'.-',
color=color,
markersize=6)
if prob[i] > threshold and truth[i] < 0.5: #fp
ax2.plot(t[:, 0],
t[:, 1],
t[:, 2],
'.-',
color=[0, 0, 0],
markersize=6)
if prob[i] < threshold and truth[i] > 0.5: #miss
ax2.plot(t[:, 0],
t[:, 1],
t[:, 2],
'.-',
color=[1, 0, 0],
markersize=6)
set_figure(ax1,
title='hit @sample%d' % indices[b],
x_limit=(0, -100),
y_limit=(-20, 20),
z_limit=(500, 1000))
set_figure(ax2,
title='error @sample%d' % indices[b],
x_limit=(0, -100),
y_limit=(-20, 20),
z_limit=(500, 1000))
plt.pause(0.01)
#fig.savefig(out_dir +'/train/%05d.png'%indices[b])
pass
net.set_mode('train')
#<debug> ===================================================================
pass #-- end of one data loader --
pass #-- end of all iterations --
if 1: #save last
torch.save(net.state_dict(),
out_dir + '/checkpoint/%d_model.pth' % (i))
torch.save(
{
'optimizer': optimizer.state_dict(),
'iter': i,
'epoch': epoch,
}, out_dir + '/checkpoint/%d_optimizer.pth' % (i))
log.write('\n')
# load the pre-shuffled train and test data
(x_train, y_train), (x_test, y_test) = cifar10.load_data()
# break training set into training and validation sets
(x_train, x_valid) = x_train[5000:], x_train[:5000]
(y_train, y_valid) = y_train[5000:], y_train[:5000]
input_size = (32, 32, 3)
embedding_dimensions = 128
batch_size = 256
gen = TripletGenerator()
train_stream = gen.flow(x_train, y_train, batch_size=batch_size)
valid_stream = gen.flow(x_valid, y_valid, batch_size=batch_size)
t = TripletNet(shape=input_size, dimensions=embedding_dimensions)
t.summary()
t.model.load_weights('model128big_batch.weights.best.hdf5', by_name=False)
checkpointer = ModelCheckpoint(filepath='model128big_batch2.weights.best.hdf5',
verbose=1,
save_best_only=True)
t.model.fit_generator(train_stream,
2500,
epochs=30,
verbose=1,
callbacks=[checkpointer],
validation_data=valid_stream,
validation_steps=20)
# t.model.fit(x_train, y_train, batch_size=64, epochs=50)
dbOut_path = sys.argv[2]
model_path = sys.argv[3]
size_embedding = int(sys.argv[4])
sizedb = [80, 80, 3]
h5 = h5py.File(db_path, 'r')
x_train = h5['X']
id_train = h5['Y_ID']
desc_train = h5['desc']
new_X, new_YID, new_desc = inizialize_dataset()
prev_id = id_train[0]
iterations = 0
t = TripletNet(shape=sizedb, dimensions=size_embedding, train=False)
t.model.load_weights(model_path, by_name=False)
borrar = 0
for index in range(len(id_train)):
curr_id = id_train[index]
if curr_id != prev_id:
borrar += 1
index_sel = (index - 1) - int(iterations / 2)
embeding = t.model.predict([
x_train[index_sel][np.newaxis, :, :, :],
x_train[index_sel][np.newaxis, :, :, :],
x_train[index_sel][np.newaxis, :, :, :]
])[1][0]
new_X.append(embeding[None])
new_YID.append(np.array([id_train[index_sel]]))
new_desc.append(desc_train[index_sel][None])
sys.stdout.flush()
else:
f.write(str(top1 / size))
f.write("," + str(top5 / size))
f.write("," + str(top10 / size) + "\n")
f.close()
if __name__ == '__main__':
os.environ["CUDA_VISIBLE_DEVICES"] = '3'
tf_config = tf.ConfigProto()
tf_config.gpu_options.allow_growth = True
sess = tf.Session(config=tf_config)
with sess.graph.as_default():
with sess.as_default():
tripletNet = TripletNet()
sess.run(tf.local_variables_initializer())
sess.run(tf.global_variables_initializer())
var_list = [
var for var in tf.global_variables() if "moving" in var.name
]
var_list += [
var for var in tf.global_variables()
if "global_step" in var.name
]
var_list += tf.trainable_variables()
saver = tf.train.Saver(var_list=var_list)
last_file = tf.train.latest_checkpoint("file/models/")
print('Restoring model from {}'.format(last_file))
saver.restore(sess, last_file)
test(tripletNet, sess, dataset="test")
def main(data_dir,
results_dir,
weights_dir,
which_dataset,
image_resize,
image_crop_size,
exp_num,
max_epochs,
batch_size,
samples_update_size,
num_workers=4,
lr=5e-6,
weight_decay=1e-5):
"""
This is the main function. You need to interface only with this function to train. (It will record all the results)
Once you have trained use create_db.py to create the embeddings and then use the inference_on_single_image.py to test
Arguments:
data_dir : parent directory for data
results_dir : directory to store the results (Make sure you create this directory first)
weights_dir : directory to store the weights (Make sure you create this directory first)
which_dataset : "oxford" or "paris"
image_resize : resize to this size
image_crop_size : square crop size
exp_num : experiment number to record the log and results
max_epochs : maximum epochs to run
batch_size : batch size (I used 5)
samples_update_size : Number of samples the network should see before it performs one parameter update (I used 64)
Keyword Arguments:
num_workers : default 4
lr : Initial learning rate (default 5e-6)
weight_decay: default 1e-5
Eg run:
if __name__ == '__main__':
main(data_dir="./data/", results_dir="./results", weights_dir="./weights",
which_dataset="oxbuild", image_resize=460, image_crop_size=448,
exp_num=3, max_epochs=10, batch_size=5, samples_update_size=64)
"""
# Define directories
labels_dir = os.path.join(data_dir, which_dataset, "gt_files")
image_dir = os.path.join(data_dir, which_dataset, "images")
# Create Query extractor object
q_train = QueryExtractor(labels_dir, image_dir, subset="train")
q_valid = QueryExtractor(labels_dir, image_dir, subset="valid")
# Create transformss
mean = [0.485, 0.456, 0.406]
std = [0.229, 0.224, 0.225]
transforms_train = transforms.Compose([
transforms.Resize(image_resize),
transforms.RandomResizedCrop(image_crop_size, scale=(0.8, 1.2)),
transforms.ColorJitter(brightness=(0.80, 1.20)),
transforms.RandomHorizontalFlip(p=0.50),
transforms.RandomChoice([
transforms.RandomRotation(15),
transforms.Grayscale(num_output_channels=3),
]),
transforms.ToTensor(),
transforms.Normalize(mean=mean, std=std),
])
transforms_valid = transforms.Compose([
transforms.Resize(image_resize),
transforms.CenterCrop(image_crop_size),
transforms.ToTensor(),
transforms.Normalize(mean=mean, std=std),
])
# Create dataset
dataset_train = VggImageRetrievalDataset(labels_dir,
image_dir,
q_train,
transforms=transforms_train)
dataset_valid = VggImageRetrievalDataset(labels_dir,
image_dir,
q_valid,
transforms=transforms_valid)
# Create dataloader
train_loader = DataLoader(dataset_train,
batch_size=batch_size,
num_workers=num_workers,
shuffle=True)
valid_loader = DataLoader(dataset_valid,
batch_size=batch_size,
num_workers=num_workers,
shuffle=False)
# Create cuda parameters
use_cuda = torch.cuda.is_available()
np.random.seed(2020)
torch.manual_seed(2020)
device = torch.device("cuda" if use_cuda else "cpu")
# Create embedding network
embedding_model = create_embedding_net()
model = TripletNet(embedding_model)
model.to(device)
# Create optimizer and scheduler
optimizer = optim.Adam(model.parameters(),
lr=lr,
weight_decay=weight_decay)
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=10)
# Create log file
log_file = open(os.path.join(results_dir, "log-{}.txt".format(exp_num)),
"w+")
log_file.write("----------Experiment {}----------\n".format(exp_num))
log_file.write("Dataset = {}, Image sizes = {}, {}\n".format(
which_dataset, image_resize, image_crop_size))
# Creat batch update value
update_batch = int(math.ceil(float(samples_update_size) / batch_size))
model_name = "{}-exp-{}.pth".format(which_dataset, exp_num)
loss_plot_save_path = os.path.join(
results_dir, "{}-loss-exp-{}.png".format(which_dataset, exp_num))
# Print stats before starting training
print("Running VGG Image Retrieval Training script")
print("Dataset used\t\t:{}".format(which_dataset))
print("Max epochs\t\t: {}".format(max_epochs))
print("Gradient update\t\t: every {} batches ({} samples)".format(
update_batch, samples_update_size))
print("Initial Learning rate\t: {}".format(lr))
print("Image resize, crop size\t: {}, {}".format(image_resize,
image_crop_size))
print("Available device \t:", device)
# Train the triplet network
tr_hist, val_hist = train_model(model,
device,
optimizer,
scheduler,
train_loader,
valid_loader,
epochs=max_epochs,
update_batch=update_batch,
model_name=model_name,
save_dir=weights_dir,
log_file=log_file)
# Close the file
log_file.close()
# Plot and save
plot_history(tr_hist,
val_hist,
"Triplet Loss",
loss_plot_save_path,
labels=["train", "validation"])
# if __name__ == '__main__':
# main(data_dir="./data/", results_dir="./results", weights_dir="./weights",
# which_dataset="oxbuild", image_resize=460, image_crop_size=448,
# exp_num=3, max_epochs=10, batch_size=5, samples_update_size=64)
cv2.cvtColor(cv2.imread(hp.image_dir + prediction[5]),
cv2.COLOR_BGR2RGB))
ax8.axis('off')
ax8.set_title('Similarity: %.4f' % score[5])
ax9.imshow(
cv2.cvtColor(cv2.imread(hp.image_dir + prediction[6]),
cv2.COLOR_BGR2RGB))
ax9.axis('off')
ax9.set_title('Similarity: %.4f' % score[6])
plt.tight_layout()
plt.show()
if __name__ == "__main__":
print('Use GPU: ', use_gpu)
model = torch.nn.DataParallel(TripletNet())
if use_gpu:
model = model.cuda()
torch.backends.cudnn.benchmark = True
if os.path.exists(hp.logdir + 'model.pkl'):
if use_gpu:
map_location = lambda storage, loc: storage.cuda()
else:
map_location = 'cpu'
ckpt = torch.load(hp.logdir + 'model.pkl', map_location=map_location)
model.load_state_dict(ckpt['state_dict'])
print('Restore model')
model.eval()
cs_func = torch.nn.CosineSimilarity(dim=1, eps=1e-6)
# Construct testing set
f_pred = options["pred_pairs_file_in"]
print("!!! Loading term pairs for prediction from: {}".format(f_pred))
pred_pairs = load_element_pairs(f_pred, with_label=False)
print("Number of term pairs for prediction: {}".format(len(pred_pairs)))
# Construct model skeleton
cuda = options["device_id"] != -1
if options["use_pair_feature"]:
point_net = PointNet(options)
pair_net = PairNetWithPairFeatures(options, point_net)
model = TripletNetWithPairFeatures(pair_net)
else:
point_net = PointNet(options)
pair_net = PairNet(options, point_net)
model = TripletNet(pair_net)
if cuda:
model.cuda()
# Load pre-trained model
model_path = options["snapshot"]
model.load_state_dict(torch.load(model_path))
print(model)
# Conduct pair prediction and dump results to file
if options["use_pair_feature"]:
prediction_scores = pair_prediction_with_pair_feature(model.pair_net, pred_pairs, pair_features, cuda, options,
batch_size=10000)
else:
prediction_scores = pair_prediction(model.pair_net, pred_pairs, cuda, options, batch_size=10000)
f_res = options["pred_pairs_file_out"]