Update app.py

app.py

@@ -693,64 +693,18 @@ def main():
     if 'chat_history' not in st.session_state:
         st.session_state.chat_history = []
     
-    # Display Hugging Face Spaces information if debug mode is on
-    if st.session_state.debug:
-        with st.expander("Hugging Face Spaces Debugging Information", expanded=True):
-            st.markdown("""
-            ### Common Issues with Hugging Face Spaces
-            
-            1. **403/404 Errors**: Often caused by permission issues when accessing files or external resources.
-            
-            2. **Memory Limits**: Free spaces have limited memory (16GB). Large models may cause OOM errors.
-            
-            3. **Disk Space**: Limited to 10GB for persistent storage.
-            
-            4. **Network Restrictions**: Some external URLs might be blocked or restricted.
-            
-            ### Accessing Logs
-            To see detailed error logs in Hugging Face Spaces:
-            
-            1. Go to your Space dashboard
-            2. Click on "Logs" in the left sidebar
-            3. Check both "Build logs" and "Running logs" tabs
-            
-            In the running logs, look for Python tracebacks or error messages.
-            
-            ### This App's Setup
-            - All image processing is now done in-memory to avoid file permission issues
-            - Debug logging is available through this interface
-            - For large model loading issues, try using smaller models or increasing RAM allocation
-            """)
-            
-            # Add a test connection button
-            if st.button("Test Network Connection"):
-                try:
-                    import requests
-                    test_urls = [
-                        "https://huggingface.co/", 
-                        "https://www.google.com/",
-                        "https://jsonplaceholder.typicode.com/todos/1"
-                    ]
-                    
-                    for url in test_urls:
-                        try:
-                            response = requests.get(url, timeout=5)
-                            st.write(f"✅ {url}: Status {response.status_code}")
-                        except Exception as e:
-                            st.write(f"❌ {url}: Error - {str(e)}")
-                except Exception as e:
-                    st.error(f"Could not perform connection test: {str(e)}")
+    # Create multi-tab interface
+    tab1, tab2, tab3 = st.tabs(["Deepfake Detection", "Image Captions", "LLM Analysis"])
     
-    # Create expanders for each stage
-    with st.expander("Stage 1: Model Loading", expanded=True):
-        st.write("Please load the models using the buttons below:")
+    # Tab 1: Deepfake Detection with Model Loading and Image Upload
+    with tab1:
+        st.header("Deepfake Detection")
         
-        # Button for loading models
-        xception_col, blip_col, llm_col = st.columns(3)
-        
-        with xception_col:
+        # Model Loading section
+        with st.expander("Load Detection Model", expanded=True):
+            st.write("Please load the Xception model for deepfake detection:")
             if not st.session_state.xception_model_loaded:
-                if st.button("📥 Load Xception Model for Detection", type="primary"):
+                if st.button("📥 Load Xception Model", type="primary"):
                     # Load Xception model
                     model, device = load_detection_model_xception()
                     if model is not None:
@@ -763,7 +717,226 @@ def main():
             else:
                 st.success("✅ Xception model loaded and ready!")
         
