import streamlit as st
import torch
from PIL import Image
from torchvision import transforms

# Load your model (ensure this is the correct path to your model file)
@st.cache(allow_output_mutation=True)
def load_model():
    model = torch.load('pretrained_vit_model_full.pth', map_location=torch.device('cpu'))
    model.eval()
    return model

model = load_model()

# Function to apply transforms to the image (update as per your model's requirement)
def transform_image(image):
    transform = transforms.Compose([
        transforms.Resize((224, 224)),  # Resize to the input size that your model expects
        transforms.ToTensor(),
        # Add other transformations as needed
    ])
    return transform(image).unsqueeze(0)  # Add batch dimension

st.title("Animal Facial Expression Recognition")

# Slider
x = st.slider('Select a value')
st.write(x, 'squared is', x * x)

# File uploader
uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "png", "jpeg"])
if uploaded_file is not None:
    image = Image.open(uploaded_file).convert('RGB')
    st.image(image, caption='Uploaded Image.', use_column_width=True)
    st.write("")
    st.write("Classifying...")

    # Transform the image
    input_tensor = transform_image(image)

    # Make prediction
    with torch.no_grad():
        prediction = model(input_tensor)

    # Display the prediction (modify as per your output)
    st.write('Predicted class:', prediction.argmax().item())