Update app.py

app.py

@@ -197,17 +197,8 @@ This tool detects deepfakes using three AI models:
 temperature = 0.7
 max_tokens = 500
 
-# Custom instruction text area in sidebar
-use_custom_instructions = st.sidebar.toggle("Enable Custom Instructions", value=False, help="Toggle to enable/disable custom instructions")
-
-if use_custom_instructions:
-    custom_instruction = st.sidebar.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 = ""
+# Define empty custom_instruction to maintain compatibility
+custom_instruction = ""
 
 # ----- GradCAM Implementation for Xception -----
 class ImageDataset(torch.utils.data.Dataset):
@@ -1027,120 +1018,124 @@ def main():
         if hasattr(st.session_state, 'current_image') and st.session_state.llm_model_loaded:
             st.subheader("Deepfake Analysis Chat")
             
-            # Display images
-            col1, col2 = st.columns([1, 2])
-            with col1:
-                st.image(st.session_state.current_image, caption="Original Image", width=300)
-                if hasattr(st.session_state, 'current_overlay'):
-                    st.image(st.session_state.current_overlay, caption="GradCAM Visualization", width=300)
-            
-            with col2:
-                # 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 reference images in a sidebar-like column
+            col_images, col_chat = st.columns([1, 3])
             
-            # Display chat history
-            for i, (question, answer) in enumerate(st.session_state.chat_history):
-                st.markdown(f"**Question {i+1}:** {question}")
-                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'):
-                caption_text += f"\n\nImage Description:\n{st.session_state.image_caption}"
+            with col_images:
+                st.write("#### Reference Images")
+                st.image(st.session_state.current_image, caption="Original", use_column_width=True)
                 
-            if hasattr(st.session_state, 'gradcam_caption'):
-                caption_text += f"\n\nGradCAM Analysis:\n{st.session_state.gradcam_caption}"
-            
-            # Default question with option to customize
-            default_question = f"This image has been classified as {{pred_label}}. Analyze all the provided images (original, GradCAM visualization, and comparison) to determine if this is a deepfake. Focus on highlighted areas in the GradCAM visualization. Provide both a technical explanation for experts and a simple explanation for non-technical users."
-            
-            # User input for new question
-            new_question = st.text_area("Ask a question about the image:", value=default_question if not st.session_state.chat_history else "", height=100)
-            
-            # Analyze button and Clear Chat button in the same row
-            col1, col2 = st.columns([3, 1])
-            with col1:
-                analyze_button = st.button("🔍 Send Question", type="primary")
-            with col2:
-                clear_button = st.button("🗑️ Clear Chat History")
-            
-            if clear_button:
-                st.session_state.chat_history = []
-                st.rerun()
+                if hasattr(st.session_state, 'current_overlay'):
+                    st.image(st.session_state.current_overlay, caption="GradCAM", use_column_width=True)
                 
-            if analyze_button and new_question:
-                try:
-                    # Add caption info if it's the first question
-                    if not st.session_state.chat_history:
-                        full_question = new_question + caption_text
-                    else:
-                        full_question = new_question
+                if hasattr(st.session_state, 'comparison_image'):
+                    st.image(st.session_state.comparison_image, caption="Comparison", use_column_width=True)
                     
