HyDiaPy/client.py at master · butvinm/HyDiaPy · GitHub

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
import itertools
from dataclasses import dataclass
from pathlib import Path

import openfhe
from face_match import norm_face_embedding

EMBEDDING_DIM = 128

SIMILARITY_THRESHOLD = 0.9


@dataclass
class ClientParameters:
    """Public Client parameters shared with other parties."""

    cc: openfhe.CryptoContext
    pk: openfhe.PublicKey
    ring_dim: int
    batch_size: int
    embedding_dim: int
    similarity_threshold: float


class Client:
    def __init__(self) -> None:
        self.cc, self.key_pair, self.ring_dim, self.batch_size = self._init_crypto_context()

    def _init_crypto_context(self) -> tuple[openfhe.CryptoContext, openfhe.KeyPair, int, int]:
        parameters = openfhe.CCParamsCKKSRNS()
        parameters.SetSecurityLevel(openfhe.SecurityLevel.HEStd_128_classic)
        parameters.SetMultiplicativeDepth(11)
        parameters.SetScalingModSize(45)
        parameters.SetScalingTechnique(openfhe.ScalingTechnique.FIXEDMANUAL)

        cc = openfhe.GenCryptoContext(parameters)
        cc.Enable(openfhe.PKESchemeFeature.PKE)
        cc.Enable(openfhe.PKESchemeFeature.KEYSWITCH)
        cc.Enable(openfhe.PKESchemeFeature.LEVELEDSHE)
        cc.Enable(openfhe.PKESchemeFeature.ADVANCEDSHE)

        key_pair = cc.KeyGen()
        cc.EvalMultKeyGen(key_pair.secretKey)
        cc.EvalRotateKeyGen(key_pair.secretKey, range(EMBEDDING_DIM))
        cc.EvalSumKeyGen(key_pair.secretKey)

        ring_dim = cc.GetRingDimension()
        batch_size = cc.GetBatchSize()

        return cc, key_pair, ring_dim, batch_size

    def setup(self) -> ClientParameters:
        return ClientParameters(
            cc=self.cc,
            pk=self.key_pair.publicKey,
            ring_dim=self.ring_dim,
            batch_size=self.batch_size,
            embedding_dim=EMBEDDING_DIM,
            similarity_threshold=SIMILARITY_THRESHOLD,
        )

    def query(
        self,
        query_image: Path,
    ) -> openfhe.Ciphertext:
        face_embedding = norm_face_embedding(query_image)
        partitions_count = self.batch_size // EMBEDDING_DIM
        query_aligned = list(itertools.chain(*itertools.repeat(face_embedding, partitions_count)))
        query_pt = self.cc.MakeCKKSPackedPlaintext(query_aligned)
        query_ct = self.cc.Encrypt(self.key_pair.publicKey, query_pt)
        return query_ct

    def extract_identities(
        self,
        labels: list[str],
        thresholds: list[openfhe.Ciphertext],
    ) -> list[str]:
        identities = []
        i = 0
        for batch in thresholds:
            thresholds_pt = self.cc.Decrypt(self.key_pair.secretKey, batch).GetCKKSPackedValue()
            for threshold in thresholds_pt:
                if abs(threshold) > 1:
                    identities.append(labels[i])

                i += 1

        return identities

    def extract_membership(
        self,
        thresholds: openfhe.Ciphertext,
    ) -> bool:
        thresholds_pt = self.cc.Decrypt(self.key_pair.secretKey, thresholds).GetCKKSPackedValue()
        return abs(thresholds_pt[0]) > 1