diff --git a/vsms/multigrain.py b/vsms/multigrain.py index 36cf7e5..8af2dbf 100644 --- a/vsms/multigrain.py +++ b/vsms/multigrain.py @@ -160,7 +160,7 @@ def augment_score2(db,tup,qvec,vec_meta,vecs, rw_coarse=1.): ious = box_iou(tup, vec_meta) vec_meta = vec_meta.assign(iou=ious.reshape(-1)) max_boxes = vec_meta.groupby('zoom_level').iou.idxmax() - max_boxes = max_boxes.sort_index(ascending=True) # largest zoom level (zoomed out) goes last + max_boxes = max_boxes.sort_index(ascending=True) # largest zoom level (zoomed out) goes last relevant_meta = vec_meta.iloc[max_boxes] relevant_iou = relevant_meta.iou > 0 #there should be at least some overlap for it to be relevant max_boxes = max_boxes[relevant_iou.values] @@ -222,119 +222,35 @@ def get_pos_negs_all_v2(dbidxs, ds, vec_meta): negidxs = pr.BitMap(np.concatenate(neg)) return posidxs, negidxs - -class AugmentedDB(object): - """Structure holding a dataset, - together with precomputed embeddings - and (optionally) data structure - """ - def __init__(self, raw_dataset : torch.utils.data.Dataset, - embedding : XEmbedding, - embedded_dataset : np.ndarray, - vector_meta : pd.DataFrame): - self.raw = raw_dataset - self.embedding = embedding - self.embedded = embedded_dataset - self.vector_meta = vector_meta.assign(**get_boxes(vector_meta)) - - all_indices = pr.BitMap() - assert len(self.raw) == vector_meta.dbidx.unique().shape[0] - assert embedded_dataset.shape[0] == vector_meta.shape[0] - all_indices.add_range(0, len(self.raw)) - self.all_indices = pr.FrozenBitMap(all_indices) - - #norms = np.linalg.norm(embedded_dataset, axis=1)[:,np.newaxis] - # if not np.isclose(norms,0).all(): - # print('warning: embeddings are not normalized?', norms.max()) - - def __len__(self): - return len(self.raw) - - def _query_prelim(self, *, vector, topk, zoom_level, exclude=None): - if exclude is None: - exclude = pr.BitMap([]) - - included_dbidx = pr.BitMap(self.all_indices).difference(exclude) - vec_meta = self.vector_meta - vecs = self.embedded # = restrict_fine_grained(self.vector_meta, self.embedded, included_dbidx) - - # breakpoint() - vec_meta = vec_meta.reset_index(drop=True) - vec_meta = vec_meta[(~vec_meta.dbidx.isin(exclude))] - if zoom_level is not None: - vec_meta = vec_meta[vec_meta.zoom_level == zoom_level] - - vecs = vecs[vec_meta.index.values] - - if len(included_dbidx) == 0: - return np.array([]) - - if len(included_dbidx) <= topk: - topk = len(included_dbidx) - - scores = vecs @ vector.reshape(-1) - topscores = vec_meta.assign(score=scores).sort_values('score', ascending=False) - # return topk unique images. - seen_idx = set() - for i,idx in enumerate(topscores.dbidx.values): - seen_idx.add(idx) - if len(seen_idx) == topk: - break - - topscores = topscores.iloc[:i] - assert topscores.dbidx.unique().shape[0] <= topk - return topscores - - def query(self, *, vector, topk, mode='dot', exclude=None, shortlist_size=None, rel_weight_coarse=1): - if shortlist_size is None: - shortlist_size = topk*5 - - db = self - qvec=vector - scmeta = self._query_prelim(vector=qvec, topk=shortlist_size, zoom_level=None, exclude=exclude) - - scs = [] - for i,tup in enumerate(scmeta.itertuples()): - relmeta = db.vector_meta[db.vector_meta.dbidx == tup.dbidx] - relvecs = db.embedded[relmeta.index.values] - tup = scmeta.iloc[i:i+1] - sc = augment_score2(db, tup, qvec, relmeta, relvecs, rw_coarse=rel_weight_coarse) - scs.append(sc) - - final = scmeta.assign(aug_score=np.array(scs)) - agg = final.groupby('dbidx').aug_score.max().reset_index().sort_values('aug_score', ascending=False) - idxs = agg.dbidx.iloc[:topk] - return idxs - import math import annoy from annoy import AnnoyIndex -class IndexedDB(object): - """Structure holding a dataset, - together with precomputed embeddings - and (optionally) data structure +class AugmentedDB(object): + """implements a two stage lookup """ def __init__(self, raw_dataset : torch.utils.data.Dataset, embedding : XEmbedding, embedded_dataset : np.ndarray, vector_meta : pd.DataFrame, - index_path : str): + index_path : str = None): self.raw = raw_dataset self.embedding = embedding self.embedded = embedded_dataset - self.vector_meta = vector_meta.assign(**get_boxes(vector_meta)) - self.vec_index = AnnoyIndex(512, 'dot') - self.vec_index.load(index_path) + self.vector_meta = vector_meta + + if index_path is not None: + self.vec_index = AnnoyIndex(512, 'dot') + self.vec_index.load(index_path) + assert self.vec_index.n_items() == len(self.embedded) + else: + self.vec_index = None all_indices = pr.BitMap() assert len(self.raw) == vector_meta.dbidx.unique().shape[0] assert embedded_dataset.shape[0] == vector_meta.shape[0] all_indices.add_range(0, len(self.raw)) self.all_indices = pr.FrozenBitMap(all_indices) - - #norms = np.linalg.norm(embedded_dataset, axis=1)[:,np.newaxis] - # if not np.isclose(norms,0).all(): - # print('warning: embeddings are not normalized?', norms.max()) + def __len__(self): return len(self.raw) @@ -384,18 +300,22 @@ def get_nns_by_vector_exact(): scorepos = np.argsort(-scores) return scorepos, scores[scorepos] + if self.vec_index is not None: + idxs, scores = get_nns_by_vector_approx() + else: + idxs, scores = get_nns_by_vector_exact() - idxs, scores = get_nns_by_vector_approx() - topscores = vec_meta.iloc[idxs] + # work only with the two columns here bc dataframe can be large + topscores = vec_meta[['dbidx']].iloc[idxs] topscores = topscores.assign(score=scores) topscores = topscores[~topscores.dbidx.isin(exclude)] - bydbidx = topscores.groupby('dbidx').score.max() - bydbidx = bydbidx.sort_values(ascending=False) - candidates = pr.BitMap(bydbidx.index.values[:topk]) - topscores = topscores[topscores.dbidx.isin(candidates)] - assert topscores.dbidx.unique().shape[0] == topk + scoresbydbidx = topscores.groupby('dbidx').score.max().sort_values(ascending=False) + score_cutoff = scoresbydbidx.iloc[topk-1] # kth largest score + topscores = topscores[topscores.score >= score_cutoff] + candidates = pr.BitMap(topscores.dbidx) + assert len(candidates) >= topk assert candidates.intersection_cardinality(exclude) == 0 - return topscores, candidates + return topscores.index.values, candidates def query(self, *, vector, topk, mode='dot', exclude=None, shortlist_size=None, rel_weight_coarse=1): if shortlist_size is None: @@ -403,18 +323,25 @@ def query(self, *, vector, topk, mode='dot', exclude=None, shortlist_size=None, db = self qvec=vector - scmeta, candidate_id = self._query_prelim(vector=qvec, topk=shortlist_size, zoom_level=None, exclude=exclude) + meta_idx, candidate_id = self._query_prelim(vector=qvec, topk=shortlist_size, zoom_level=None, exclude=exclude) + + fullmeta = self.vector_meta[self.vector_meta.dbidx.isin(candidate_id)] + fullmeta = fullmeta.assign(**get_boxes(fullmeta)) + scmeta = self.vector_meta.iloc[meta_idx] + scmeta = scmeta.assign(**get_boxes(scmeta)) nframes = len(candidate_id) dbidxs = np.zeros(nframes)*-1 dbscores = np.zeros(nframes) - for i,(dbidx,gp) in enumerate(scmeta.groupby('dbidx')): + + ## for each frame, compute augmented scores for each tile and record max + for i,(dbidx,frame_vec_meta) in enumerate(scmeta.groupby('dbidx')): dbidxs[i] = dbidx - relmeta = db.vector_meta[db.vector_meta.dbidx == dbidx] + relmeta = fullmeta[fullmeta.dbidx == dbidx] # get metadata for all boxes in frame. relvecs = db.embedded[relmeta.index.values] - boxscs = np.zeros(gp.shape[0]) - for j in range(gp.shape[0]): - tup = gp.iloc[j:j+1] + boxscs = np.zeros(frame_vec_meta.shape[0]) + for j in range(frame_vec_meta.shape[0]): + tup = frame_vec_meta.iloc[j:j+1] boxscs[j] = augment_score2(db, tup, qvec, relmeta, relvecs, rw_coarse=rel_weight_coarse) dbscores[i] = np.max(boxscs) diff --git a/vsms/vloop_dataset_loaders.py b/vsms/vloop_dataset_loaders.py index 09355b0..4be5eb8 100644 --- a/vsms/vloop_dataset_loaders.py +++ b/vsms/vloop_dataset_loaders.py @@ -62,7 +62,7 @@ def get_qgt(box_data, min_box_size): def make_evdataset(*, root, paths, embedded_dataset, query_ground_truth, box_data, embedding, fine_grained_embedding=None, fine_grained_meta=None, min_box_size=0.): - query_ground_truth = make_qgt(box_data, min_box_size) # boxes are still there + query_ground_truth = get_qgt(box_data, min_box_size) # boxes are still there # for use by algorithms but not used for accuracy metrics #query_ground_truth = query_ground_truth.clip(0.,1.) root = os.path.abspath(root) diff --git a/vsms/vls_benchmark_tools.py b/vsms/vls_benchmark_tools.py index 3b329ca..0e6af4d 100644 --- a/vsms/vls_benchmark_tools.py +++ b/vsms/vls_benchmark_tools.py @@ -211,17 +211,11 @@ def run_loop6(*, ev :EvDataset, category, qstr, interactive, warm_start, n_batch hdb = FineEmbeddingDB(raw_dataset=ev.image_dataset, embedding=ev.embedding, embedded_dataset=vecs, vector_meta=vec_meta) elif granularity == 'multi': - # assert False, 'not using this at the moment' vec_meta = ev.fine_grained_meta vecs = ev.fine_grained_embedding - # dbidxs = np.arange(len(ev)).astype('int') - # vec_meta_coarse = pd.DataFrame({'iis': np.zeros_like(dbidxs), 'jjs':np.zeros_like(dbidxs), 'dbidx':dbidxs}) - # vec_meta_coarse = vec_meta_coarse.assign(scale='coarse') - # vecs_coarse = ev.embedded_dataset - # vec_meta = pd.concat([vec_meta_fine, vec_meta_coarse], ignore_index=True) - # vecs = np.concatenate([vecs_fine, vecs_coarse]) - index_path = './data/bdd_10k_allgrains_index.ann' - hdb = IndexedDB(raw_dataset=ev.image_dataset, embedding=ev.embedding, + #index_path = './data/bdd_10k_allgrains_index.ann' + index_path = None + hdb = AugmentedDB(raw_dataset=ev.image_dataset, embedding=ev.embedding, embedded_dataset=vecs, vector_meta=vec_meta, index_path=index_path) elif granularity == 'coarse': @@ -268,7 +262,7 @@ def run_loop6(*, ev :EvDataset, category, qstr, interactive, warm_start, n_batch for i in tqdm(range(n_batches), leave=False, disable=tqdm_disabled): idxbatch = bfq.query_stateful(mode=tmode, vector=tvec, batch_size=batch_size) - acc_indices.append(idxbatch) + acc_indices.append(idxbatch.copy()) # trying to figure out leak acc_results.append(gt[idxbatch]) if interactive != 'plain':