Update app.py

app.py

@@ -1,28 +1,9 @@
 import streamlit as st
-import torch
-
-#streamlit clean
-#streamlit run app.py
-
-#pip install --upgrade pip
-#curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh
-
-
 from PIL import Image
-#from transformers import ViTForImageClassification, ViTImageProcessor
-
+from transformers import pipeline
 
-from transformers import AutoImageProcessor, AutoModelForImageClassification
-
-# Load the model
-model_name = "trpakov/vit-face-expression"
-image_processor = AutoImageProcessor.from_pretrained(model_name)
-model = AutoModelForImageClassification.from_pretrained(model_name)
-
-# Load the model
-model_name = "trpakov/vit-face-expression"
-#model = ViTForImageClassification.from_pretrained(model_name)
-#image_processor = ViTImageProcessor.from_pretrained(model_name)
+# Create an image classification pipeline
+pipe = pipeline("image-classification", model="trpakov/vit-face-expression")
 
 # Streamlit app
 st.title("Emotion Recognition with vit-face-expression")
@@ -36,21 +17,15 @@ uploaded_image = st.file_uploader("Upload an image", type=["jpg", "png"])
 
 if uploaded_image:
     image = Image.open(uploaded_image)
-    inputs = image_processor(images=image, return_tensors="pt")
-    pixel_values = inputs.pixel_values
-
-    # Predict emotion
-    with torch.no_grad():
-        outputs = model(pixel_values)
-        predicted_class = torch.argmax(outputs.logits, dim=1).item()
-
-    emotion_labels = ["Angry", "Disgust", "Fear", "Happy", "Sad", "Surprise", "Neutral"]
-    predicted_emotion = emotion_labels[predicted_class]
+    # Predict emotion using the pipeline
+    results = pipe(image)
+    predicted_class = results[0]["label"]
+    predicted_emotion = predicted_class.split("_")[-1].capitalize()
 
     st.image(image, caption=f"Predicted emotion: {predicted_emotion}", use_column_width=True)
 
 
     # Display scores for each category
-    st.write("Emotion Scores:")
-    for label, score in zip(emotion_labels, outputs.logits[0]):
-        st.write(f"{label}: {score:.4f}")
+    #st.write("Emotion Scores:")
+    #for label, score in zip(emotion_labels, outputs.logits[0]):
+    #    st.write(f"{label}: {score:.4f}")