app.py

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

# Global variable to store the pipeline
deepfake_pipe = None

def load_model():
    """Load the model once and cache it"""
    global deepfake_pipe
    if deepfake_pipe is None:
        st.info("Loading Deepfake Detection Model...")
        deepfake_pipe = pipeline(
            "image-classification", 
            model="prithivMLmods/Deep-Fake-Detector-v2-Model"
        )
        st.success("Model loaded successfully!")
    return deepfake_pipe

def predict_deepfake(image):
    """Predict if image is deepfake or real"""
    try:
        # Load model if not already loaded
        pipe = load_model()
        
        # Ensure image is in PIL format
        if not isinstance(image, Image.Image):
            image = Image.open(image)
        
        # Make prediction
        results = pipe(image)
        
        # Format results
        prediction = {result['label']: result['score'] for result in results}
        
        # Determine final verdict
        deepfake_score = prediction.get('Deepfake', 0)
        realism_score = prediction.get('Realism', 0)
        
        if deepfake_score > realism_score:
            verdict = f"🚨 DEEPFAKE DETECTED ({deepfake_score:.2%} confidence)"
            color = "red"
        else:
            verdict = f"✅ REAL IMAGE ({realism_score:.2%} confidence)"
            color = "green"
        
        return prediction, verdict, color
        
    except Exception as e:
        return {"Error": 1.0}, f"❌ Prediction failed: {str(e)}", "red"

# Streamlit app configuration
st.set_page_config(
    page_title="Jerry - Deepfake Detector",
    page_icon="🕵️",
    layout="wide"
)

def main():
    st.title("🕵️ Jerry - Deepfake Detection Tool")
    st.markdown("**Upload an image to check if it's real or AI-generated!**")
    
    # Use session state to store results and prevent unnecessary re-renders
    if 'results' not in st.session_state:
        st.session_state.results = None
        st.session_state.verdict = None
        st.session_state.color = None
    
    # Create two columns with fixed widths
    col1, col2 = st.columns([1, 1], gap="medium")
    
    with col1:
        st.subheader("Upload Image")
        # Fixed-height container for image to prevent layout shifts
        with st.container(height=400, border=True):
            uploaded_file = st.file_uploader(
                "Choose an image",
                type=['png', 'jpg', 'jpeg'],
                label_visibility="collapsed",
                key="file_uploader"
            )
            
            if uploaded_file is not None:
                # Display uploaded image with fixed dimensions
                image = Image.open(uploaded_file)
                st.image(
                    image,
                    caption="Uploaded Image",
                    use_container_width=True,
                    clamp=True
                )
    
    with col2:
        st.subheader("Detection Results")
        # Fixed-height container for results
        with st.container(height=400, border=True):
            if uploaded_file is not None:
                if st.button("🔍 Analyze Image", type="primary"):
                    with st.spinner("Analyzing image..."):
                        # Perform prediction and store in session state
                        prediction, verdict, color = predict_deepfake(uploaded_file)
                        st.session_state.results = prediction
                        st.session_state.verdict = verdict
                        st.session_state.color = color
                    
                    # Display results
                    st.write("**Prediction Scores:**")
                    for label, score in st.session_state.results.items():
                        st.write(f"{label}: {score:.2%}")
                    
                    # Display verdict with colored text
                    st.markdown(
                        f"<p style='color:{st.session_state.color};font-size:18px;'><b>{st.session_state.verdict}</b></p>",
                        unsafe_allow_html=True
                    )
            else:
                st.write("Please upload an image to analyze.")
    
    # Additional information
    st.markdown(
        """
        ---
        **How it works:**
        - Upload any facial image
        - Jerry analyzes it and gives appropriate verdict
        - Returns confidence scores for "Deepfake" vs "Realism"
        
        **Note:** The model loads only once when first used, then runs quickly for subsequent predictions!
        """
    )

if __name__ == "__main__":
    print("🚀 Starting Jerry - Deepfake Detection App...")
    print("Model will load on first prediction...")
    main()