Merge pull request #6 from kalcohol/main-fix-sample

fix classification
AXERA-TECH · Dec 5, 2024 · 6d4192d · 6d4192d
2 parents 916d2d9 + 2ba26e2
commit 6d4192d
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/examples/classification.py b/examples/classification.py
@@ -52,8 +52,8 @@ def preprocess_image(image_path, target_size=(256, 256), crop_size=(224, 224)):
 
 def get_top_k_predictions(output, k=5):
     # Get top k predictions
-    top_k_indices = np.argsort(output[0])[-k:][::-1]
-    top_k_scores = output[0][top_k_indices]
+    top_k_indices = np.argsort(output[0].flatten())[-k:][::-1]
+    top_k_scores = output[0].flatten()[top_k_indices]
     return top_k_indices, top_k_scores