def Get_labels_for_evaluation(self, dataloader):
LogText('Predictions for evaluation FAN', self.experiment_name,
self.log_path)
self.model.eval()
keypoints = {}
for i_batch, sample in enumerate(dataloader):
input = Cuda(sample['image'])
bsize = input.size(0)
name = sample['filename']
groundtruth = sample['groundtruth']
is_test_sample = sample['is_it_test_sample']
with torch.no_grad():
output = self.model.forward(input)
output = output[:, torch.from_numpy(self.active_channels)]
for i in range(bsize):
sampleKeypoints = Utils.GetPointsFromHeatmaps(
output[i])[:, :3].detach().cpu().numpy()
sampleKeypoints[
sampleKeypoints[:, 2] < self.confidence_thres_FAN] = np.nan
sampleKeypoints = sampleKeypoints[:, :2]
samplegroundtruth = groundtruth[i].detach().cpu().numpy()
keypoints[name[i]] = {
'prediction': sampleKeypoints,
'groundtruth': samplegroundtruth,
'is_it_test_sample': is_test_sample[i]
}
self.save_keypoints(keypoints,
f'EvaluateStep2Keypoints{self.iterations}.pickle')
return keypoints
def Train_step2(self, dataloader):
self.model.train()
count = 0
LogText(f"Epoch {self.epoch} Training Begins", self.experiment_name,
self.log_path)
for i_batch, sample in enumerate(dataloader):
if (self.iterations > 0 and self.iterations in self.lrstep):
self.schedualer.step()
LogText('LR ' + str(self.optimizer.param_groups[0]['lr']),
self.experiment_name, self.log_path)
self.iterations += 1
if (count == 0):
self.optimizer.zero_grad()
count = self.batch_multiplier
input = Cuda(sample['image'])
GaussianShape = Cuda(sample['GaussianShape'])
GaussianShape = GaussianShape[:, self.active_channels, :, :]
heatmaps_with_keypoints = Cuda(sample['heatmaps_with_keypoints'])
heatmaps_with_keypoints = heatmaps_with_keypoints[:, self.
active_channels]
output_shape = self.model(input)
output_shape = output_shape[:, self.active_channels, :, :]
loss = torch.mean(
self.criterion(output_shape,
GaussianShape)[heatmaps_with_keypoints])
loss.backward()
count -= 1
if (count == 0):
self.optimizer.step()
self.iterations += 1
LogText(
'Epoch ' + str(self.epoch) + ' completed, iterations ' +
str(self.iterations), self.experiment_name, self.log_path)
self.save_step2()
self.epoch += 1
def __init__(self,
number_of_clusters,
confidence_thres_superpoint,
nms_thres_superpoint,
path_to_pretrained_superpoint,
experiment_name,
log_path,
remove_superpoint_outliers_percentage,
use_box=False,
UseScales=False,
RemoveBackgroundClusters=False):
self.path_to_pretrained_superpoint = path_to_pretrained_superpoint
self.use_box = use_box
self.confidence_thres_superpoint = confidence_thres_superpoint
self.nms_thres_superpoint = nms_thres_superpoint
self.log_path = log_path
self.remove_superpoint_outliers_percentage = remove_superpoint_outliers_percentage
self.experiment_name = experiment_name
self.number_of_clusters = number_of_clusters
self.model = Cuda(SuperPointNet())
self.UseScales = UseScales
self.RemoveBackgroundClusters = RemoveBackgroundClusters
if (self.UseScales):
self.SuperpointUndoScaleDistill1 = iaa.Affine(scale={
"x": 1 / 1.3,
"y": 1 / 1.3
})
self.SuperpointUndoScaleDistill2 = iaa.Affine(scale={
"x": 1 / 1.6,
"y": 1 / 1.6
})
try:
checkpoint = torch.load(path_to_pretrained_superpoint,
map_location='cpu')
self.model.load_state_dict(checkpoint['state_dict'])
LogText(
f"Superpoint Network from checkpoint {path_to_pretrained_superpoint}",
self.experiment_name, self.log_path)
except:
LogText(f"Superpoint network failed to load", self.experiment_name,
self.log_path)
self.softmax = torch.nn.Softmax(dim=1)
self.pixelSuffle = torch.nn.PixelShuffle(8)
self.model.eval()
def load_trained_secondstep_model(self, checkpoint_filename):
LogText(
f"Pretrained Second Step model loaded from : {checkpoint_filename}",
self.experiment_name, self.log_path)
checkpoint = torch.load(checkpoint_filename, map_location='cpu')
self.iterations = checkpoint['iteration']
self.active_channels = checkpoint['active_channels']
self.model.load_state_dict(checkpoint['state_dict'])
self.optimizer.load_state_dict(checkpoint['optimizer'])
self.schedualer.load_state_dict(checkpoint['schedualer'])
def load_trained_fiststep_model(self, checkpoint_filename):
LogText(
f"Pretrained First Step model loaded from : {checkpoint_filename}",
self.experiment_name, self.log_path)
checkpoint = torch.load(checkpoint_filename, map_location='cpu')
self.iterations = checkpoint['iteration']
self.centroid = checkpoint['centroid']
if (self.centroid is not None):
self.KmeansClustering = clustering.Kmeans(self.number_of_clusters,
self.centroid)
self.model.load_state_dict(checkpoint['state_dict'])
self.optimizer.load_state_dict(checkpoint['optimizer'])
self.schedualer = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(
self.optimizer, T_0=5, T_mult=1, eta_min=5e-6)
def init_secondstep(self,
lr,
weight_decay,
batch_multiplier,
number_of_clusters,
lrstep,
clusteroverlap,
path_to_checkpoint=None):
self.iterations = 0
self.epoch = 0
self.batch_multiplier = batch_multiplier
self.weight_decay = weight_decay
self.lr = lr
self.lrstep = lrstep
if (path_to_checkpoint is not None):
try:
LogText(
f"Fan Initiated from weights of : {path_to_checkpoint}",
self.experiment_name, self.log_path)
checkpoint = torch.load(path_to_checkpoint, map_location='cpu')
self.model.load_state_dict(checkpoint['state_dict'])
except:
raise Exception(
f'Loading weights for FAN from {path_to_checkpoint} failed.'
)
self.number_of_clusters = number_of_clusters
self.clusteroverlap = clusteroverlap
self.active_channels = np.arange(self.number_of_clusters)
newlayer1 = Cuda(
nn.Conv2d(256,
self.number_of_clusters,
kernel_size=1,
stride=1,
padding=0))
self.model._modules['l1'] = newlayer1
self.optimizer = torch.optim.RMSprop(self.model.parameters(),
lr=self.lr,
weight_decay=self.weight_decay)
self.schedualer = torch.optim.lr_scheduler.StepLR(self.optimizer,
step_size=1,
gamma=0.1)
def init_firststep(self, lr, weight_decay, number_of_clusters,
training_iterations_before_first_clustering):
LogText(f"Training model initiated", self.experiment_name,
self.log_path)
self.weight_decay = weight_decay
self.lr = lr
self.training_iterations_before_first_clustering = training_iterations_before_first_clustering
self.number_of_clusters = number_of_clusters
self.optimizer = torch.optim.RMSprop(self.model.parameters(),
lr=self.lr,
weight_decay=self.weight_decay)
self.schedualer = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(
self.optimizer, T_0=5, T_mult=1, eta_min=5e-6)
self.centroid = None
self.margin = 0.8
self.eps = 1e-9
self.KmeansClustering = clustering.Kmeans(self.number_of_clusters)
self.iterations = 0
def ShowVisualRes(keypoints,log_path,experiment_name,number_of_clusters,dataset_name):
fig = plt.figure(figsize=(34,55))
gs1 = gridspec.GridSpec(13, 8)
gs1.update(wspace=0.0, hspace=0.0)
filenames=[k for k in keypoints.keys() if keypoints[k]['is_it_test_sample']]
filenames.sort()
filenames=filenames[:13*8]
dataset = Database( dataset_name, number_of_clusters,test=True)
for i in range(len(filenames)):
ax = plt.subplot(gs1[i])
plt.axis('off')
pointstoshow = keypoints[filenames[i]]['prediction']
image = dataset.getimage_FAN(dataset, filenames[i])
ax.imshow(image)
colors = [Utils.colorlist[int(i)] for i in np.arange(len(pointstoshow))]
ax.scatter(pointstoshow[:, 0], pointstoshow[:, 1], s=400, c=colors, marker='P',edgecolors='black', linewidths=0.3)
fig.show()
fig.savefig(log_path+ 'Logs/' + experiment_name + f'/Step2.jpg')
LogText(f"Step2 results created in {log_path+ 'Logs/' + experiment_name + f'/Step2.jpg'}", experiment_name,log_path)
def ShowClusters(keypoints,log_path,experiment_name,number_of_clusters,dataset_name):
dataset = Database( dataset_name, number_of_clusters )
image_names=list(keypoints.keys())
random.shuffle(image_names)
for cluster_number in range(number_of_clusters):
counter_figureimages=0
counter_datasetimages=0
fig, subplots= plt.subplots(8,8,figsize=(15,15))
subplots=subplots.reshape(-1)
fig.subplots_adjust(wspace=0,hspace=0)
for s in subplots:
s.set_axis_off()
while counter_figureimages<64:
#for the case where cluster has less than 64 instances
if(counter_datasetimages>len(keypoints)-1):
fig.savefig(log_path+ 'Logs/' + experiment_name + f'/Cluster{cluster_number}.jpg')
break
imagename=image_names[counter_datasetimages]
imagepoints = keypoints[imagename]
#if cluster exists in image
if(sum(imagepoints[:, 2]==cluster_number)>0):
image = dataset.getimage_FAN(dataset,imagename)
ax=subplots[counter_figureimages]
ax.imshow(image)
ax.scatter(4*imagepoints[imagepoints[:, 2]==cluster_number,0], 4*imagepoints[imagepoints[:, 2]==cluster_number, 1])
counter_figureimages+=1
counter_datasetimages+=1
fig.savefig(log_path+ 'Logs/' + experiment_name + f'/Cluster{cluster_number}.jpg')
LogText(f"Cluster images created in {log_path+ 'Logs/' + experiment_name + f'/Cluster{cluster_number}.jpg'}", experiment_name,log_path)
def ShowKeypoints(keypoints,log_path,experiment_name,number_of_clusters,dataset_name):
dataset = Database( dataset_name, number_of_clusters )
count=0
image_names=list(keypoints.keys())
random.shuffle(image_names)
fig, subplots= plt.subplots(8,8,figsize=(15,15))
subplots=subplots.reshape(-1)
fig.subplots_adjust(wspace=0,hspace=0)
for s in subplots:
s.set_axis_off()
while count<8*8:
imagepoints = keypoints[image_names[count]]
image = dataset.getimage_FAN(dataset,image_names[count])
ax=subplots[count]
ax.imshow(image)
ax.scatter(4*imagepoints[:,0], 4*imagepoints[:, 1])
count+=1
fig.savefig(log_path+ 'Logs/' + experiment_name + f'/Keypoints.jpg')
LogText(f"Keypoint images created in {log_path+ 'Logs/' + experiment_name + f'/Keypoints.jpg'}", experiment_name,log_path)
def MergeClusters(self, dataloader, Image_Keypoints):
LogText('Predictions for evaluation FAN', self.experiment_name,
self.log_path)
self.model.eval()
Image_shapeKeypoints = {}
for i_batch, sample in enumerate(dataloader):
input = Cuda(sample['image'])
name = sample['filename']
with torch.no_grad():
output = self.model.forward(input)
output = output[:, torch.from_numpy(self.active_channels)]
bsize = output.size(0)
for i in range(bsize):
Image_shapeKeypoints[name[i]] = Utils.GetPointsFromHeatmaps(
output[i])
# get points per cluster
points_per_cluster = np.zeros(self.number_of_clusters)
for index in Image_Keypoints:
cluster_assignments = Image_Keypoints[index][:, 2]
points_per_cluster[cluster_assignments.astype(int)] += 1
points_per_channel = points_per_cluster[self.active_channels]
totalDistanceMatrix = np.zeros(
(len(self.active_channels), len(self.active_channels)))
# get spatial distance between clusters
numberofconfidentpointspercluster = np.zeros(len(self.active_channels))
for image in Image_shapeKeypoints.keys():
points = Image_shapeKeypoints[image].detach().cpu().numpy()
distancematrix = squareform(pdist(points[:, :2]))
numberofconfidentpointspercluster[points[:, 2] > 0.2] += 1
distancematrix[points[:, 2] < self.confidence_thres_FAN, :] = 300
distancematrix[:, points[:, 2] < self.confidence_thres_FAN] = 300
#5.7 corresponds to nms of size 1 on the 64x64 dimension
totalDistanceMatrix = totalDistanceMatrix + (distancematrix <
5.7).astype(int)
confident_points_per_channel = np.diag(totalDistanceMatrix).copy()
np.fill_diagonal(totalDistanceMatrix, 0)
indexes_sorted = np.argsort(points_per_channel)[::-1]
points_of_smaller_cluster = np.zeros(
(len(self.active_channels), len(self.active_channels)))
for x in range(len(self.active_channels)):
for y in range(len(self.active_channels)):
points_of_smaller_cluster[x, y] = min(
numberofconfidentpointspercluster[x],
numberofconfidentpointspercluster[y])
indexes_channels_to_extend = np.array([])
indexes_channels_merged = []
while (len(indexes_sorted) > 0):
channel = indexes_sorted[0]
is_remaining = True
for i in range(len(indexes_channels_to_extend)):
element = indexes_channels_to_extend[i]
if (totalDistanceMatrix[int(element),
int(channel)] > self.clusteroverlap *
points_of_smaller_cluster[int(element),
int(channel)]):
indexes_channels_merged[i] = np.append(
indexes_channels_merged[i], int(channel)).astype(int)
is_remaining = False
indexes_sorted = np.delete(indexes_sorted, 0)
break
if (is_remaining):
indexes_channels_to_extend = np.append(
indexes_channels_to_extend, int(channel))
indexes_channels_merged.append(np.array([]))
indexes_sorted = np.delete(indexes_sorted, 0)
extendclusters = self.active_channels[
indexes_channels_to_extend.astype(int)]
clusters_merged = []
for el in indexes_channels_merged:
clusters_merged.append(self.active_channels[el.astype(int)])
pairs_to_keep = np.array([len(f) > 0 for f in clusters_merged])
extendclusters = extendclusters[pairs_to_keep].astype(int)
clusters_merged = np.array(clusters_merged)[pairs_to_keep]
count = 0
if (len(extendclusters) > 0):
LogText("Clusters merged:", self.experiment_name, self.log_path)
for s in range(len(extendclusters)):
LogText(f"{extendclusters[s]} -> {clusters_merged[s]}",
self.experiment_name, self.log_path)
# substitute merged clusters
for index in Image_Keypoints:
keypoint = Image_Keypoints[index]
for p in range(len(extendclusters)):
indeces_of_keypoints_to_merge = np.in1d(
keypoint[:, 2], clusters_merged[p])
if (sum(indeces_of_keypoints_to_merge) == 0):
continue
elif (sum(indeces_of_keypoints_to_merge) > 0):
indeces_of_keypoints_to_merge = np.in1d(
keypoint[:, 2],
np.append(clusters_merged[p], extendclusters[p]))
clusterinimage = keypoint[:, 2][
indeces_of_keypoints_to_merge].astype(int)
index_of_bigger_cluster = np.argmax(
points_per_cluster[clusterinimage])
cluster_to_remove = np.delete(clusterinimage,
index_of_bigger_cluster)
indexes_to_keep = np.in1d(keypoint[:, 2],
cluster_to_remove) == False
keypoint = keypoint[indexes_to_keep]
indeces_of_keypoints_to_merge = np.in1d(
keypoint[:, 2], clusters_merged[p])
keypoint[:, 2][
indeces_of_keypoints_to_merge] = extendclusters[p]
Image_Keypoints[index] = keypoint
self.active_channels = self.active_channels[
indexes_channels_to_extend.astype(int)]
self.active_channels.sort()
LogText(f"Remaining Clusters: {len(self.active_channels)}",
self.experiment_name, self.log_path)
self.save_keypoints(Image_Keypoints,
f'MergedKeypoints{self.iterations}.pickle')
return Image_Keypoints
def Update_pseudoLabels(self, dataloader, oldkeypoints=None):
LogText(f"Clustering stage for iteration {self.iterations}",
self.experiment_name, self.log_path)
self.model.eval()
imagesize = 256
heatmapsize = 64
numberoffeatures = 256
buffersize = 500000
# allocation of 2 buffers for temporal storing of keypoints and descriptors.
Keypoint_buffer = torch.zeros(buffersize, 3)
Descriptor__buffer = torch.zeros(buffersize, numberoffeatures)
# arrays on which we save buffer content periodically. Corresponding files are temporal and
# will be deleted after the completion of the process
CreateFileArray(
str(
GetCheckPointsPath(self.experiment_name, self.log_path) /
'keypoints'), 3)
CreateFileArray(
str(
GetCheckPointsPath(self.experiment_name, self.log_path) /
'descriptors'), numberoffeatures)
# intermediate variables
first_index = 0
last_index = 0
buffer_first_index = 0
buffer_last_index = 0
keypoint_indexes = {}
pointsperimage = 0
LogText(f"Inference of keypoints and descriptors begins",
self.experiment_name, self.log_path)
for i_batch, sample in enumerate(dataloader):
input = Cuda(sample['image'])
names = sample['filename']
with torch.no_grad():
output = self.model.forward(input)
outputHeatmap = output[0]
descriptors_volume = output[1]
batch_keypoints = GetBatchMultipleHeatmap(
outputHeatmap, self.confidence_thres_FAN)
for i in range(input.size(0)):
indexes = batch_keypoints[:, 0] == i
sample_keypoints = batch_keypoints[indexes, 1:][:, :3]
pointsperimage += len(sample_keypoints)
if (oldkeypoints is not None):
if (names[i] in oldkeypoints):
keypoints_previous_round = Cuda(
torch.from_numpy(
oldkeypoints[names[i]].copy())).float()
sample_keypoints = MergePoints(
sample_keypoints, keypoints_previous_round)
descriptors = GetDescriptors(descriptors_volume[i],
sample_keypoints[:, :2],
heatmapsize, heatmapsize)
numofpoints = sample_keypoints.shape[0]
last_index += numofpoints
buffer_last_index += numofpoints
Keypoint_buffer[buffer_first_index:buffer_last_index, :
2] = sample_keypoints.cpu()[:, :2]
Descriptor__buffer[
buffer_first_index:buffer_last_index, :] = descriptors
keypoint_indexes[names[i]] = [first_index, last_index]
first_index += numofpoints
buffer_first_index += numofpoints
# periodically we store the buffer in file
if buffer_last_index > int(buffersize * 0.8):
AppendFileArray(
np.array(Keypoint_buffer[:buffer_last_index]),
str(
GetCheckPointsPath(self.experiment_name, self.log_path)
/ 'keypoints'))
AppendFileArray(
np.array(Descriptor__buffer[:buffer_last_index]),
str(
GetCheckPointsPath(self.experiment_name, self.log_path)
/ 'descriptors'))
Keypoint_buffer = torch.zeros(buffersize, 3)
Descriptor__buffer = torch.zeros(buffersize, numberoffeatures)
buffer_first_index = 0
buffer_last_index = 0
# store any keypoints left on the buffers
AppendFileArray(
np.array(Keypoint_buffer[:buffer_last_index]),
str(
GetCheckPointsPath(self.experiment_name, self.log_path) /
'keypoints'))
AppendFileArray(
np.array(Descriptor__buffer[:buffer_last_index]),
str(
GetCheckPointsPath(self.experiment_name, self.log_path) /
'descriptors'))
# load handlers to the Keypoints and Descriptor files
Descriptors, fileHandler1 = OpenreadFileArray(
str(
GetCheckPointsPath(self.experiment_name, self.log_path) /
'descriptors'))
Keypoints, fileHandler2 = OpenreadFileArray(
str(
GetCheckPointsPath(self.experiment_name, self.log_path) /
'keypoints'))
Keypoints = Keypoints[:, :]
LogText(
f"Keypoints Detected per image Only detector {pointsperimage / len(keypoint_indexes)}",
self.experiment_name, self.log_path)
LogText(f"Inference of keypoints and descriptors completed",
self.experiment_name, self.log_path)
LogText(
f"Keypoints Detected per image {len(Keypoints)/len(keypoint_indexes)}",
self.experiment_name, self.log_path)
# we use a subset of all the descriptors for clustering based on the recomendation of the Faiss repository
numberOfPointsForClustering = 500000
descriptors = clustering.preprocess_features(
Descriptors[:numberOfPointsForClustering])
_, self.centroid = self.KmeansClustering.cluster(descriptors,
verbose=False)
self.KmeansClustering.clus.nredo = 1
thresholds = self.GetThresholdsPerCluster(Descriptors)
Image_Keypoints = {}
averagepointsperimage = 0
for image in keypoint_indexes:
start, end = keypoint_indexes[image]
keypoints = Keypoints[start:end, :]
image_descriptors = clustering.preprocess_features(
Descriptors[start:end])
# calculate distance of each keypoints to each centroid
distanceMatrix, clustering_assignments = self.KmeansClustering.index.search(
image_descriptors, self.number_of_clusters)
distanceMatrix = np.take_along_axis(
distanceMatrix, np.argsort(clustering_assignments), axis=-1)
# assign keypoints to centroids using the Hungarian algorithm. This ensures that each
# image has at most one instance of each cluster
keypointIndex, clusterAssignment = linear_sum_assignment(
distanceMatrix)
tempKeypoints = np.zeros((len(keypointIndex), 3))
tempKeypoints = keypoints[keypointIndex]
clusterAssignmentDistance = distanceMatrix[keypointIndex,
clusterAssignment]
clusterstokeep = np.zeros(len(clusterAssignmentDistance))
clusterstokeep = clusterstokeep == 1
# keep only points that lie in their below a cluster specific theshold
clusterstokeep[clusterAssignmentDistance <
thresholds[clusterAssignment]] = True
tempKeypoints[:, 2] = clusterAssignment
Image_Keypoints[image] = tempKeypoints[clusterstokeep]
averagepointsperimage += sum(clusterstokeep)
#initialise centroids for next clustering round
self.KmeansClustering = clustering.Kmeans(self.number_of_clusters,
self.centroid)
LogText(
f"Keypoints Detected per image {averagepointsperimage/len(Image_Keypoints)}",
self.experiment_name, self.log_path)
self.save_keypoints(Image_Keypoints,
f'UpdatedKeypoints{self.iterations}.pickle')
ClosereadFileArray(
fileHandler1,
str(
GetCheckPointsPath(self.experiment_name, self.log_path) /
'keypoints'))
ClosereadFileArray(
fileHandler2,
str(
GetCheckPointsPath(self.experiment_name, self.log_path) /
'descriptors'))
LogText(f"Clustering stage completed", self.experiment_name,
self.log_path)
return Image_Keypoints
def Train_step1(self, dataloader):
LogText(f"Training Begins", self.experiment_name, self.log_path)
self.model.train()
while (True):
for i_batch, sample in enumerate(dataloader):
self.optimizer.zero_grad()
if (self.iterations % 2000 == 0):
LogText(f"Iterations : {self.iterations}",
self.experiment_name, self.log_path)
if (self.iterations ==
self.training_iterations_before_first_clustering):
LogText(f"Intial training stage completed",
self.experiment_name, self.log_path)
self.iterations += 1
self.save_step1()
return
if (self.iterations % 2000 == 0 and self.iterations >
self.training_iterations_before_first_clustering):
self.schedualer.step()
LogText(
f'Current Learning rate :' +
str(self.optimizer.param_groups[0]['lr']),
self.experiment_name, self.log_path)
if (self.optimizer.param_groups[0]['lr'] == self.lr):
self.iterations += 1
self.save_step1()
return
input = Cuda(sample['image'])
descriptorpairs = Cuda(sample['keypoints'])
keypointHeatmaps = (Cuda(sample['keypointHeatmaps']))
bsize = input.size(0)
number_of_pairs = descriptorpairs.size(1)
batchid = Cuda(
torch.arange(bsize).repeat(number_of_pairs).reshape(
number_of_pairs, bsize).transpose(1, 0))
target = Cuda(descriptorpairs[:, :, 4].reshape(-1).clone())
output1_detector, output1_descriptor = self.model(
input[:, 0:3, :, :])
output2_detector, output2_descriptor = self.model(
input[:, 3:, :, :])
loss_detector1 = self.criterion(output1_detector,
keypointHeatmaps[:, 0:1, :, :])
loss_detector2 = self.criterion(output2_detector,
keypointHeatmaps[:, 1:2, :, :])
output1features = output1_descriptor[
batchid.reshape(-1).long(), :,
descriptorpairs[:, :, 1].reshape(-1).long(),
descriptorpairs[:, :, 0].reshape(-1).long()]
output2features = output2_descriptor[
batchid.reshape(-1).long(), :,
descriptorpairs[:, :, 3].reshape(-1).long(),
descriptorpairs[:, :, 2].reshape(-1).long()]
distances = (
output2features[descriptorpairs[:, :, 0].reshape(-1) != -1]
-
output1features[descriptorpairs[:, :, 0].reshape(-1) != -1]
).pow(2).sum(1)
descriptor_losses = (
target[descriptorpairs[:, :, 0].reshape(-1) != -1].float()
* distances +
(1 +
-1 * target[descriptorpairs[:, :, 0].reshape(-1) != -1]
).float() * torch.nn.functional.relu(
self.margin - (distances + self.eps).sqrt()).pow(2))
descriptor_losses = descriptor_losses.mean()
loss = 10 * descriptor_losses + loss_detector1 + loss_detector2
loss.backward()
self.optimizer.step()
self.iterations += 1
def CreateInitialPseudoGroundtruth(self, dataloader):
LogText(f"Extraction of initial Superpoint pseudo groundtruth",
self.experiment_name, self.log_path)
imagesize = 256
heatmapsize = 64
numberoffeatures = 256
buffersize = 500000
#allocation of 2 buffers for temporal storing of keypoints and descriptors.
Keypoint_buffer = torch.zeros(buffersize, 3)
Descriptor__buffer = torch.zeros(buffersize, numberoffeatures)
#arrays on which we save buffer content periodically. Corresponding files are temporal and
#will be deleted after the completion of the process
CreateFileArray(
str(
GetCheckPointsPath(self.experiment_name, self.log_path) /
'keypoints'), 3)
CreateFileArray(
str(
GetCheckPointsPath(self.experiment_name, self.log_path) /
'descriptors'), numberoffeatures)
#intermediate variables
first_index = 0
last_index = 0
buffer_first_index = 0
buffer_last_index = 0
keypoint_indexes = {}
LogText(f"Inference of Keypoints begins", self.experiment_name,
self.log_path)
for i_batch, sample in enumerate(dataloader):
input = Cuda(sample['image_gray'])
names = sample['filename']
bsize = input.size(0)
if (self.UseScales):
input = input.view(-1, 1, input.shape[2], input.shape[3])
with torch.no_grad():
detectorOutput, descriptorOutput = self.GetSuperpointOutput(
input)
if (self.UseScales):
detectorOutput = detectorOutput.view(bsize, -1,
detectorOutput.shape[2],
detectorOutput.shape[3])
input = input.view(bsize, -1, input.shape[2], input.shape[3])
descriptorOutput = descriptorOutput.view(
bsize, -1, descriptorOutput.size(1),
descriptorOutput.size(2), descriptorOutput.size(3))[:, 0]
for i in range(0, bsize):
keypoints = self.GetPoints(detectorOutput[i].unsqueeze(0),
self.confidence_thres_superpoint,
self.nms_thres_superpoint)
if (self.RemoveBackgroundClusters):
bounding_box = sample['bounding_box'][i]
pointsinbox = torch.ones(len(keypoints))
pointsinbox[(keypoints[:, 0] < int(bounding_box[0]))] = -1
pointsinbox[(keypoints[:, 1] < int(bounding_box[1]))] = -1
pointsinbox[(keypoints[:, 0] > int(bounding_box[2]))] = -1
pointsinbox[(keypoints[:, 1] > int(bounding_box[3]))] = -1
elif (self.use_box):
bounding_box = sample['bounding_box'][i]
pointsinbox = torch.ones(len(keypoints))
pointsinbox[(keypoints[:, 0] < int(bounding_box[0]))] = -1
pointsinbox[(keypoints[:, 1] < int(bounding_box[1]))] = -1
pointsinbox[(keypoints[:, 0] > int(bounding_box[2]))] = -1
pointsinbox[(keypoints[:, 1] > int(bounding_box[3]))] = -1
keypoints = keypoints[pointsinbox == 1]
descriptors = GetDescriptors(descriptorOutput[i], keypoints,
input.shape[3], input.shape[2])
#scale image keypoints to FAN resolution
keypoints = dataloader.dataset.keypointsToFANResolution(
dataloader.dataset, names[i], keypoints)
keypoints = ((heatmapsize / imagesize) * keypoints).round()
last_index += len(keypoints)
buffer_last_index += len(keypoints)
Keypoint_buffer[
buffer_first_index:buffer_last_index, :2] = keypoints
Descriptor__buffer[
buffer_first_index:buffer_last_index] = descriptors
if (self.RemoveBackgroundClusters):
Keypoint_buffer[buffer_first_index:buffer_last_index,
2] = pointsinbox
keypoint_indexes[names[i]] = [first_index, last_index]
first_index += len(keypoints)
buffer_first_index += len(keypoints)
#periodically we store the buffer in file
if buffer_last_index > int(buffersize * 0.8):
AppendFileArray(
np.array(Keypoint_buffer[:buffer_last_index]),
str(
GetCheckPointsPath(self.experiment_name, self.log_path)
/ 'keypoints'))
AppendFileArray(
np.array(Descriptor__buffer[:buffer_last_index]),
str(
GetCheckPointsPath(self.experiment_name, self.log_path)
/ 'descriptors'))
Keypoint_buffer = torch.zeros(buffersize, 3)
Descriptor__buffer = torch.zeros(buffersize, numberoffeatures)
buffer_first_index = 0
buffer_last_index = 0
LogText(f"Inference of Keypoints completed", self.experiment_name,
self.log_path)
#store any keypoints left on the buffers
AppendFileArray(
np.array(Keypoint_buffer[:buffer_last_index]),
str(
GetCheckPointsPath(self.experiment_name, self.log_path) /
'keypoints'))
AppendFileArray(
np.array(Descriptor__buffer[:buffer_last_index]),
str(
GetCheckPointsPath(self.experiment_name, self.log_path) /
'descriptors'))
#load handlers to the Keypoints and Descriptor files
Descriptors, fileHandler1 = OpenreadFileArray(
str(
GetCheckPointsPath(self.experiment_name, self.log_path) /
'descriptors'))
Keypoints, fileHandler2 = OpenreadFileArray(
str(
GetCheckPointsPath(self.experiment_name, self.log_path) /
'keypoints'))
Keypoints = Keypoints[:, :]
LogText(
f"Keypoints Detected per image {len(Keypoints)/len(keypoint_indexes)}",
self.experiment_name, self.log_path)
#perform outlier detection
inliersindexes = np.ones(len(Keypoints)) == 1
if (self.remove_superpoint_outliers_percentage > 0):
inliersindexes = self.Indexes_of_inliers(Keypoints, Descriptors,
buffersize)
#extend outliers with background points for constant background datasets
if (self.RemoveBackgroundClusters):
foregroundpointindex = self.Indexes_of_BackgroundPoints(
Keypoints, Descriptors, keypoint_indexes)
inliersindexes = np.logical_and(inliersindexes,
foregroundpointindex)
LogText(
f"Keypoints Detected per image(filtering) {sum(inliersindexes) / len(keypoint_indexes)}",
self.experiment_name, self.log_path)
#we use a subset of all the descriptors for clustering based on the recomendation of the Faiss repository
numberOfPointsForClustering = 500000
LogText(f"Clustering of keypoints", self.experiment_name,
self.log_path)
#clustering of superpoint features
KmeansClustering = clustering.Kmeans(self.number_of_clusters,
centroids=None)
descriptors = clustering.preprocess_features(
Descriptors[:numberOfPointsForClustering][
inliersindexes[:numberOfPointsForClustering]])
KmeansClustering.cluster(descriptors, verbose=False)
thresholds = self.GetThresholdsPerCluster(inliersindexes, Descriptors,
KmeansClustering)
Image_Keypoints = {}
averagepointsperimage = 0
for image in keypoint_indexes:
start, end = keypoint_indexes[image]
inliersinimage = inliersindexes[start:end]
keypoints = Keypoints[start:end, :]
inliersinimage[np.sum(keypoints[:, :2] < 0, 1) > 0] = False
inliersinimage[np.sum(keypoints[:, :2] > 64, 1) > 0] = False
keypoints = keypoints[inliersinimage]
image_descriptors = clustering.preprocess_features(
Descriptors[start:end])
image_descriptors = image_descriptors[inliersinimage]
#calculate distance of each keypoints to each centroid
distanceMatrix, clustering_assignments = KmeansClustering.index.search(
image_descriptors, self.number_of_clusters)
distanceMatrix = np.take_along_axis(
distanceMatrix, np.argsort(clustering_assignments), axis=-1)
#assign keypoints to centroids using the Hungarian algorithm. This ensures that each
#image has at most one instance of each cluster
keypointIndex, clusterAssignment = linear_sum_assignment(
distanceMatrix)
tempKeypoints = keypoints[keypointIndex]
clusterAssignmentDistance = distanceMatrix[keypointIndex,
clusterAssignment]
clusterstokeep = np.zeros(len(clusterAssignmentDistance))
clusterstokeep = clusterstokeep == 1
# keep only points that lie in their below a cluster specific theshold
clusterstokeep[clusterAssignmentDistance <
thresholds[clusterAssignment]] = True
tempKeypoints[:, 2] = clusterAssignment
Image_Keypoints[image] = tempKeypoints[clusterstokeep]
averagepointsperimage += sum(clusterstokeep)
LogText(
f"Keypoints Detected per image(clusteringAssignment) {averagepointsperimage / len(Image_Keypoints)}",
self.experiment_name, self.log_path)
ClosereadFileArray(
fileHandler1,
str(
GetCheckPointsPath(self.experiment_name, self.log_path) /
'keypoints'))
ClosereadFileArray(
fileHandler2,
str(
GetCheckPointsPath(self.experiment_name, self.log_path) /
'descriptors'))
self.save_keypoints(Image_Keypoints, "SuperPointKeypoints.pickle")
LogText(f"Extraction of Initial pseudoGroundtruth completed",
self.experiment_name, self.log_path)
return Image_Keypoints
def main():
step=2
experiment_options=Options()
global args
args = experiment_options.args
# config parameters
args = experiment_options.args
experiment_name=args.experiment_name
dataset_name = args.dataset_name
number_of_workers = args.num_workers
resume =args.resume
hyperparameters=experiment_options.GetHyperparameters(step,dataset_name)
# training parameters
iter_before_merging = hyperparameters.iter_before_merging
batchSize= hyperparameters.batchSize
weight_decay= hyperparameters.weight_decay
lr= hyperparameters.lr
batch_multiplier= hyperparameters.batch_multiplier
number_of_clusters= hyperparameters.number_of_clusters
totalIterations= hyperparameters.totalIterations
lrstep=hyperparameters.lrstep
confidence_thres_FAN=hyperparameters.confidence_thres_FAN
clusteroverlap=hyperparameters.clusteroverlap
#load paths
with open('paths/main.yaml') as file:
paths = yaml.load(file, Loader=yaml.FullLoader)
log_path=paths['log_path']
Utils.initialize_log_dirs(experiment_name,log_path)
LogText(f"Experiment Name {experiment_name}\n"
f"Database {dataset_name}\n"
"Training Parameters \n"
f"Batch size {batch_multiplier*batchSize} \n"
f"Learning rate {lr} \n"
f"Weight Decay {weight_decay} \n"
f"Iterations Before Mergins {iter_before_merging} \n"
f"Total Iterations {totalIterations} \n"
f"Number of Clusters {number_of_clusters} \n"
, experiment_name,log_path)
LogText("Training of Second step begins", experiment_name,log_path)
#augmentations for Second step
augmentations = iaa.Sequential([
iaa.Sometimes(0.5,
iaa.GaussianBlur(sigma=(0, 0.5))
),
iaa.ContrastNormalization((0.75, 1.5)),
iaa.AdditiveGaussianNoise(loc=0, scale=(0.0, 0.05 * 255), per_channel=0.5),
iaa.Multiply((0.8, 1.2), per_channel=0.2),
iaa.Sometimes(0.5,
iaa.MultiplyHueAndSaturation((0.5, 1.5), per_channel=True),
),
iaa.Affine(
scale={"x": (0.7, 1.3), "y": (0.7, 1.3)},
translate_percent={"x": (-0.05, 0.05), "y": (-0.05, 0.05)},
rotate=(-40, 40),
)
])
criterion = nn.MSELoss(reduce=False).cuda()
#model initialisation from weights of the first step
FAN = FAN_Model(number_of_clusters,criterion,experiment_name,confidence_thres_FAN,log_path,step)
if(resume):
FAN.init_secondstep(lr,weight_decay,batch_multiplier,number_of_clusters,lrstep,clusteroverlap)
path_to_checkpoint,path_to_keypoints=Utils.GetPathsResumeSecondStep(experiment_name,log_path)
FAN.load_trained_secondstep_model(path_to_checkpoint)
keypoints = Utils.load_keypoints(path_to_keypoints)
else:
path_to_checkpoint,path_to_keypoints=Utils.GetPathsTrainSecondStep(experiment_name,log_path)
FAN.init_secondstep(lr,weight_decay,batch_multiplier,number_of_clusters,lrstep,clusteroverlap,path_to_checkpoint=path_to_checkpoint)
keypoints = Utils.load_keypoints(path_to_keypoints)
FAN.RemoveSmallClusters(keypoints)
#initial dataloader
dataset = Database( dataset_name, FAN.number_of_clusters, image_keypoints=keypoints,
function_for_dataloading=Database.get_FAN_secondStep_train, augmentations=augmentations)
dataloader = DataLoader(dataset, batch_size=batchSize, shuffle=True, num_workers=number_of_workers,drop_last=True)
epochsbetweenMerging=0
while FAN.iterations < totalIterations:
FAN.Train_step2(dataloader)
#merging operation is performed initially after $iter_before_merging iterations and
#after that every epoch
if (FAN.iterations > iter_before_merging and epochsbetweenMerging==0):
LogText("Performing Cluster Merging", experiment_name,log_path)
#create dataloader for cluster merge operation
dataset_merging = Database(dataset_name, FAN.number_of_clusters,image_keypoints=keypoints,function_for_dataloading=Database.get_FAN_inference)
dataloader_merging = DataLoader(dataset_merging, batch_size=batchSize, shuffle=False, num_workers=number_of_workers, drop_last=False)
#form new set of keypoints with merged clusters
keypoints = FAN.MergeClusters(dataloader_merging, keypoints)
epochsbetweenMerging=1
#update training dataloader with new set of keypoints
dataset = Database( dataset_name, FAN.number_of_clusters, image_keypoints=keypoints,
function_for_dataloading=Database.get_FAN_secondStep_train, augmentations=augmentations)
dataloader = DataLoader(dataset, batch_size=batchSize, shuffle=True, num_workers=number_of_workers,
drop_last=True)
#merging is performed every 2 epochs (epochsbetweenMerging+1)
elif(FAN.iterations > iter_before_merging and epochsbetweenMerging>0):
epochsbetweenMerging=epochsbetweenMerging-1
def main():
step = 1
experiment_options = Options()
global args
# config parameters
args = experiment_options.args
experiment_name = args.experiment_name
dataset_name = args.dataset_name
number_of_workers = args.num_workers
resume = args.resume
hyperparameters = experiment_options.GetHyperparameters(step, dataset_name)
# training parameters
batchSize = hyperparameters.batchSize
weight_decay = hyperparameters.weight_decay
lr = hyperparameters.lr
number_of_clusters = hyperparameters.number_of_clusters
number_of_clustering_rounds = hyperparameters.number_of_clustering_rounds
nms_thres_superpoint = hyperparameters.nms_thres_superpoint
confidence_thres_superpoint = hyperparameters.confidence_thres_superpoint
use_box = hyperparameters.use_box
remove_superpoint_outliers_percentage = hyperparameters.remove_superpoint_outliers_percentage
training_iterations_before_first_clustering = hyperparameters.training_iterations_before_first_clustering
confidence_thres_FAN = hyperparameters.confidence_thres_FAN
UseScales = hyperparameters.UseScales
RemoveBackgroundClusters = hyperparameters.RemoveBackgroundClusters
#load paths
with open('paths/main.yaml') as file:
paths = yaml.load(file, Loader=yaml.FullLoader)
CheckPaths(paths, dataset_name)
log_path = paths['log_path']
path_to_superpoint_checkpoint = paths['path_to_superpoint_checkpoint']
#This funcion will create the directories /Logs and a /CheckPoints at log_path
Utils.initialize_log_dirs(experiment_name, log_path)
LogText(
f"Experiment Name {experiment_name}\n"
f"Database {dataset_name}\n"
"Training Parameters: \n"
f"Batch size {batchSize} \n"
f"Learning rate {lr} \n"
f"Weight Decay {weight_decay} \n"
f"Training iterations before first clustering {training_iterations_before_first_clustering} \n"
f"Number of clustering rounds {number_of_clustering_rounds} \n"
f"FAN detection threshold {confidence_thres_FAN} \n"
f"Number of Clusters {number_of_clusters} \n"
f"Outlier removal {remove_superpoint_outliers_percentage} \n",
experiment_name, log_path)
LogText("Training of First step begins", experiment_name, log_path)
#augmentations for first step
augmentations = iaa.Sequential([
iaa.Sometimes(0.3, iaa.GaussianBlur(sigma=(0, 0.5))),
iaa.ContrastNormalization((0.85, 1.3)),
iaa.Sometimes(
0.5,
iaa.AdditiveGaussianNoise(loc=0,
scale=(0.0, 0.05 * 255),
per_channel=0.5)),
iaa.Multiply((0.9, 1.1), per_channel=0.2),
iaa.Sometimes(
0.3,
iaa.MultiplyHueAndSaturation((0.5, 1.5), per_channel=True),
),
iaa.Affine(
scale={
"x": (0.8, 1.2),
"y": (0.8, 1.2)
},
translate_percent={
"x": (-0.05, 0.05),
"y": (-0.05, 0.05)
},
rotate=(-40, 40),
)
])
#selection of the dataloading function
superpoint_dataloading_function = Database.get_image_superpoint
if (UseScales):
superpoint_dataloading_function = Database.get_image_superpoint_multiple_scales
superpoint = SuperPoint(
number_of_clusters,
confidence_thres_superpoint,
nms_thres_superpoint,
path_to_superpoint_checkpoint,
experiment_name,
log_path,
remove_superpoint_outliers_percentage,
use_box,
UseScales,
RemoveBackgroundClusters,
)
superpoint_dataset = Database(
dataset_name,
number_of_clusters,
function_for_dataloading=superpoint_dataloading_function,
augmentations=augmentations,
use_box=use_box)
dataloader = DataLoader(superpoint_dataset,
batch_size=batchSize,
shuffle=False,
num_workers=number_of_workers,
drop_last=True)
criterion = nn.MSELoss().cuda()
FAN = FAN_Model(number_of_clusters, criterion, experiment_name,
confidence_thres_FAN, log_path, step)
FAN.init_firststep(lr, weight_decay, number_of_clusters,
training_iterations_before_first_clustering)
if (resume):
path_to_checkpoint, path_to_keypoints = Utils.GetPathsResumeFirstStep(
experiment_name, log_path)
if (path_to_checkpoint is not None):
FAN.load_trained_fiststep_model(path_to_checkpoint)
keypoints = Utils.load_keypoints(path_to_keypoints)
else:
#get initial pseudo-groundtruth by applying superpoint on the training data
keypoints = superpoint.CreateInitialPseudoGroundtruth(dataloader)
dataset = Database(
dataset_name,
FAN.number_of_clusters,
image_keypoints=keypoints,
function_for_dataloading=Database.get_FAN_firstStep_train,
augmentations=augmentations)
dataloader = DataLoader(dataset,
batch_size=batchSize,
shuffle=True,
num_workers=number_of_workers,
drop_last=True)
database_clustering = Database(
dataset_name,
FAN.number_of_clusters,
function_for_dataloading=Database.get_FAN_inference)
dataloader_clustering = DataLoader(database_clustering,
batch_size=batchSize,
shuffle=False,
num_workers=number_of_workers,
drop_last=True)
for i in range(number_of_clustering_rounds):
FAN.Train_step1(dataloader)
keypoints = FAN.Update_pseudoLabels(dataloader_clustering, keypoints)
dataset = Database(
dataset_name,
FAN.number_of_clusters,
image_keypoints=keypoints,
function_for_dataloading=Database.get_FAN_firstStep_train,
augmentations=augmentations)
dataloader = DataLoader(dataset,
batch_size=batchSize,
shuffle=True,
num_workers=number_of_workers,
drop_last=True)