Avinash-05/FHE

1"Client-server interface custom implementation for filter models."2 3from concrete import fhe4 5from filters import Filter6 7 8class FHEServer:9    """Server interface run a FHE circuit."""10 11    def __init__(self, path_dir):12        """Initialize the FHE interface.13 14        Args:15            path_dir (Path): The path to the directory where the circuit is saved.16        """17        self.path_dir = path_dir18 19        # Load the FHE circuit20        self.server = fhe.Server.load(self.path_dir / "server.zip")21 22    def run(self, serialized_encrypted_image, serialized_evaluation_keys):23        """Run the filter on the server over an encrypted image.24 25        Args:26            serialized_encrypted_image (bytes): The encrypted and serialized image.27            serialized_evaluation_keys (bytes): The serialized evaluation keys.28 29        Returns:30            bytes: The filter's output.31        """32        # Deserialize the encrypted input image and the evaluation keys33        encrypted_image = fhe.Value.deserialize(serialized_encrypted_image)34        evaluation_keys = fhe.EvaluationKeys.deserialize(serialized_evaluation_keys)35 36        # Execute the filter in FHE37        encrypted_output = self.server.run(encrypted_image, evaluation_keys=evaluation_keys)38 39        # Serialize the encrypted output image40        serialized_encrypted_output = encrypted_output.serialize()41 42        return serialized_encrypted_output43 44 45class FHEDev:46    """Development interface to save and load the filter."""47 48    def __init__(self, filter, path_dir):49        """Initialize the FHE interface.50 51        Args:52            filter (Filter): The filter to use in the FHE interface.53            path_dir (str): The path to the directory where the circuit is saved.54        """55 56        self.filter = filter57        self.path_dir = path_dir58 59        self.path_dir.mkdir(parents=True, exist_ok=True)60 61    def save(self):62        """Export all needed artifacts for the client and server interfaces."""63        64        assert self.filter.fhe_circuit is not None, (65            "The model must be compiled before saving it."66        )67 68        # Save the circuit for the server, using the via_mlir in order to handle cross-platform69        # execution70        path_circuit_server = self.path_dir / "server.zip"71        self.filter.fhe_circuit.server.save(path_circuit_server, via_mlir=True)72 73        # Save the circuit for the client74        path_circuit_client = self.path_dir / "client.zip"75        self.filter.fhe_circuit.client.save(path_circuit_client)76 77 78class FHEClient:79    """Client interface to encrypt and decrypt FHE data associated to a Filter."""80 81    def __init__(self, path_dir, filter_name, key_dir=None):82        """Initialize the FHE interface.83 84        Args:85            path_dir (Path): The path to the directory where the circuit is saved.86            filter_name (str): The filter's name to consider.87            key_dir (Path): The path to the directory where the keys are stored. Default to None.88        """89        self.path_dir = path_dir90        self.key_dir = key_dir91 92        # If path_dir does not exist raise93        assert path_dir.exists(), f"{path_dir} does not exist. Please specify a valid path."94 95        # Load the client96        self.client = fhe.Client.load(self.path_dir / "client.zip", self.key_dir)97 98        # Instantiate the filter99        self.filter = Filter(filter_name)100 101    def generate_private_and_evaluation_keys(self, force=False):102        """Generate the private and evaluation keys.103 104        Args:105            force (bool): If True, regenerate the keys even if they already exist.106        """107        self.client.keygen(force)108 109    def get_serialized_evaluation_keys(self):110        """Get the serialized evaluation keys.111 112        Returns:113            bytes: The evaluation keys.114        """115        return self.client.evaluation_keys.serialize()116 117    def encrypt_serialize(self, input_image):118        """Encrypt and serialize the input image in the clear.119 120        Args:121            input_image (numpy.ndarray): The image to encrypt and serialize.122 123        Returns:124            bytes: The pre-processed, encrypted and serialized image.125        """126        # Encrypt the image127        encrypted_image = self.client.encrypt(input_image)128 129        # Serialize the encrypted image to be sent to the server130        serialized_encrypted_image = encrypted_image.serialize()131        return serialized_encrypted_image132 133    def deserialize_decrypt_post_process(self, serialized_encrypted_output_image):134        """Deserialize, decrypt and post-process the output image in the clear.135 136        Args:137            serialized_encrypted_output_image (bytes): The serialized and encrypted output image.138 139        Returns:140            numpy.ndarray: The decrypted, deserialized and post-processed image.141        """142        # Deserialize the encrypted image143        encrypted_output_image = fhe.Value.deserialize(144            serialized_encrypted_output_image145        )146 147        # Decrypt the image148        output_image = self.client.decrypt(encrypted_output_image)149 150        # Post-process the image151        post_processed_output_image = self.filter.post_processing(output_image)152 153        return post_processed_output_image154