Argsort operation is now integrated into the ONNX model

app.py

@@ -4,11 +4,11 @@ import random
 import gradio as gr
 import matplotlib.pyplot as plt
 import numpy as np
+import onnxruntime as ort
 import polars as pl
 import seaborn as sns
 from huggingface_hub import hf_hub_download
 from PIL import Image
-import onnxruntime as ort
 
 # Seaborn configuration
 sns.set_theme()
@@ -43,12 +43,12 @@ def compute_word_weights(vocabulary):
     total_train_count = train_counts["train_count"].sum()
     word_weights = [(vocab["word"], vocab["train_count"] / total_train_count)
                     for vocab in vocabulary.rows(named=True)]
-    words = [word for word, _ in word_weights]
-    weights = [weight for _, weight in word_weights]
-    return words, weights
+    return [word for word, _ in word_weights], [weight for _, weight in word_weights]
 
 
 ort_session = load_model()
+input_name = ort_session.get_inputs()[0].name
+
 vocabulary = load_vocabulary()
 words, weights = compute_word_weights(vocabulary)
 
@@ -60,13 +60,10 @@ def get_random_word():
 # Process the image drawn on canvas
 def process_image(image, current_word):
     input_img = image["composite"].resize((224, 224))
-    inputs = {ort_session.get_inputs()[0].name: np.array(input_img)}
-    outputs = ort_session.run(None, inputs)[0]
-    indices_i = np.argsort(outputs)[::-1]
-    preds_i = outputs[indices_i]
+    indices, preds = ort_session.run(None, {input_name: np.array(input_img)})
 
     predictions = []
-    for pred, idx in zip(preds_i, indices_i):
+    for pred, idx in zip(preds, indices):
         vocab = vocabulary.row(idx, named=True)
         predictions.append({"word": vocab["word"], "category": vocab["category_name"], "prob": pred})
 
@@ -98,15 +95,9 @@ def process_image(image, current_word):
     ax.set_title("Top 10 Predictions", pad=15)
     ax.set_xlabel("Probability")
     ax.set_ylabel(None)
-    plt.close(fig)
     return fig, current_word
 
 
-def update_image(image):
-    image = Image.fromarray(image["composite"])
-    return image
-
-
 def create_initial_image():
     data = np.full((520, 700, 3), 255, dtype=np.uint8)  # White image
     return Image.fromarray(data)