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

# Set the title and sidebar
st.set_page_config(page_title="Deepfake vs Real Image Detection", page_icon="🤖")

st.title("Deepfake vs Real Image Detection")
st.sidebar.title("Options")

# Description
st.markdown(
    """
    Welcome to the Deepfake vs Real Image Detection app! 
    Upload an image and let our model determine if it is real or a deepfake.
    """
)

# Load the pipeline
model_name = "vm24bho/net_dfm_myimg"
try:
    pipe = pipeline('image-classification', model=model_name)
except Exception as e:
    st.error(f"Error loading model: {e}")
    st.stop()

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

# Add a sample image option
sample_image = None
if st.sidebar.button("Use Sample Image"):
    sample_url = "https://drive.google.com/file/d/15zh_XjwH9gGAzNdNUEgHraYAzE4GdSRU/view?usp=drive_link"  # Replace with a valid sample image URL
    try:
        response = requests.get(sample_url, stream=True)
        sample_image = Image.open(response.raw)
    except Exception as e:
        st.error(f"Error loading sample image: {e}")
        st.stop()

if uploaded_file is not None:
    image = Image.open(uploaded_file)
elif sample_image is not None:
    image = sample_image
else:
    image = None

# Display the uploaded image
if image is not None:
    st.image(image, caption='Uploaded Image', use_column_width=True)
    st.write("")
    st.write("Classifying...")

    # Apply the model
    try:
        result = pipe(image)
    except Exception as e:
        st.error(f"Error classifying image: {e}")
        st.stop()

    # Determine if the image is classified as real or fake
    if result[0]['label'] == 'FAKE':
        st.write("This Image is: Fake")
    else:
        st.write("This Image is: Real")

    # Display the result
    st.subheader("Classification Result")
    for res in result:
        st.write(f"**{res['label']}**: {res['score']*100:.2f}%")
else:
    st.sidebar.info("Upload an image to get started or use the sample image.")