Update app.py

app.py

@@ -546,133 +546,175 @@ def main():
     # Image upload section
     with st.expander("Stage 2: Image Upload & Initial Detection", expanded=True):
         st.subheader("Upload an Image")
-        uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
         
-        if uploaded_file is not None:
-            try:
-                # Load and display the image (with controlled size)
-                image = Image.open(uploaded_file).convert("RGB")
-                
-                # Display the image with a controlled width
-                col1, col2 = st.columns([1, 2])
-                with col1:
-                    st.image(image, caption="Uploaded Image", width=300)
-                
-                # Generate detailed caption for original image if BLIP model is loaded
-                if st.session_state.blip_model_loaded:
-                    with st.spinner("Generating image description..."):
-                        caption = generate_image_caption(
-                            image, 
-                            st.session_state.original_processor, 
-                            st.session_state.original_model
-                        )
-                        st.session_state.image_caption = caption
-                
-                # Detect with Xception model if loaded
-                if st.session_state.xception_model_loaded:
-                    try:
-                        with st.spinner("Analyzing image with Xception model..."):
-                            # Preprocess image for Xception
-                            st.write("Starting Xception processing...")
-                            input_tensor, original_image, face_box = preprocess_image_xception(image)
-                            
-                            # Get device and model
-                            device = st.session_state.device
-                            model = st.session_state.xception_model
-                            
-                            # Ensure model is in eval mode
-                            model.eval()
-                            
-                            # Move tensor to device
-                            input_tensor = input_tensor.to(device)
-                            st.write(f"Input tensor on device: {device}")
-                            
-                            # Forward pass with proper error handling
-                            try:
-                                with torch.no_grad():
-                                    st.write("Running model inference...")
-                                    logits = model(input_tensor)
-                                    st.write(f"Raw logits: {logits}")
-                                    probabilities = torch.softmax(logits, dim=1)[0]
-                                    st.write(f"Probabilities: {probabilities}")
-                                    pred_class = torch.argmax(probabilities).item()
-                                    confidence = probabilities[pred_class].item()
-                                    st.write(f"Predicted class: {pred_class}, Confidence: {confidence:.4f}")
-                                    
-                                    # Explicit class mapping - adjust if needed based on your model
-                                    pred_label = "Fake" if pred_class == 0 else "Real"
-                                    st.write(f"Mapped to label: {pred_label}")
-                            except Exception as e:
-                                st.error(f"Error in model inference: {str(e)}")
-                                import traceback
-                                st.error(traceback.format_exc())
-                                # Set default values
-                                pred_class = 0
-                                confidence = 0.5
-                                pred_label = "Error in prediction"
-                            
-                            # Display results
-                            with col2:
-                                st.markdown("### Detection Result")
-                                st.markdown(f"**Classification:** {pred_label} (Confidence: {confidence:.2%})")
+        # Add alternative upload methods
+        upload_tab1, upload_tab2, upload_tab3 = st.tabs(["File Upload", "URL Input", "Sample Images"])
+        
+        uploaded_image = None
+        
+        with upload_tab1:
+            uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
+            if uploaded_file is not None:
+                try:
+                    # Try to load directly from bytes to avoid file system issues
+                    file_bytes = uploaded_file.getvalue()
+                    uploaded_image = Image.open(io.BytesIO(file_bytes)).convert("RGB")
+                    st.session_state.upload_method = "file"
+                except Exception as e:
+                    st.error(f"Error loading file: {str(e)}")
+        
+        with upload_tab2:
+            url = st.text_input("Enter image URL:")
+            if url and url.strip():
+                try:
+                    import requests
+                    response = requests.get(url, stream=True)
+                    if response.status_code == 200:
+                        uploaded_image = Image.open(io.BytesIO(response.content)).convert("RGB")
+                        st.session_state.upload_method = "url"
+                    else:
+                        st.error(f"Failed to load image from URL: Status code {response.status_code}")
+                except Exception as e:
+                    st.error(f"Error loading image from URL: {str(e)}")
+        
+        with upload_tab3:
+            sample_images = {
+                "Real face sample": "https://raw.githubusercontent.com/deepfakes/faceswap/master/docs/full_guide/images/test_face.jpg",
+                "Likely deepfake sample": "https://storage.googleapis.com/deepfake-detection/example_deepfake.jpg",
+                "Neutral face sample": "https://t4.ftcdn.net/jpg/02/19/63/31/360_F_219633151_BW6TD8D1EA9OqZu4JgdmeJGg4JBaiAHj.jpg",
+            }
+            selected_sample = st.selectbox("Select a sample image:", list(sample_images.keys()))
+            if st.button("Load Sample"):
+                try:
+                    import requests
+                    response = requests.get(sample_images[selected_sample], stream=True)
+                    if response.status_code == 200:
+                        uploaded_image = Image.open(io.BytesIO(response.content)).convert("RGB")
+                        st.session_state.upload_method = "sample"
+                    else:
+                        st.error(f"Failed to load sample image: Status code {response.status_code}")
+                except Exception as e:
+                    st.error(f"Error loading sample image: {str(e)}")
+        
+        # If we have an uploaded image, process it
+        if uploaded_image is not None:
+            # Display the image
+            image = uploaded_image
+            col1, col2 = st.columns([1, 2])
+            with col1:
+                st.image(image, caption="Uploaded Image", width=300)
+            
+            # Generate detailed caption for original image if BLIP model is loaded
+            if st.session_state.blip_model_loaded:
+                with st.spinner("Generating image description..."):
+                    caption = generate_image_caption(
+                        image, 
+                        st.session_state.original_processor, 
+                        st.session_state.original_model
+                    )
+                    st.session_state.image_caption = caption
+            
+            # Continue with your existing code for the Xception model analysis
+            if st.session_state.xception_model_loaded:
+                try:
+                    with st.spinner("Analyzing image with Xception model..."):
+                        # Preprocess image for Xception
+                        st.write("Starting Xception processing...")
+                        input_tensor, original_image, face_box = preprocess_image_xception(image)
+                        
+                        # Get device and model
+                        device = st.session_state.device
+                        model = st.session_state.xception_model
+                        
+                        # Ensure model is in eval mode
+                        model.eval()
+                        
+                        # Move tensor to device
+                        input_tensor = input_tensor.to(device)
+                        st.write(f"Input tensor on device: {device}")
+                        
+                        # Forward pass with proper error handling
+                        try:
+                            with torch.no_grad():
+                                st.write("Running model inference...")
+                                logits = model(input_tensor)
+                                st.write(f"Raw logits: {logits}")
+                                probabilities = torch.softmax(logits, dim=1)[0]
+                                st.write(f"Probabilities: {probabilities}")
+                                pred_class = torch.argmax(probabilities).item()
+                                confidence = probabilities[pred_class].item()
+                                st.write(f"Predicted class: {pred_class}, Confidence: {confidence:.4f}")
                                 
