Update app.py

app.py

@@ -2,42 +2,70 @@ import streamlit as st
 import torch
 from PIL import Image
 from torchvision import transforms
+from typing import List, Tuple
 
-# Load your model (ensure this is the correct path to your model file)
+# Assuming your model and class names are set up correctly
+pretrained_vit_path = 'pretrained_vit_model_full.pth'
+class_names = ['Angry', 'Other', 'Sad', 'happy']
+
+# Load your model
 @st.cache(allow_output_mutation=True)
 def load_model():
-    model = torch.load('pretrained_vit_model_full.pth', map_location=torch.device('cpu'))
+    model = torch.load(pretrained_vit_path, 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):
+# Function to apply transforms to the image
+def transform_image(image, size=(224, 224)):
     transform = transforms.Compose([
-        transforms.Resize((224, 224)),  # Resize to the input size that your model expects
+        transforms.Resize(size),  
+        transforms.CenterCrop(size),  # Ensures image is square
         transforms.ToTensor(),
-        # Add other transformations as needed
+        # Add other transformations as needed, such as normalization
     ])
     return transform(image).unsqueeze(0)  # Add batch dimension
 
+# Prediction function
+def predict(model, image_tensor):
+    with torch.no_grad():
+        outputs = model(image_tensor)
+        _, predicted = torch.max(outputs, 1)
+        probabilities = torch.nn.functional.softmax(outputs, dim=1)
+        top_prob, top_catid = torch.topk(probabilities, 1)
+    return class_names[predicted[0]], top_prob[0].item()
+
+# Streamlit interface
 st.title("Animal Facial Expression Recognition")
 
+# Create two columns for the layout
+col1, col2 = st.columns([1, 1])
 
-# 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...")
+# First column for the uploader
+with col1:
+    uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "png", "jpeg"])
 
-    # Transform the image
-    input_tensor = transform_image(image)
+# Second column for the prediction results
+with col2:
+    if uploaded_file is not None:
+        # Display "Classifying..." text
+        st.write("Classifying...")
+    else:
+        # Show a message when no image is uploaded
+        st.write("Upload an image to see the classification result.")
 
-    # Make prediction
-    with torch.no_grad():
-        prediction = model(input_tensor)
+# If an image has been uploaded, display it and run the prediction
+if uploaded_file is not None:
+    # Display the uploaded image in the first column
+    with col1:
+        image = Image.open(uploaded_file).convert('RGB')
+        st.image(image, caption='Uploaded Image', use_column_width=True)
+        
+    # Transform the image and make prediction in the second column
+    with col2:
+        image_tensor = transform_image(image)
+        predicted_class, probability = predict(model, image_tensor)
+        st.write(f'Predicted class: {predicted_class}')
+        st.write(f'Probability: {probability:.3f}')
 
-    # Display the prediction (modify as per your output)
-    st.write('Predicted class:', prediction.argmax().item())