# Yolov5 weight encryp
### 適用pytorch model

Step1. 加密套件準備，使用cryptography的Fernet進行加密。

pip install cryptography

Step2. 生成Key，可使用generate_key()函式，也可自定義。

from cryptography.fernet import Fernet

key = Fernet.generate_key()

Step2-1. 選擇性把key保存起來。

with open('key', 'wb') as f:
    f.write(key)

Step2-2. key讀出來。

with open('key', 'rb') as f:
    key = f.read()

Step3. 加密weight。

import io
import torch
from cryptography.fernet import Fernet

model_weight = torch.load(weight路徑)
b = io.BytesIO()
torch.save(model_weight, b)
b.seek(0)
weight_bytes = b.read()
encrypted_weight_bytes = Fernet(key).encrypt(weight_bytes)

with open(加密後weight保存路徑, 'wb') as f:
    f.write(encrypted_weight_bytes)

Step4. 解密weight。

在yolo v5的程式碼中，models資料夾內experimental.py為load pytorch model程式庫，
其中函式attempt_load是這邊要改寫的地方。
原程式碼：

from models.yolo import Detect, Model

# Loads an ensemble of models weights=[a,b,c] or a single model weights=[a] or weights=a
model = Ensemble()
for w in weights if isinstance(weights, list) else [weights]:
    ckpt = torch.load(attempt_download(w), map_location=map_location)  # load
    if fuse:
        model.append(ckpt['ema' if ckpt.get('ema') else 'model'].float().fuse().eval())  # FP32 model
    else:
        model.append(ckpt['ema' if ckpt.get('ema') else 'model'].float().eval())  # without layer fuse

改寫後程式碼：

from models.yolo import Detect, Model
-from cryptography.fernet import Fernet
-import io

# Loads an ensemble of models weights=[a,b,c] or a single model weights=[a] or weights=a
model = Ensemble()
for w in weights if isinstance(weights, list) else [weights]:
-   with open(w, 'rb') as f:
-       encrypted_data = f.read()
-    decrypted_data = Fernet(key).decrypt(encrypted_data) # key可讀取也可貼byte，例如：b'RtqikMWYdz45HAuXx-dguw6lmNXgOyT0XIMV8r0TtvU='
-    b = io.BytesIO(decrypted_data)
-    b.seek(0)
-    ckpt = torch.load(b, map_location=map_location)  # load
    if fuse:
        model.append(ckpt['ema' if ckpt.get('ema') else 'model'].float().fuse().eval())  # FP32 model
    else:
        model.append(ckpt['ema' if ckpt.get('ema') else 'model'].float().eval())  # without layer fuse

5. 測試如果key錯誤，錯誤訊息如下：

cryptography.fernet.InvalidToken