Update pages/11_Cifar10_HF.py

pages/11_Cifar10_HF.py

@@ -1,12 +1,10 @@
 import streamlit as st
 from transformers import ViTFeatureExtractor, ViTForImageClassification
 from PIL import Image
-import requests
-import numpy as np
 import torch
 
-# Load pre-trained model and feature extractor
-model_name = "google/vit-base-patch16-224"
+# Load pre-trained model and feature extractor for CIFAR-10
+model_name = "aaraki/vit-base-patch16-224-in21k-finetuned-cifar10"
 feature_extractor = ViTFeatureExtractor.from_pretrained(model_name)
 model = ViTForImageClassification.from_pretrained(model_name)
 
@@ -32,4 +30,9 @@ if uploaded_file is not None:
             outputs = model(**inputs)
             logits = outputs.logits
             predicted_class_idx = logits.argmax(-1).item()
-            st.write(f"Predicted Class: {predicted_class_idx} ({class_names[predicted_class_idx]})")
+            
+            # Check if the predicted_class_idx is within bounds
+            if 0 <= predicted_class_idx < len(class_names):
+                st.write(f"Predicted Class: {predicted_class_idx} ({class_names[predicted_class_idx]})")
+            else:
+                st.error("Prediction index out of range.")