Spaces:
Sleeping
Sleeping
File size: 3,577 Bytes
9282c17 bdd2b03 9282c17 bdd2b03 9282c17 | 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 87 88 89 90 91 92 93 94 95 96 97 98 | import os
import zvec
from typing import Dict, Any, Optional, List, Tuple
class DataManager:
def __init__(self, collection_path: str = "data/performers.zvec"):
self.collection_path = collection_path
self._metadata_cache: Dict[str, Dict] = {}
self.collection = None
self._load_collection()
def _load_collection(self):
try:
self.collection = zvec.open(
path=self.collection_path,
option=zvec.CollectionOption(read_only=True, enable_mmap=True),
)
print(f"Collection loaded OK from {self.collection_path}")
except Exception as e:
import traceback
print(f"FATAL: failed to load collection from {self.collection_path}: {e}")
traceback.print_exc()
def get_performer_info(self, stash_id: str, confidence: float, distance: float = 0.0) -> Optional[Dict[str, Any]]:
meta = self._metadata_cache.get(stash_id)
if not meta:
return None
confidence_int = int(confidence * 100)
return {
'id': stash_id,
'name': meta.get("name", ""),
'confidence': confidence_int,
'distance': round(distance, 4),
'image': meta.get("image", ""),
'country': meta.get("country") or None,
'hits': 1,
'performer_url': f"https://stashdb.org/performers/{stash_id}"
}
def _query(self, field_name: str, embedding, limit: int) -> Tuple[List[str], List[float]]:
if self.collection is None:
return [], []
results = self.collection.query(
vectors=zvec.VectorQuery(field_name=field_name, vector=embedding.tolist()),
topk=limit,
)
ids = []
distances = []
for doc in results:
doc_id = doc.id if hasattr(doc, 'id') else doc['id']
doc_score = doc.score if hasattr(doc, 'score') else doc['score']
doc_fields = doc.fields if hasattr(doc, 'fields') else doc.get('fields', {})
ids.append(doc_id)
distances.append(doc_score)
self._metadata_cache[doc_id] = doc_fields
return ids, distances
def query_facenet_index(self, embedding, limit: int) -> Tuple[List[str], List[float]]:
return self._query("facenet", embedding, limit)
def query_arc_index(self, embedding, limit: int) -> Tuple[List[str], List[float]]:
return self._query("arc", embedding, limit)
def query_multi(self, facenet_emb, arc_emb, limit: int) -> Tuple[List[str], List[float]]:
if self.collection is None:
return [], []
results = self.collection.query(
vectors=[
zvec.VectorQuery(field_name="facenet", vector=facenet_emb.tolist()),
zvec.VectorQuery(field_name="arc", vector=arc_emb.tolist()),
],
topk=limit,
reranker=zvec.WeightedReRanker(
topn=limit,
metric=zvec.MetricType.COSINE,
weights={"facenet": 1.0, "arc": 1.0},
),
)
ids = []
scores = []
for doc in results:
doc_id = doc.id if hasattr(doc, 'id') else doc['id']
doc_score = doc.score if hasattr(doc, 'score') else doc['score']
doc_fields = doc.fields if hasattr(doc, 'fields') else doc.get('fields', {})
ids.append(doc_id)
scores.append(doc_score)
self._metadata_cache[doc_id] = doc_fields
return ids, scores
|