faceapp.py

import streamlit as st
from PIL import Image
from transformers import ViTForImageClassification, ViTImageProcessor

# Load the model
model_name = "trpakov/vit-face-expression"
model = ViTForImageClassification.from_pretrained(model_name)
image_processor = ViTImageProcessor.from_pretrained(model_name)

# Streamlit app
st.title("Emotion Recognition with vit-face-expression")

# Slider example
x = st.slider('Select a value')
st.write(f"{x} squared is {x * x}")

# Upload image
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]

    st.image(image, caption=f"Predicted emotion: {predicted_emotion}", use_column_width=True)