fix: Fixed resizing and switched to ONNX

app.py

@@ -1,39 +1,71 @@
+# Copyright (C) 2022, François-Guillaume Fernandez.
+
+# This program is licensed under the Apache License 2.0.
+# See LICENSE or go to <https://www.apache.org/licenses/LICENSE-2.0> for full license details.
+
 import argparse
+import json
 
 import gradio as gr
-import torch
+import numpy as np
+import onnxruntime
+from huggingface_hub import hf_hub_download
 from PIL import Image
-from torchvision.transforms import Compose, ConvertImageDtype, Normalize, PILToTensor, Resize
-from torchvision.transforms.functional import InterpolationMode
 
-from holocron import models
 
+REPO = "frgfm/rexnet1_0x"
+
+# Download model config & checkpoint
+with open(hf_hub_download(args.repo, filename="config.json"), "rb") as f:
+    cfg = json.load(f)
+
+ort_session = onnxruntime.InferenceSession(hf_hub_download(args.repo, filename="model.onnx"))
+
+def preprocess_image(pil_img: Image.Image) -> np.ndarray:
+    """Preprocess an image for inference
+
+    Args:
+        pil_img: a valid pillow image
+
+    Returns:
+        the resized and normalized image of shape (1, C, H, W)
+    """
+
+    # Resizing (PIL takes (W, H) order for resizing)
+    img = pil_img.resize(cfg["input_shape"][-2:][::-1], Image.BILINEAR)
+    # (H, W, C) --> (C, H, W)
+    img = np.asarray(img).transpose((2, 0, 1)).astype(np.float32) / 255
+    # Normalization
+    img -= np.array(cfg["mean"])[:, None, None]
+    img /= np.array(cfg["std"])[:, None, None]
+
+    return img[None, ...]
 
-model = models.rexnet1_3x(pretrained=True).eval()
+def predict(image):
+    # Preprocessing
+    np_img = preprocess_image(image)
+    ort_input = {ort_session.get_inputs()[0].name: np_img}
 
-preprocessor = Compose([
-    Resize(model.default_cfg['input_shape'][1:], interpolation=InterpolationMode.BILINEAR),
-    PILToTensor(),
-    ConvertImageDtype(torch.float32),
-    Normalize(model.default_cfg['mean'], model.default_cfg['std'])
-])
+    # Inference
+    ort_out = ort_session.run(None, ort_input)
+    # Post-processing
+    out_exp = np.exp(ort_out[0][0])
+    probs = out_exp / out_exp.sum()
 
-def predict(img):
-    img = Image.fromarray(img.astype('uint8'), 'RGB')
-    img = preprocessor(img)
-    with torch.inference_mode():
-        prediction = torch.nn.functional.softmax(model(img.unsqueeze(0))[0], dim=0)
-    return {class_name: float(conf) for class_name, conf in zip(model.default_cfg['classes'], prediction)}
+    return {class_name: float(conf) for class_name, conf in zip(cfg["classes"], probs)}
 
-image = gr.inputs.Image()
+img = gr.inputs.Image(type="pil")
 outputs = gr.outputs.Label(num_top_classes=3)
 
 gr.Interface(
     fn=predict,
-    inputs=[image],
+    inputs=[img],
     outputs=outputs,
     title="Holocron: image classification demo",
-    article=("<p style='text-align: center'><a href='https://github.com/frgfm/Holocron'>" "Github Repo</a> | "
-             "<a href='https://frgfm.github.io/Holocron/'>Documentation</a></p>"),
+    article=(
+        "<p style='text-align: center'><a href='https://github.com/frgfm/Holocron'>"
+        "Github Repo</a> | "
+        "<a href='https://frgfm.github.io/Holocron/'>Documentation</a></p>"
+    ),
     live=True,
 ).launch()

requirements.txt

@@ -1,2 +1,5 @@
--e git+https://github.com/frgfm/Holocron.git#egg=pylocron
-gradio>=3.0.2
+gradio>=3.0.2,<4.0.0
+Pillow>=8.4.0
+onnxruntime>=1.10.0,<2.0.0
+huggingface-hub>=0.4.0,<1.0.0
+numpy>=1.19.5,<2.0.0