-                                # Display face box on image if detected
-                                if face_box:
-                                    img_to_show = original_image.copy()
-                                    img_draw = np.array(img_to_show)
-                                    x, y, w, h = face_box
-                                    cv2.rectangle(img_draw, (x, y), (x + w, y + h), (0, 255, 0), 2)
-                                    st.image(Image.fromarray(img_draw), caption="Detected Face", width=300)
+                                # Explicit class mapping - adjust if needed based on your model
+                                pred_label = "Fake" if pred_class == 0 else "Real"
+                                st.write(f"Mapped to label: {pred_label}")
+                        except Exception as e:
+                            st.error(f"Error in model inference: {str(e)}")
+                            import traceback
+                            st.error(traceback.format_exc())
+                            # Set default values
+                            pred_class = 0
+                            confidence = 0.5
+                            pred_label = "Error in prediction"
+                        
+                        # Display results
+                        with col2:
+                            st.markdown("### Detection Result")
+                            st.markdown(f"**Classification:** {pred_label} (Confidence: {confidence:.2%})")
                             
-                            # GradCAM visualization with error handling
-                            st.subheader("GradCAM Visualization")
-                            try:
-                                st.write("Generating GradCAM visualization...")
-                                cam, overlay, comparison, detected_face_box = process_image_with_xception_gradcam(
-                                    image, model, device, pred_class
-                                )
-                                
-                                if comparison:
-                                    # Display GradCAM results (controlled size)
-                                    st.image(comparison, caption="Original | CAM | Overlay", width=700)
-                                    
-                                    # Save for later use
-                                    st.session_state.comparison_image = comparison
-                                else:
-                                    st.error("GradCAM visualization failed - comparison image not generated")
-                                
-                                # Generate caption for GradCAM overlay image if BLIP model is loaded
-                                if st.session_state.blip_model_loaded and overlay:
-                                    with st.spinner("Analyzing GradCAM visualization..."):
-                                        gradcam_caption = generate_gradcam_caption(
-                                            overlay,
-                                            st.session_state.finetuned_processor,
-                                            st.session_state.finetuned_model
-                                        )
-                                        st.session_state.gradcam_caption = gradcam_caption
-                            except Exception as e:
-                                st.error(f"Error generating GradCAM: {str(e)}")
-                                import traceback
-                                st.error(traceback.format_exc())
+                            # Display face box on image if detected
+                            if face_box:
+                                img_to_show = original_image.copy()
+                                img_draw = np.array(img_to_show)
+                                x, y, w, h = face_box
+                                cv2.rectangle(img_draw, (x, y), (x + w, y + h), (0, 255, 0), 2)
+                                st.image(Image.fromarray(img_draw), caption="Detected Face", width=300)
+                        
+                        # GradCAM visualization with error handling
+                        st.subheader("GradCAM Visualization")
+                        try:
+                            st.write("Generating GradCAM visualization...")
+                            cam, overlay, comparison, detected_face_box = process_image_with_xception_gradcam(
+                                image, model, device, pred_class
+                            )
                             
