CelestialSurveyor

1
from cryptography.fernet import Fernet
2
from slow_fast_exp import SlowFast_body, bottleneck
3
from tensorflow.keras.models import Model
4

5

6
def build_model() -> Model:
7
    """
8
    Builds and returns a SlowFast model with specific configurations.
9
    You can update this function with your own model.
10

11
    Returns:
12
        Model: The built SlowFast model.
13
    """
14
    model = SlowFast_body([3, 4, 6, 3], bottleneck)
15
    return model
16

17

18
def encrypt_model(model_name: str, key: bytes = b'J17tdv3zz2nemLNwd17DV33-sQbo52vFzl2EOYgtScw=') -> None:
19
    """
20
    Encrypts the model weights and saves them to a binary file.
21
    Preliminary version of model decryption. It was done when I was thinking to make this project open source or not.
22

23
    Args:
24
        model_name: Name of the model file to be encrypted.
25
        key: Encryption key to be used. Default is a predefined key.
26

27
    Returns:
28
        None
29
    """
30
    with open(f"{model_name}.h5", "rb") as file:
31
        model_bytes = file.read()
32

33
    # Use the provided key to create a cipher
34
    cipher = Fernet(key)
35

36
    # Encrypt the entire model
37
    encrypted_model = cipher.encrypt(model_bytes)
38

39
    # Save the encrypted weights to a file
40
    with open(f"{model_name}.bin", "wb") as file:
41
        file.write(encrypted_model)
42

43

44
if __name__ == '__main__':
45
    build_model().summary()
46