Spaces:
Sleeping
Sleeping
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) | |
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()) | |