Cape handler returns top 5 classes

README.md

@@ -16,6 +16,8 @@ $ mkdir $TARGET
 $ cp app.py $TARGET
 # Add ONNX resnet model 
 $ cp -r onnx_model $TARGET
+# Add imagenet classes file
+$ cp imagenet_classes.txt $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
 ```
@@ -42,4 +44,9 @@ $ export TOKEN=<copied from above>
 Invoke the image classification model with:
 ```
 $ python run_prediction.py
+golden retriever: 39.7%
+Labrador retriever: 7.1%
+tennis ball: 1.3%
+clumber: 0.9%
+Brittany spaniel: 0.7%
 ```

app.py

@@ -1,12 +1,32 @@
+import json
 import numpy as np
 import onnxruntime
 
 
 onnx_sess = onnxruntime.InferenceSession("./onnx_model/resnet50.onnx")
 
+with open("imagenet_classes.txt", "r") as f:
+    imagenet_classes = [s.strip() for s in f.readlines()]
+
+
+def softmax(x):
+    e_x = np.exp(x - np.max(x))
+    return e_x / e_x.sum()
+
+
+def get_top5_classes(onnx_output, imagenet_classes):
+    onnx_output = softmax(onnx_output.flatten())
+    top5_cat_id = np.argsort(-onnx_output)[:5]
+    top5_class_score = {
+        imagenet_classes[class_id]: onnx_output[class_id].tolist()
+        for class_id in top5_cat_id
+    }
+    return top5_class_score
+
 
 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()
+    top5_classes = get_top5_classes(ort_outs[0], imagenet_classes)
+    return json.dumps(top5_classes)

convert_resnet_to_onnx.py

@@ -26,10 +26,17 @@ def export_model_to_onnx(model, dummy_tensor_shape, onnx_file_path):
     if not onnx_path.exists():
         onnx_path.mkdir()
 
+    # Export model to ONNX
     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",
     )
+
+    # Export imagnet classes to a txt file
+    imagenet_classes = weights.meta["categories"]
+    with open("imagenet_classes.txt", "w") as f:
+        for class_name in imagenet_classes:
+            f.write(str(class_name) + "\n")