File size: 2,887 Bytes
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 |
"""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(),
):
encrypted_image_path = get_server_file_path("encrypted_image", filter, user_id)
evaluation_key_path = get_server_file_path("evaluation_key", filter, user_id)
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(),
):
encrypted_image_path = get_server_file_path("encrypted_image", filter, user_id)
evaluation_key_path = get_server_file_path("evaluation_key", filter, user_id)
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 model
fhe_model = CustomFHEServer(FILTERS_PATH / f"{filter}/deployment")
# Run the FHE execution
start = time.time()
encrypted_output_image = fhe_model.run(encrypted_image, evaluation_key)
fhe_execution_time = round(time.time() - start, 2)
encrypted_output_path = get_server_file_path("encrypted_output", filter, user_id)
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(),
):
encrypted_output_path = get_server_file_path("encrypted_output", filter, user_id)
with encrypted_output_path.open("rb") as encrypted_output_file:
encrypted_output = encrypted_output_file.read()
return Response(encrypted_output)
|