Update app.py

app.py

@@ -4,6 +4,10 @@ import torchvision.transforms as T
 from PIL import Image
 from module import myModule
 
+CLASS_TO_IDX = ['AMMAN', 'AYYAPPA', 'BHAIRAV', 'BRAHMA', 'BUDDHA', 'DURGA', 'GANESHA', 'HANUMAN', 'KAALI', 
+                'KRISHNA', 'KURMA', 'LAKSHMI', 'LINGA', 'MATSYA', 'MURUGA', 'NARASIMHA', 'NATARAJA', 'PARASURAMA', 
+                'RAMA', 'SARASWATI', 'SHIVA', 'THIRTHANKARA', 'VAMANA', 'VARAHA', 'VISHNU']
+
 IMG_SIZE = (224, 224)
 STATS = dict(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
 # Define the transformation for the input image
@@ -66,5 +70,6 @@ if img is not None:
     # Make predictions when the user clicks the "Predict" button
     if st.button("Predict"):
         values, indices = predict(img)
+        classes = [CLASS_TO_IDX[index] for index in indices]
         # Display the top 3 predictions as a bar chart
-        st.bar_chart({label: prob for label, prob in zip(indices, values)}, color="#FFC101")
+        st.bar_chart({label: prob for label, prob in zip(classes, values)}, color="#FFC101")