-                            # Save results in session state for LLM analysis
-                            st.session_state.current_image = image
-                            st.session_state.current_overlay = overlay if 'overlay' in locals() else None
-                            st.session_state.current_face_box = detected_face_box if 'detected_face_box' in locals() else None
-                            st.session_state.current_pred_label = pred_label
-                            st.session_state.current_confidence = confidence
+                            if comparison:
+                                # Display GradCAM results (controlled size)
+                                st.image(comparison, caption="Original | CAM | Overlay", width=700)
+                                
+                                # Save for later use
+                                st.session_state.comparison_image = comparison
+                            else:
+                                st.error("GradCAM visualization failed - comparison image not generated")
                             
-                            st.success("✅ Initial detection and GradCAM visualization complete!")
-                    except Exception as e:
-                        st.error(f"Overall error in Xception processing: {str(e)}")
-                        import traceback
-                        st.error(traceback.format_exc())
-                else:
-                    st.warning("⚠️ Please load the Xception model first to perform initial detection.")
-            except Exception as e:
-                st.error(f"Error processing image: {str(e)}")
-                import traceback
-                st.error(traceback.format_exc())  # This will show the full error traceback
+                            # Generate caption for GradCAM overlay image if BLIP model is loaded
+                            if st.session_state.blip_model_loaded and overlay:
+                                with st.spinner("Analyzing GradCAM visualization..."):
+                                    gradcam_caption = generate_gradcam_caption(
+                                        overlay,
+                                        st.session_state.finetuned_processor,
+                                        st.session_state.finetuned_model
+                                    )
+                                    st.session_state.gradcam_caption = gradcam_caption
+                        except Exception as e:
+                            st.error(f"Error generating GradCAM: {str(e)}")
+                            import traceback
+                            st.error(traceback.format_exc())
+                        
+                        # Save results in session state for LLM analysis
+                        st.session_state.current_image = image
+                        st.session_state.current_overlay = overlay if 'overlay' in locals() else None
+                        st.session_state.current_face_box = detected_face_box if 'detected_face_box' in locals() else None
+                        st.session_state.current_pred_label = pred_label
+                        st.session_state.current_confidence = confidence
+                        
+                        st.success("✅ Initial detection and GradCAM visualization complete!")
+                except Exception as e:
+                    st.error(f"Overall error in Xception processing: {str(e)}")
+                    import traceback
+                    st.error(traceback.format_exc())
+            else:
+                st.warning("⚠️ Please load the Xception model first to perform initial detection.")
     
     # Image Analysis Summary section - AFTER Stage 2
     if hasattr(st.session_state, 'current_image') and (hasattr(st.session_state, 'image_caption') or hasattr(st.session_state, 'gradcam_caption')):