File size: 3,684 Bytes
aac350d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
99
100
101
102
103
104
105
106
107
108
109
"""
Reference store — known-faces gallery.

Stores per-person face embeddings (numpy arrays) on disk plus a JSON
manifest mapping person_name -> list of embedding filenames.  Used by
recognition providers as the "known faces" database.

This is the ONLY place gallery mutations happen.  Recognition providers
read from here; they never write.
"""

from __future__ import annotations

import json
import threading
from pathlib import Path
from typing import Dict, List

import numpy as np
from loguru import logger


class ReferenceStore:
    """Disk-backed known-faces gallery."""

    def __init__(self, root: Path) -> None:
        self._root = Path(root)
        self._root.mkdir(parents=True, exist_ok=True)
        self._manifest_path = self._root / "manifest.json"
        self._lock = threading.RLock()
        self._manifest: Dict[str, List[str]] = {}
        self._load()

    # ------------------------------------------------------------------ #
    # Manifest persistence
    # ------------------------------------------------------------------ #
    def _load(self) -> None:
        if self._manifest_path.exists():
            try:
                self._manifest = json.loads(self._manifest_path.read_text())
            except json.JSONDecodeError:
                logger.warning("Gallery manifest corrupted, starting fresh")
                self._manifest = {}
        else:
            self._manifest = {}

    def _save(self) -> None:
        self._manifest_path.write_text(json.dumps(self._manifest, indent=2))

    # ------------------------------------------------------------------ #
    # CRUD
    # ------------------------------------------------------------------ #
    def add(self, name: str, embedding: np.ndarray) -> None:
        """Add an embedding for a person."""
        name = name.strip().lower().replace(" ", "_")
        with self._lock:
            if name not in self._manifest:
                self._manifest[name] = []
            idx = len(self._manifest[name])
            emb_path = self._root / f"{name}_{idx}.npy"
            np.save(emb_path, embedding.astype(np.float32))
            self._manifest[name].append(emb_path.name)
            self._save()

    def remove(self, name: str) -> bool:
        name = name.strip().lower().replace(" ", "_")
        with self._lock:
            if name not in self._manifest:
                return False
            for fname in self._manifest[name]:
                (self._root / fname).unlink(missing_ok=True)
            del self._manifest[name]
            self._save()
            return True

    def clear(self) -> int:
        with self._lock:
            n = len(self._manifest)
            for name, files in self._manifest.items():
                for f in files:
                    (self._root / f).unlink(missing_ok=True)
            self._manifest = {}
            self._save()
            return n

    def list_persons(self) -> List[dict]:
        with self._lock:
            return [
                {"name": name, "num_embeddings": len(files)}
                for name, files in self._manifest.items()
            ]

    def get_embeddings(self, name: str) -> List[np.ndarray]:
        name = name.strip().lower().replace(" ", "_")
        with self._lock:
            files = self._manifest.get(name, [])
            return [np.load(self._root / f) for f in files]

    def get_all(self) -> Dict[str, List[np.ndarray]]:
        with self._lock:
            return {
                name: [np.load(self._root / f) for f in files]
                for name, files in self._manifest.items()
            }

    @property
    def root(self) -> Path:
        return self._root