generate_dev_files.py

"A script to generate all development files necessary for the image filtering demo."

import shutil
import onnx
from common import AVAILABLE_FILTERS, FILTERS_PATH
from custom_client_server import CustomFHEDev
from filters import Filter

print("Generating deployment files for all available filters")

for image_filter in AVAILABLE_FILTERS:
    print("Filter:", image_filter, "\n")

    # Create the filter instance
    filter = Filter(image_filter)

    # Compile the model on a representative inputset
    filter.compile()

    # Define the directory path associated to this filter
    filter_path = FILTERS_PATH / image_filter

    # Define the directory path associated to this filter's deployment files
    deployment_path = filter_path / "deployment"

    # Delete the deployment folder and its content if it already exists
    if deployment_path.is_dir():
        shutil.rmtree(deployment_path)

    # Save the files needed for deployment
    fhe_dev_filter = CustomFHEDev(deployment_path, filter)
    fhe_dev_filter.save()

    # Save the ONNX model
    onnx.save(filter.onnx_model, filter_path / "server.onnx")

print("Done !")