import streamlit as st
from transformers import ViTImageProcessor, AutoModelForImageClassification
from PIL import Image
import requests
from io import BytesIO

# Load the model and processor
processor = ViTImageProcessor.from_pretrained('AdamCodd/vit-base-nsfw-detector')
model = AutoModelForImageClassification.from_pretrained('AdamCodd/vit-base-nsfw-detector')

# Define prediction function
def predict_image(image):
    try:
        # Process the image and make prediction
        inputs = processor(images=image, return_tensors="pt")
        outputs = model(**inputs)
        logits = outputs.logits

        # Get predicted class
        predicted_class_idx = logits.argmax(-1).item()
        predicted_label = model.config.id2label[predicted_class_idx]

        return predicted_label
    except Exception as e:
        return str(e)

# Streamlit app
st.title("NSFW Image Classifier")

# Upload image file
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)
    st.write("")
    st.write("Classifying...")

    # Predict and display result
    prediction = predict_image(image)
    st.write(f"Predicted Class: {prediction}")