-                    result = analyze_image_with_llm(
-                        st.session_state.current_image,
-                        st.session_state.current_overlay,
-                        st.session_state.current_face_box,
-                        st.session_state.current_pred_label,
-                        st.session_state.current_confidence,
-                        full_question,
-                        st.session_state.llm_model,
-                        st.session_state.tokenizer,
-                        temperature=temperature,
-                        max_tokens=max_tokens,
-                        custom_instruction=custom_instruction
+                if hasattr(st.session_state, 'current_pred_label'):
+                    st.info(f"**Classification:** {st.session_state.current_pred_label} (Confidence: {st.session_state.current_confidence:.2%})")
+            
+            with col_chat:
+                # Display chat history
+                for i, (question, answer) in enumerate(st.session_state.chat_history):
+                    st.markdown(f"**Question {i+1}:** {question}")
+                    st.markdown(f"**Answer:** {answer}")
+                    st.markdown("---")
+                
+                # Custom instruction in the chat column
+                use_custom_instructions = st.toggle("Enable Custom Instructions", key="llm_custom_instructions", value=False)
+                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'):
+                    caption_text += f"\n\nImage Description:\n{st.session_state.image_caption}"
                     
-                    # Add to chat history
-                    st.session_state.chat_history.append((new_question, result))
-                    
-                    # Display the latest result too
-                    st.success("✅ Analysis complete!")
+                if hasattr(st.session_state, 'gradcam_caption'):
+                    caption_text += f"\n\nGradCAM Analysis:\n{st.session_state.gradcam_caption}"
+                
+                # Default question with option to customize
+                default_question = f"This image has been classified as {{pred_label}}. Analyze all the provided images (original, GradCAM visualization, and comparison) to determine if this is a deepfake. Focus on highlighted areas in the GradCAM visualization. Provide both a technical explanation for experts and a simple explanation for non-technical users."
+                
+                # User input for new question
+                new_question = st.text_area("Ask a question about the image:", value=default_question if not st.session_state.chat_history else "", height=100)
+                
+                # Analyze button and Clear Chat button in the same row
+                col1, col2 = st.columns([3, 1])
+                with col1:
+                    analyze_button = st.button("🔍 Send Question", type="primary")
+                with col2:
+                    clear_button = st.button("🗑️ Clear Chat History")
+                
+                if clear_button:
+                    st.session_state.chat_history = []
+                    st.rerun()
                     
-                    # Check if the result contains both technical and non-technical explanations
-                    if "Technical" in result and "Non-Technical" in result:
-                        try:
-                            # Split the result into technical and non-technical sections
-                            parts = result.split("Non-Technical")
-                            technical = parts[0]
-                            non_technical = "Non-Technical" + parts[1]
-                            
-                            # Display in two columns
-                            tech_col, simple_col = st.columns(2)
-                            with tech_col:
-                                st.subheader("Technical Analysis")
-                                st.markdown(technical)
-                            
-                            with simple_col:
-                                st.subheader("Simple Explanation")
-                                st.markdown(non_technical)
-                        except Exception as e:
-                            # Fallback if splitting fails
+                if analyze_button and new_question:
+                    try:
+                        # Add caption info if it's the first question
+                        if not st.session_state.chat_history:
+                            full_question = new_question + caption_text
+                        else:
+                            full_question = new_question
+                        
+                        result = analyze_image_with_llm(
+                            st.session_state.current_image,
+                            st.session_state.current_overlay,
+                            st.session_state.current_face_box,
+                            st.session_state.current_pred_label,
+                            st.session_state.current_confidence,
+                            full_question,
+                            st.session_state.llm_model,
+                            st.session_state.tokenizer,
+                            temperature=temperature,
+                            max_tokens=max_tokens,
+                            custom_instruction=custom_instruction
+                        )
+                        
+                        # Add to chat history
+                        st.session_state.chat_history.append((new_question, result))
+                        
+                        # Display the latest result too
+                        st.success("✅ Analysis complete!")
+                        
+                        # Check if the result contains both technical and non-technical explanations
+                        if "Technical" in result and "Non-Technical" in result:
+                            try:
+                                # Split the result into technical and non-technical sections
+                                parts = result.split("Non-Technical")
+                                technical = parts[0]
+                                non_technical = "Non-Technical" + parts[1]
+                                
+                                # Display in two columns
+                                tech_col, simple_col = st.columns(2)
+                                with tech_col:
+                                    st.subheader("Technical Analysis")
+                                    st.markdown(technical)
+                                
+                                with simple_col:
+                                    st.subheader("Simple Explanation")
+                                    st.markdown(non_technical)
+                            except Exception as e:
+                                # Fallback if splitting fails
+                                st.subheader("Analysis Result")
+                                st.markdown(result)
+                        else:
+                            # Just display the whole result
                             st.subheader("Analysis Result")
                             st.markdown(result)
-                    else:
-                        # Just display the whole result
-                        st.subheader("Analysis Result")
-                        st.markdown(result)
-                    
-                    # Rerun to update the chat history display
-                    st.rerun()
-                    
-                except Exception as e:
-                    st.error(f"Error during LLM analysis: {str(e)}")
+                        
+                        # Rerun to update the chat history display
+                        st.rerun()
+                        
+                    except Exception as e:
+                        st.error(f"Error during LLM analysis: {str(e)}")
         else:
             if not hasattr(st.session_state, 'current_image'):
                 st.warning("⚠️ Please upload an image in the Detection tab first.")