-        with blip_col:
+        # Image upload section
+        with st.expander("Upload and Analyze Image", expanded=True):
+            st.subheader("Upload an Image")
+            
+            # Add alternative upload methods
+            upload_tab1, upload_tab2 = st.tabs(["File Upload", "URL Input"])
+            
+            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:
+                        # Simple direct approach - load the image directly
+                        image = Image.open(uploaded_file).convert("RGB")
+                        uploaded_image = image
+                        st.session_state.upload_method = "file"
+                    except Exception as e:
+                        st.error(f"Error loading image: {str(e)}")
+                        import traceback
+                        st.error(traceback.format_exc())
+            
+            with upload_tab2:
+                url = st.text_input("Enter image URL:")
+                if url and url.strip():
+                    try:
+                        import requests
+                        # Simplified URL handling with more robust approach
+                        headers = {
+                            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36',
+                            'Accept': 'image/jpeg, image/png, image/*, */*',
+                            'Referer': 'https://huggingface.co/'
+                        }
+                        
+                        # Try three different methods to handle various API restrictions
+                        try_methods = True
+                        
+                        # Method 1: Direct requests
+                        if try_methods:
+                            try:
+                                response = requests.get(url, stream=True, headers=headers, timeout=10)
+                                if response.status_code == 200 and 'image' in response.headers.get('Content-Type', ''):
+                                    try:
+                                        image = Image.open(io.BytesIO(response.content)).convert("RGB")
+                                        uploaded_image = image
+                                        st.session_state.upload_method = "url_direct"
+                                        try_methods = False
+                                        st.success("✅ Image loaded via direct request")
+                                    except Exception as e:
+                                        st.warning(f"Direct method received data but couldn't process as image: {str(e)}")
+                                else:
+                                    st.info(f"Direct method failed: Status {response.status_code}, trying alternative method...")
+                            except Exception as e:
+                                st.info(f"Direct method error: {str(e)}, trying alternative method...")
+                        
+                        # Method 2: Use Python's urllib as fallback
+                        if try_methods:
+                            try:
+                                import urllib.request
+                                from urllib.error import HTTPError
+                                
+                                opener = urllib.request.build_opener()
+                                opener.addheaders = [('User-agent', headers['User-Agent'])]
+                                urllib.request.install_opener(opener)
+                                
+                                with urllib.request.urlopen(url, timeout=10) as response:
+                                    image_data = response.read()
+                                    image = Image.open(io.BytesIO(image_data)).convert("RGB")
+                                    uploaded_image = image
+                                    st.session_state.upload_method = "url_urllib"
+                                    try_methods = False
+                                    st.success("✅ Image loaded via urllib")
+                            except HTTPError as e:
+                                st.info(f"urllib method failed: HTTP error {e.code}, trying next method...")
+                            except Exception as e:
+                                st.info(f"urllib method error: {str(e)}, trying next method...")
+                        
+                        # Method 3: Use a proxy service as last resort
+                        if try_methods:
+                            try:
+                                # This uses an image proxy service to bypass CORS issues
+                                # Only as last resort since it depends on external service
+                                proxy_url = f"https://images.weserv.nl/?url={url}"
+                                response = requests.get(proxy_url, stream=True, timeout=10)
+                                if response.status_code == 200:
+                                    image = Image.open(io.BytesIO(response.content)).convert("RGB")
+                                    uploaded_image = image
+                                    st.session_state.upload_method = "url_proxy"
+                                    try_methods = False
+                                    st.success("✅ Image loaded via proxy service")
+                                else:
+                                    st.error(f"All methods failed to load the image from URL. Last status: {response.status_code}")
+                            except Exception as e:
+                                st.error(f"All methods failed. Final error: {str(e)}")
+                        
+                        if not uploaded_image:
+                            st.error("Failed to load image using all available methods.")
+                            
+                    except Exception as e:
+                        st.error(f"Error processing URL: {str(e)}")
+                        if st.session_state.debug:
+                            import traceback
+                            st.error(traceback.format_exc())
+            
+            # 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)
+                
+                # Continue with Xception model analysis
+                if st.session_state.xception_model_loaded:
+                    try:
+                        with st.spinner("Analyzing image with Xception model..."):
+                            # Preprocess image for Xception
+                            input_tensor, original_image, face_box = preprocess_image_xception(image)
+                            
+                            if input_tensor is None:
+                                st.error("Failed to preprocess image. Please try another image.")
+                                st.stop()
+                                
+                            # 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)
+                            
+                            # Forward pass with proper error handling
+                            try:
+                                with torch.no_grad():
+                                    logits = model(input_tensor)
+                                    probabilities = torch.softmax(logits, dim=1)[0]
+                                    pred_class = torch.argmax(probabilities).item()
+                                    confidence = probabilities[pred_class].item()
+                                    
+                                    # Explicit class mapping - adjust if needed based on your model
+                                    pred_label = "Fake" if pred_class == 0 else "Real"
+                            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%})")
+                                
+                                # 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:
+                                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
+                                        
+                                        # Display the caption directly here as well for immediate feedback
+                                        st.markdown("### GradCAM Analysis")
+                                        st.markdown(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 use in other tabs
+                            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.")
+    
+    # Tab 2: Image Captions with BLIP models
+    with tab2:
+        st.header("Image Captions")
+        
+        # Model Loading section
+        with st.expander("Load Captioning Models", expanded=True):
             if not st.session_state.blip_model_loaded:
                 if st.button("📥 Load BLIP for Captioning", type="primary"):
                     # Load BLIP models
@@ -779,10 +952,65 @@ def main():
                         st.error("❌ Failed to load BLIP models.")
             else:
                 st.success("✅ BLIP captioning models loaded and ready!")
+        
+        # Image Caption Display
+        if hasattr(st.session_state, 'current_image'):
+            col1, col2 = st.columns([1, 2])
+            
+            with col1:
+                st.image(st.session_state.current_image, caption="Image", width=300)
                 
-        with llm_col:
+                if hasattr(st.session_state, 'current_overlay'):
+                    st.image(st.session_state.current_overlay, caption="GradCAM Visualization", width=300)
+            
+            with col2:
+                if not st.session_state.blip_model_loaded:
+                    st.warning("⚠️ Please load the BLIP models first to see captions.")
+                else:
+                    # Button to generate captions if not already generated
+                    if not hasattr(st.session_state, 'image_caption') or st.button("Regenerate Image Caption"):
+                        with st.spinner("Generating image description..."):
+                            caption = generate_image_caption(
+                                st.session_state.current_image, 
+                                st.session_state.original_processor, 
+                                st.session_state.original_model
+                            )
+                            st.session_state.image_caption = caption
+                    
+                    # Display original image caption
+                    if hasattr(st.session_state, 'image_caption'):
+                        st.markdown("### Image Description")
+                        st.markdown(st.session_state.image_caption)
+                        st.markdown("---")
+                    
+                    # Display GradCAM caption if available
+                    if hasattr(st.session_state, 'gradcam_caption'):
+                        st.markdown("### GradCAM Analysis")
+                        st.markdown(st.session_state.gradcam_caption)
+                        
+                        # Button to regenerate GradCAM caption
+                        if hasattr(st.session_state, 'current_overlay') and st.button("Regenerate GradCAM Caption"):
+                            with st.spinner("Reanalyzing GradCAM visualization..."):
+                                gradcam_caption = generate_gradcam_caption(
+                                    st.session_state.current_overlay,
+                                    st.session_state.finetuned_processor,
+                                    st.session_state.finetuned_model
+                                )
+                                st.session_state.gradcam_caption = gradcam_caption
+                                st.rerun()
+                    else:
+                        st.info("GradCAM caption not available. Visit the Detection tab to generate it.")
+        else:
+            st.info("Please upload and analyze an image in the Detection tab first.")
+    
+    # Tab 3: LLM Analysis
+    with tab3:
+        st.header("LLM Analysis")
+        
+        # Model Loading section
+        with st.expander("Load LLM Model", expanded=True):
             if not st.session_state.llm_model_loaded:
-                if st.button("📥 Load Vision LLM for Analysis", type="primary"):
+                if st.button("📥 Load Vision LLM", type="primary"):
                     # Load LLM model
                     model, tokenizer = load_llm_model()
                     if model is not None and tokenizer is not None:
@@ -794,263 +1022,23 @@ def main():
                         st.error("❌ Failed to load Vision LLM.")
             else:
                 st.success("✅ Vision LLM loaded and ready!")
-    
-    # Image upload section
-    with st.expander("Stage 2: Image Upload & Initial Detection", expanded=True):
-        st.subheader("Upload an Image")
-        
-        # Add alternative upload methods
-        upload_tab1, upload_tab2 = st.tabs(["File Upload", "URL Input"])
-        
-        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:
-                    # Simple direct approach - load the image directly
-                    image = Image.open(uploaded_file).convert("RGB")
-                    uploaded_image = image
-                    st.session_state.upload_method = "file"
-                except Exception as e:
-                    st.error(f"Error loading image: {str(e)}")
-                    import traceback
-                    st.error(traceback.format_exc())
-        
-        with upload_tab2:
-            url = st.text_input("Enter image URL:")
-            if url and url.strip():
-                try:
-                    import requests
-                    # Simplified URL handling with more robust approach
-                    headers = {
-                        'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36',
-                        'Accept': 'image/jpeg, image/png, image/*, */*',
-                        'Referer': 'https://huggingface.co/'
-                    }
-                    
-                    # Try three different methods to handle various API restrictions
-                    try_methods = True
-                    
-                    # Method 1: Direct requests
-                    if try_methods:
-                        try:
-                            response = requests.get(url, stream=True, headers=headers, timeout=10)
-                            if response.status_code == 200 and 'image' in response.headers.get('Content-Type', ''):
-                                try:
-                                    image = Image.open(io.BytesIO(response.content)).convert("RGB")
-                                    uploaded_image = image
-                                    st.session_state.upload_method = "url_direct"
-                                    try_methods = False
-                                    st.success("✅ Image loaded via direct request")
-                                except Exception as e:
-                                    st.warning(f"Direct method received data but couldn't process as image: {str(e)}")
-                            else:
-                                st.info(f"Direct method failed: Status {response.status_code}, trying alternative method...")
-                        except Exception as e:
-                            st.info(f"Direct method error: {str(e)}, trying alternative method...")
-                    
-                    # Method 2: Use Python's urllib as fallback
-                    if try_methods:
-                        try:
-                            import urllib.request
-                            from urllib.error import HTTPError
-                            
-                            opener = urllib.request.build_opener()
-                            opener.addheaders = [('User-agent', headers['User-Agent'])]
-                            urllib.request.install_opener(opener)
-                            
-                            with urllib.request.urlopen(url, timeout=10) as response:
-                                image_data = response.read()
-                                image = Image.open(io.BytesIO(image_data)).convert("RGB")
-                                uploaded_image = image
-                                st.session_state.upload_method = "url_urllib"
-                                try_methods = False
-                                st.success("✅ Image loaded via urllib")
-                        except HTTPError as e:
-                            st.info(f"urllib method failed: HTTP error {e.code}, trying next method...")
-                        except Exception as e:
-                            st.info(f"urllib method error: {str(e)}, trying next method...")
-                    
-                    # Method 3: Use a proxy service as last resort
-                    if try_methods:
-                        try:
-                            # This uses an image proxy service to bypass CORS issues
-                            # Only as last resort since it depends on external service
-                            proxy_url = f"https://images.weserv.nl/?url={url}"
-                            response = requests.get(proxy_url, stream=True, timeout=10)
-                            if response.status_code == 200:
-                                image = Image.open(io.BytesIO(response.content)).convert("RGB")
-                                uploaded_image = image
-                                st.session_state.upload_method = "url_proxy"
-                                try_methods = False
-                                st.success("✅ Image loaded via proxy service")
-                            else:
-                                st.error(f"All methods failed to load the image from URL. Last status: {response.status_code}")
-                        except Exception as e:
-                            st.error(f"All methods failed. Final error: {str(e)}")
-                    
-                    if not uploaded_image:
-                        st.error("Failed to load image using all available methods.")
-                    
-                except Exception as e:
-                    st.error(f"Error processing URL: {str(e)}")
-                    if st.session_state.debug:
-                        import traceback
-                        st.error(traceback.format_exc())
-        
-        # 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
+        # Chat Interface
+        if hasattr(st.session_state, 'current_image') and st.session_state.llm_model_loaded:
+            st.subheader("Deepfake Analysis Chat")
             
-            # 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
-                        input_tensor, original_image, face_box = preprocess_image_xception(image)
-                        
-                        if input_tensor is None:
-                            st.error("Failed to preprocess image. Please try another image.")
-                            st.stop()
-                        
-                        # 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)
-                        
-                        # Forward pass with proper error handling
-                        try:
-                            with torch.no_grad():
-                                logits = model(input_tensor)
-                                probabilities = torch.softmax(logits, dim=1)[0]
-                                pred_class = torch.argmax(probabilities).item()
-                                confidence = probabilities[pred_class].item()
-                                
-                                # Explicit class mapping - adjust if needed based on your model
-                                pred_label = "Fake" if pred_class == 0 else "Real"
-                        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%})")
-                            
-                            # 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:
-                            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
-                                    
-                                    # Display the caption directly here as well for immediate feedback
-                                    st.markdown("### GradCAM Analysis")
-                                    st.markdown(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')):
-        with st.expander("Image Analysis Summary", expanded=True):
-            # Display images and analysis in organized layout
+            # Display images
             col1, col2 = st.columns([1, 2])
-            
             with col1:
-                # Display original image
                 st.image(st.session_state.current_image, caption="Original Image", width=300)
-                # Display GradCAM overlay
                 if hasattr(st.session_state, 'current_overlay'):
                     st.image(st.session_state.current_overlay, caption="GradCAM Visualization", width=300)
             
             with col2:
-                # Image description
-                if hasattr(st.session_state, 'image_caption'):
-                    st.markdown("### Image Description")
-                    st.markdown(st.session_state.image_caption)
-                    st.markdown("---")
-                
-                # GradCAM analysis
-                if hasattr(st.session_state, 'gradcam_caption'):
-                    st.markdown("### GradCAM Analysis")
-                    st.markdown(st.session_state.gradcam_caption)
-                    st.markdown("---")
-                else:
-                    st.warning("GradCAM caption not found in session state.")
-    
-    # LLM Analysis section - AFTER Image Analysis Summary
-    with st.expander("Stage 3: Detailed Analysis with Vision LLM", expanded=False):
-        if hasattr(st.session_state, 'current_image') and st.session_state.llm_model_loaded:
-            st.subheader("Detailed Deepfake Analysis")
+                # Display detection result if available
+                if hasattr(st.session_state, 'current_pred_label'):
+                    st.markdown("### Detection Result")
+                    st.markdown(f"**Classification:** {st.session_state.current_pred_label} (Confidence: {st.session_state.current_confidence:.2%})")
             
             # Display chat history
             for i, (question, answer) in enumerate(st.session_state.chat_history):
@@ -1058,6 +1046,17 @@ def main():
                 st.markdown(f"**Answer:** {answer}")
                 st.markdown("---")
             
+            # Custom instruction
+            use_custom_instructions = st.toggle("Enable Custom Instructions", value=False, help="Toggle to enable/disable custom instructions")
+            if use_custom_instructions:
+                custom_instruction = st.text_area(
+                    "Custom Instructions (Advanced)",
+                    value="Specify your preferred style of explanation (e.g., 'Provide technical, detailed explanations' or 'Use simple, non-technical language'). You can also specify what aspects of the image to focus on.",
+                    help="Add specific instructions for the analysis"
+                )
+            else:
+                custom_instruction = ""
+            
             # Include both captions in the prompt if available
             caption_text = ""
             if hasattr(st.session_state, 'image_caption'):
@@ -1142,11 +1141,11 @@ def main():
                     
                 except Exception as e:
                     st.error(f"Error during LLM analysis: {str(e)}")
-                    
-        elif not hasattr(st.session_state, 'current_image'):
-            st.warning("⚠️ Please upload an image and complete the initial detection first.")
         else:
-            st.warning("⚠️ Please load the Vision LLM to perform detailed analysis.")
+            if not hasattr(st.session_state, 'current_image'):
+                st.warning("⚠️ Please upload an image in the Detection tab first.")
+            else:
+                st.warning("⚠️ Please load the Vision LLM to perform detailed analysis.")
     
     # Footer
     st.markdown("---")