|
import os |
|
import json |
|
import pyzipper |
|
from typing import Dict, Any, Optional |
|
from voyager import Index, Space, StorageDataType |
|
|
|
class DataManager: |
|
def __init__(self, faces_path: str = "data/faces.json", |
|
performers_zip: str = "data/persons.zip", |
|
facenet_index_path: str = "data/face_facenet.voy", |
|
arc_index_path: str = "data/face_arc.voy"): |
|
""" |
|
Initialize the data manager. |
|
|
|
Parameters: |
|
faces_path: Path to the faces.json file |
|
performers_zip: Path to the performers zip file |
|
facenet_index_path: Path to the facenet index file |
|
arc_index_path: Path to the arc index file |
|
""" |
|
self.faces_path = faces_path |
|
self.performers_zip = performers_zip |
|
self.facenet_index_path = facenet_index_path |
|
self.arc_index_path = arc_index_path |
|
|
|
|
|
self.index_arc = Index(Space.Cosine, num_dimensions=512, storage_data_type=StorageDataType.E4M3) |
|
self.index_facenet = Index(Space.Cosine, num_dimensions=512, storage_data_type=StorageDataType.E4M3) |
|
|
|
|
|
self.faces = {} |
|
self.performer_db = {} |
|
self.load_data() |
|
|
|
def load_data(self): |
|
"""Load all data from files""" |
|
self._load_faces() |
|
self._load_performer_db() |
|
self._load_indices() |
|
|
|
def _load_faces(self): |
|
"""Load faces from JSON file""" |
|
try: |
|
with open(self.faces_path, 'r') as f: |
|
self.faces = json.load(f) |
|
except Exception as e: |
|
print(f"Error loading faces: {e}") |
|
self.faces = {} |
|
|
|
def _load_performer_db(self): |
|
"""Load performer database from encrypted zip file""" |
|
try: |
|
with pyzipper.AESZipFile(self.performers_zip) as zf: |
|
password = os.getenv("VISAGE_KEY", "").encode('ascii') |
|
zf.setpassword(password) |
|
self.performer_db = json.loads(zf.read('performers.json')) |
|
except Exception as e: |
|
print(f"Error loading performer database: {e}") |
|
self.performer_db = {} |
|
|
|
def _load_indices(self): |
|
"""Load face recognition indices""" |
|
try: |
|
with open(self.arc_index_path, 'rb') as f: |
|
self.index_arc = self.index_arc.load(f) |
|
|
|
with open(self.facenet_index_path, 'rb') as f: |
|
self.index_facenet = self.index_facenet.load(f) |
|
except Exception as e: |
|
print(f"Error loading indices: {e}") |
|
|
|
def get_performer_info(self, stash_id: str, confidence: float) -> Optional[Dict[str, Any]]: |
|
""" |
|
Get performer information from the database |
|
|
|
Parameters: |
|
stash_id: Stash ID of the performer |
|
confidence: Confidence score (0-1) |
|
|
|
Returns: |
|
Dictionary with performer information or None if not found |
|
""" |
|
performer = self.performer_db.get(stash_id, []) |
|
if not performer: |
|
return None |
|
|
|
confidence_int = int(confidence * 100) |
|
return { |
|
'id': stash_id, |
|
"name": performer['name'], |
|
"confidence": confidence_int, |
|
'image': performer['image'], |
|
'country': performer['country'], |
|
'hits': 1, |
|
'distance': confidence_int, |
|
'performer_url': f"https://stashdb.org/performers/{stash_id}" |
|
} |
|
|
|
def query_facenet_index(self, embedding, limit): |
|
"""Query the facenet index with an embedding""" |
|
return self.index_facenet.query(embedding, limit) |
|
|
|
def query_arc_index(self, embedding, limit): |
|
"""Query the arc index with an embedding""" |
|
return self.index_arc.query(embedding, limit) |
