Add resnet onnx model

README.md

@@ -1 +1,45 @@
-# image-classification-onnx
+# image-classification-onnx
+
+### Generate ONNX model
+```
+$ python convert_resnet_to_onnx.py
+```
+
+### Deploy with Cape
+
+First, you need to create a deployment folder containing your dependencies and a cape_handler in an `app.py` file.
+```
+# Create a deployment folder
+$ export TARGET="onnx_resnet_deploy"
+$ mkdir $TARGET
+# Add function script
+$ cp app.py $TARGET
+# Add ONNX resnet model 
+$ cp -r onnx_model $TARGET
+# Add onnxrumtime dependency.
+$ docker run -v `pwd`:/build -w /build --rm -it python:3.9-slim-bullseye pip install onnxruntime==1.13.1 --target /build/$TARGET
+```
+
+Then you can deploy your function using the Cape cli:
+```
+$ cape deploy onnx_resnet_deploy
+Deploying function to Cape ...
+Success! Deployed function to Cape.
+Function ID ➜  <FUNCTION_ID>
+Function Checksum ➜  <FUNCTION_CHECKSUM>
+$ export FUNCTION_ID=<copied from above>
+```
+
+### Run Secure Prediction
+
+Generate a personal access token for your account by running:
+```
+$ cape token create --name resnet
+Success! Your token: eyJhtGckO12...(token omitted)
+$ export TOKEN=<copied from above>
+```
+
+Invoke the image classification model with:
+```
+$ python run_prediction.py
+```

app.py

@@ -0,0 +1,12 @@
+import numpy as np
+import onnxruntime
+
+
+onnx_sess = onnxruntime.InferenceSession("./onnx_model/resnet50.onnx")
+
+
+def cape_handler(input_bytes: bytes):
+    input = np.frombuffer(input_bytes, dtype=np.float32).reshape(1, 3, 224, 224)
+    ort_inputs = {onnx_sess.get_inputs()[0].name: input}
+    ort_outs = onnx_sess.run(None, ort_inputs)
+    return ort_outs[0].tobytes()

convert_resnet_to_onnx.py

@@ -0,0 +1,35 @@
+import pathlib
+
+import torch
+import torchvision
+
+
+def export_model_to_onnx(model, dummy_tensor_shape, onnx_file_path):
+    model.eval()
+    dummy_tensor = torch.ones(
+        dummy_tensor_shape, dtype=torch.float32, requires_grad=True
+    )
+
+    torch.onnx.export(
+        model,
+        dummy_tensor,
+        onnx_file_path,
+        export_params=True,
+        input_names=["input"],
+        output_names=["output"],
+        dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}},
+    )
+
+
+if __name__ == "__main__":
+    onnx_path = pathlib.Path("./onnx_model")
+    if not onnx_path.exists():
+        onnx_path.mkdir()
+
+    weights = torchvision.models.ResNet50_Weights.DEFAULT
+    model = torchvision.models.resnet50(weights)
+    export_model_to_onnx(
+        model,
+        dummy_tensor_shape=(1, 3, 224, 224),
+        onnx_file_path=onnx_path / "resnet50.onnx",
+    )

requirements.txt

@@ -0,0 +1,5 @@
+--extra-index-url https://download.pytorch.org/whl/cpu
+onnxruntime==1.13.1
+torch
+torchvision
+pycape

run_prediction.py

@@ -0,0 +1,48 @@
+import os
+import warnings
+
+import numpy as np
+import torch
+from torchvision.io import read_image
+from torchvision.models import ResNet50_Weights
+
+from pycape import Cape
+
+warnings.filterwarnings("ignore")
+
+token_env = os.environ.get("TOKEN")
+function_id_env = os.environ.get("FUNCTION_ID")
+
+
+def process_image(file):
+    img = read_image(file)
+    weights = ResNet50_Weights.DEFAULT
+    preprocess = weights.transforms()
+    batch = preprocess(img).unsqueeze(0)
+    batch_numpy = batch.detach().numpy()
+    batch_numpy_bytes = batch_numpy.tobytes()
+    return batch_numpy_bytes
+
+
+def process_onnx_output(onnx_output):
+    onnx_prediction = torch.from_numpy(
+        np.frombuffer(onnx_output, dtype=np.float32).reshape(1, 1000)
+    )
+    prediction = onnx_prediction.squeeze(0).softmax(0)
+    class_id = prediction.argmax().item()
+    score = prediction[class_id].item()
+    category_name = ResNet50_Weights.DEFAULT.meta["categories"][class_id]
+    return category_name, score
+
+
+if __name__ == "__main__":
+    cape = Cape()
+    f = cape.function(function_id_env)
+    t = cape.token(token_env)
+
+    input_bytes = process_image("./images_sample/dog.jpeg")
+
+    onnx_output = cape.run(f, t, input_bytes)
+
+    category_name, score = process_onnx_output(onnx_output)
+    print(f"{category_name}: {100 * score:.1f}%")