File size: 3,457 Bytes
21c7197
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
3cf0931
 
21c7197
 
3cf0931
 
21c7197
 
 
 
 
 
 
 
 
 
 
 
3cf0931
 
21c7197
 
 
3cf0931
21c7197
 
 
 
 
 
3cf0931
 
21c7197
 
 
3cf0931
21c7197
 
3cf0931
21c7197
 
3cf0931
21c7197
 
 
 
 
 
 
 
 
 
 
3cf0931
 
21c7197
 
3cf0931
21c7197
 
 
 
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
"""Server that will listen for GET and POST requests from the client."""

import time
from typing import List

from common import FILTERS_PATH, SERVER_TMP_PATH
from custom_client_server import CustomFHEServer
from fastapi import FastAPI, File, Form, UploadFile
from fastapi.responses import JSONResponse, Response
from pydantic import BaseModel


def get_server_file_path(name, user_id, image_filter):
    """Get the correct temporary file path for the server.

    Args:
        name (str): The desired file name.
        user_id (int): The current user's ID.
        image_filter (str): The filter chosen by the user

    Returns:
        pathlib.Path: The file path.
    """
    return SERVER_TMP_PATH / f"{name}_{image_filter}_{user_id}"


class FilterRequest(BaseModel):
    filter: str


# Initialize an instance of FastAPI
app = FastAPI()

# Define the default route
@app.get("/")
def root():
    return {"message": "Welcome to Your Image FHE Filter Server!"}


@app.post("/send_input")
def send_input(
    user_id: str = Form(),
    filter: str = Form(),
    files: List[UploadFile] = File(),
):
    """Send the inputs to the server."""
    # Retrieve the encrypted input image and the evaluation key paths
    encrypted_image_path = get_server_file_path("encrypted_image", filter, user_id)
    evaluation_key_path = get_server_file_path("evaluation_key", filter, user_id)
    
    # Write the files using the above paths
    with encrypted_image_path.open("wb") as encrypted_image, evaluation_key_path.open(
        "wb"
    ) as evaluation_key:
        encrypted_image.write(files[0].file.read())
        evaluation_key.write(files[1].file.read())


@app.post("/run_fhe")
def run_fhe(
    user_id: str = Form(),
    filter: str = Form(),
):
    """Execute the filter on the encrypted input image using FHE."""
    # Retrieve the encrypted input image and the evaluation key paths
    encrypted_image_path = get_server_file_path("encrypted_image", filter, user_id)
    evaluation_key_path = get_server_file_path("evaluation_key", filter, user_id)

    # Read the files using the above paths
    with encrypted_image_path.open("rb") as encrypted_image_file, evaluation_key_path.open(
        "rb"
    ) as evaluation_key_file:
        encrypted_image = encrypted_image_file.read()
        evaluation_key = evaluation_key_file.read()

    # Load the FHE server
    fhe_server = CustomFHEServer(FILTERS_PATH / f"{filter}/deployment")

    # Run the FHE execution
    start = time.time()
    encrypted_output_image = fhe_server.run(encrypted_image, evaluation_key)
    fhe_execution_time = round(time.time() - start, 2)

    # Retrieve the encrypted output image path
    encrypted_output_path = get_server_file_path("encrypted_output", filter, user_id)

    # Write the file using the above path
    with encrypted_output_path.open("wb") as encrypted_output:
        encrypted_output.write(encrypted_output_image)

    return JSONResponse(content=fhe_execution_time)


@app.post("/get_output")
def get_output(
    user_id: str = Form(),
    filter: str = Form(),
):
    """Retrieve the encrypted output image."""
    # Retrieve the encrypted output image path
    encrypted_output_path = get_server_file_path("encrypted_output", filter, user_id)

    # Read the file using the above path
    with encrypted_output_path.open("rb") as encrypted_output_file:
        encrypted_output = encrypted_output_file.read()

    return Response(encrypted_output)