Update app.py

app.py

@@ -1,7 +1,7 @@
 import streamlit as st
 from transformers import pipeline
-import os
 from PIL import Image
+import os
 
 # Global variable to store the pipeline
 deepfake_pipe = None
@@ -24,6 +24,10 @@ def predict_deepfake(image):
         # 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)
         
@@ -57,37 +61,61 @@ def main():
     st.title("🕵️ Jerry - Deepfake Detection Tool")
     st.markdown("**Upload an image to check if it's real or AI-generated!**")
     
-    # Create two columns
-    col1, col2 = st.columns(2)
+    # 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")
-        uploaded_file = st.file_uploader(
-            "Choose an image",
-            type=['png', 'jpg', 'jpeg'],
-            label_visibility="collapsed"
-        )
-        
-        if uploaded_file is not None:
-            # Display uploaded image
-            image = Image.open(uploaded_file)
-            st.image(image, caption="Uploaded Image", use_column_width=True)
+        # 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")
-        if uploaded_file is not None:
-            if st.button("🔍 Analyze Image", type="primary"):
-                with st.spinner("Analyzing image..."):
-                    # Convert uploaded file to filepath for pipeline
-                    prediction, verdict, color = predict_deepfake(uploaded_file)
+        # 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 prediction.items():
+                    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:{color};font-size:18px;'><b>{verdict}</b></p>", unsafe_allow_html=True)
+                    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(