Spaces:
Running
on
Zero
Running
on
Zero
File size: 3,107 Bytes
e6f20b8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 |
"""Kernel functionality implementation - aggregation and similarity computation"""
import numpy as np
from . import functional, hamming, io_helpers
class ASMKKernel:
"""Kernel for ASMK with the option of binarization."""
binary_shortcuts = {"bin": True, "nobin": False}
def __init__(self, codebook, *, binary):
self.params = {
"binary": binary,
}
self.binary = self.binary_shortcuts.get(binary, binary)
assert self.binary in self.binary_shortcuts.values()
self.codebook = codebook
#
# Aggregation
#
def aggregate_image(self, des, word_ids):
"""Aggregate descriptors (with corresponding visual word ids) for a single image"""
unique_ids = np.unique(word_ids)
ades = np.empty((unique_ids.shape[0], des.shape[1]), dtype=np.float32)
for i, word in enumerate(unique_ids):
ades[i] = (des[(word_ids==word).any(axis=1)] - self.codebook.centroids[word]).sum(0)
if self.binary:
ades = hamming.binarize_and_pack_2D(ades)
else:
ades = functional.normalize_vec_l2(ades)
return ades, unique_ids
def aggregate(self, des, word_ids, image_ids, *, progress=None, **kwargs):
"""Aggregate descriptors with corresponding visual word ids for corresponding image ids"""
acc = []
slices = list(io_helpers.slice_unique(image_ids))
for imid, seq in io_helpers.progress(slices, frequency=progress, header="Aggregate"):
ades, ids = self.aggregate_image(des[seq], word_ids[seq], **kwargs)
acc.append((ades, ids, np.full(ids.shape[0], imid)))
agg_des, agg_words, agg_imids = zip(*acc)
return np.vstack(agg_des), np.hstack(agg_words), np.hstack(agg_imids)
#
# Similarity
#
def similarity(self, qvec, vecs, image_ids, *, alpha, similarity_threshold):
"""Compute similarity between given query vector and database feature vectors with their
corresponding image ids. Alpha is the similarity exponent after the similarity
threshold is applied."""
# Compute similarity with vw residuals for all other images
if self.binary:
norm_hdist = hamming.hamming_cdist_packed(qvec.reshape(1, -1), vecs)
sim = -2*norm_hdist.squeeze(0) + 1 # normalized hamming dist -> similarity in [-1, 1]
else:
sim = np.matmul(vecs, qvec)
return functional.asmk_kernel(sim, image_ids, alpha=alpha,
similarity_threshold=similarity_threshold)
#
# Load and save
#
def state_dict(self):
"""Return state dict which is a checkpoint of current state for future recovery"""
return {
"type": self.__class__.__name__,
"params": self.params,
}
@classmethod
def initialize_from_state(cls, state, codebook):
"""Initialize from a previously stored state_dict given a codebook"""
assert state["type"] == cls.__name__
return cls(**state["params"], codebook=codebook)
|