import streamlit as st
from transformers import pipeline
from PIL import Image
import requests
from io import BytesIO

# Load the image classification pipeline
clf_pipeline = pipeline("image-classification", model="microsoft/resnet-50")

# Streamlit app
st.title("Image Classification with ResNet-50")

# Option to choose between file upload or URL input
option = st.radio("Choose image source:", ("Upload Image", "Image URL"))

image = None

if option == "Upload Image":
    uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
    if uploaded_file is not None:
        image = Image.open(uploaded_file)
        st.image(image, caption='Uploaded Image.', use_column_width=True)

elif option == "Image URL":
    img_url = st.text_input("Enter image URL:")
    if img_url:
        try:
            response = requests.get(img_url)
            image = Image.open(BytesIO(response.content))
            st.image(image, caption='Image from URL.', use_column_width=True)
        except Exception as e:
            st.error("Error loading image. Please check the URL.")

if image is not None:
    # Perform image classification
    st.write("Classifying the image...")
    results = clf_pipeline(image)
    
    # Display the classification results
    st.write("Results:")
    for result in results:
        st.write(f"Label: {result['label']}, Confidence: {result['score']:.4f}")