import streamlit as st
from PIL import Image
import torch
import torchvision.transforms as transforms
from torchvision import models
import torch.nn.functional as F

# Set up a title for the app
st.title("Simple Image Recognition App")

# Load a pre-trained model from torchvision (e.g., ResNet50)
model = models.resnet50(pretrained=True)
model.eval()  # Set the model to evaluation mode

# Upload an image
uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "png", "jpeg"])

if uploaded_file is not None:
    # Display the uploaded image
    image = Image.open(uploaded_file).convert("RGB")
    st.image(image, caption='Uploaded Image.', use_column_width=True)

    # Preprocess the image
    preprocess = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])
    image_tensor = preprocess(image).unsqueeze(0)  # Add a batch dimension

    # Run the model to make predictions
    st.write("Classifying...")
    with torch.no_grad():
        outputs = model(image_tensor)
        probabilities = F.softmax(outputs[0], dim=0)
        top3_prob, top3_classes = torch.topk(probabilities, 3)

        # Display the top 3 predictions
        st.write("Predictions:")
        for i in range(3):
            st.write(f"Label: {top3_classes[i].item()}, Confidence: {top3_prob[i].item():.2f}")