src/app.py

import streamlit as st
import numpy as np
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing import image
from PIL import Image
import io

# Set page config
st.set_page_config(
    page_title="AI Image Detection",
    page_icon="🔍",
    layout="centered"
)

# Configure Streamlit to handle file uploads better
if 'file_uploader_key' not in st.session_state:
    st.session_state.file_uploader_key = 0

# Load the model (cached to avoid reloading)
@st.cache_resource
def load_ai_detection_model():
    try:
        model = load_model('src/best_model.keras')
        return model
    except Exception as e:
        st.error(f"Error loading model: {str(e)}")
        return None

def preprocess_image(uploaded_file):
    """Preprocess the uploaded image for model prediction"""
    try:
        st.write("📸 Starting image preprocessing...")
        
        # Read file bytes
        file_bytes = uploaded_file.read()
        st.write(f"✅ Read {len(file_bytes)} bytes from file")
        
        # Reset file pointer for later use
        uploaded_file.seek(0)
        
        # Open image using PIL from bytes
        img = Image.open(io.BytesIO(file_bytes))
        st.write(f"✅ Opened image: {img.size}, mode: {img.mode}")
        
        # Convert to RGB if necessary (in case of RGBA or other formats)
        if img.mode != 'RGB':
            img = img.convert('RGB')
            st.write("✅ Converted image to RGB")
        
        # Resize to model's expected input size (300x300 based on your code)
        img = img.resize((300, 300), Image.Resampling.LANCZOS)
        st.write("✅ Resized image to 300x300")
        
        # Convert to array and normalize
        img_array = np.array(img, dtype=np.float32) / 255.0
        st.write(f"✅ Converted to array: {img_array.shape}, dtype: {img_array.dtype}")
        st.write(f"✅ Value range: {img_array.min():.3f} to {img_array.max():.3f}")
        
        # Add batch dimension
        img_array = np.expand_dims(img_array, axis=0)
        st.write(f"✅ Added batch dimension: {img_array.shape}")
        
        return img_array, img
    except Exception as e:
        st.error(f"Error preprocessing image: {str(e)}")
        import traceback
        st.error(f"Full error: {traceback.format_exc()}")
        return None, None

def predict_image(model, img_array):
    """Make prediction on the preprocessed image"""
    try:
        st.write("🔄 Making prediction...")
        prediction = model.predict(img_array, verbose=1)
        st.write(f"✅ Prediction completed: {prediction}")
        return prediction
    except Exception as e:
        st.error(f"Error making prediction: {str(e)}")
        st.write(f"Image array shape: {img_array.shape}")
        st.write(f"Image array dtype: {img_array.dtype}")
        return None

def main():
    # App title and description
    st.title("🔍 AI Image Detection")
    st.markdown("Upload an image to detect if it's AI-generated or real!")
    
    # Load model
    model = load_ai_detection_model()
    
    if model is None:
        st.error("Failed to load the model. Please check if 'best_model.keras' is in the correct location.")
        return
    
    # File uploader with size limit
    uploaded_file = st.file_uploader(
        "Choose an image file", 
        type=['jpg', 'jpeg', 'png', 'bmp', 'tiff'],
        help="Upload an image to analyze (max 5MB)",
        accept_multiple_files=False
    )
    
    # Check file size
    if uploaded_file is not None:
        file_size = len(uploaded_file.getvalue())
        max_size = 5 * 1024 * 1024  # 5MB in bytes
        
        if file_size > max_size:
            st.error(f"File size ({file_size/1024/1024:.2f}MB) exceeds the maximum limit of 5MB. Please upload a smaller image.")
            return
        
        st.success(f"✅ File uploaded successfully! Processing image...")
        
        # Reset file pointer after reading size
        uploaded_file.seek(0)
        # Display uploaded image
        col1, col2 = st.columns([1, 1])
        
        with col1:
            st.subheader("Uploaded Image")
            # Display file info
            st.write(f"**File name:** {uploaded_file.name}")
            st.write(f"**File size:** {file_size/1024:.2f} KB")
            st.image(uploaded_file, caption="Uploaded Image", use_column_width=True)
        
        # Reset file pointer before preprocessing
        uploaded_file.seek(0)
        
        # Preprocess image
        with st.spinner("Preprocessing image..."):
            img_array, processed_img = preprocess_image(uploaded_file)
        
        if img_array is not None:
            st.success("✅ Image preprocessed successfully!")
            
            # Make prediction
            with st.spinner("🤖 Analyzing image with AI model..."):
                prediction = predict_image(model, img_array)
            
            if prediction is not None:
                with col2:
                    st.subheader("🔍 Analysis Results")
                    
                    # Display raw prediction
                    st.write(f"**Raw Prediction Score:** {prediction[0][0]:.6f}")
                    
                    # Debug: Show prediction shape and values
                    st.write(f"**Prediction Shape:** {prediction.shape}")
                    st.write(f"**All Values:** {prediction[0]}")
                    
                    # Interpret prediction (adjust threshold based on your model)
                    # Assuming: prediction > 0.5 means AI-generated, < 0.5 means real
                    threshold = 0.5
                    confidence_score = float(prediction[0][0])
                    
                    if confidence_score > threshold:
                        result = "🤖 AI-Generated"
                        certainty = confidence_score
                        st.error(f"**Result:** {result}")
                        st.error("This image appears to be artificially generated!")
                    else:
                        result = "✅ Real/Human-made"
                        certainty = 1 - confidence_score
                        st.success(f"**Result:** {result}")
                        st.success("This image appears to be real!")
                    
                    # Display confidence level
                    st.write(f"**Confidence Level:**")
                    st.progress(min(certainty, 1.0))
                    st.write(f"**Confidence Percentage:** {certainty * 100:.2f}%")
                    
                    # Additional information
                    st.info("**Note:** This prediction is based on the trained model's analysis. "
                           "Results may vary depending on image quality and type.")
            else:
                st.error("❌ Failed to make prediction. Please try again.")
        else:
            st.error("❌ Failed to preprocess image. Please try uploading a different image.")
    
    # Add some information about the app
    with st.expander("ℹ️ About this app"):
        st.markdown("""
        This application uses a deep learning model to detect whether an image is AI-generated or real.
        
        **How to use:**
        1. Upload an image using the file uploader above
        2. Wait for the model to analyze the image
        3. View the prediction results and confidence score
        
        **Supported formats:** JPG, JPEG, PNG, BMP, TIFF
        
        **Model details:**
        - Input size: 300x300 pixels
        - The model outputs a probability score between 0 and 1
        - Scores > 0.5 typically indicate AI-generated images
        - Scores < 0.5 typically indicate real/human-made images
        """)

if __name__ == "__main__":
    main()