Update app.py

app.py

@@ -3,7 +3,7 @@ import torch
 import torch.nn as nn
 from torch.utils.data import DataLoader
 from torchvision import transforms
-from transformers import CLIPModel
+from transformers import CLIPModel, BlipProcessor, BlipForConditionalGeneration
 from transformers.models.clip import CLIPModel  
 from PIL import Image
 import numpy as np
@@ -26,8 +26,8 @@ st.set_page_config(
 )
 
 # Main title and description
-st.title("Advanced Deepfake Image Analyzer")
-st.markdown("Analyze images for deepfake manipulation with multi-stage analysis")
+st.title("Deepfake Image Analyser")
+st.markdown("Analyse images for deepfake manipulation")
 
 # Check for GPU availability
 def check_gpu():
@@ -40,52 +40,37 @@ def check_gpu():
         return False
 
 # Sidebar components
-st.sidebar.title("Options")
-
-# Temperature slider
-temperature = st.sidebar.slider(
-    "Temperature", 
-    min_value=0.1, 
-    max_value=1.0, 
-    value=0.7, 
-    step=0.1,
-    help="Higher values make output more random, lower values more deterministic"
-)
+st.sidebar.title("About")
+st.sidebar.markdown("""
+This tool detects deepfakes using four AI models:
+- **CLIP**: Initial Real/Fake classification
+- **GradCAM**: Highlights suspicious regions
+- **BLIP**: Describes image content
+- **Llama 3.2**: Explains potential manipulations
+
+### Quick Start
+1. **Load Models** - Start with CLIP, add others as needed
+2. **Upload Image** - View classification and heat map
+3. **Analyze** - Get explanations and ask questions
+
+*GPU recommended for better performance*
+""")
 
-# Max response length slider
-max_tokens = st.sidebar.slider(
-    "Maximum Response Length", 
-    min_value=100, 
-    max_value=1000, 
-    value=500, 
-    step=50,
-    help="The maximum number of tokens in the response"
-)
+# Fixed values for temperature and max tokens
+temperature = 0.7
+max_tokens = 500
 
 # Custom instruction text area in sidebar
-custom_instruction = st.sidebar.text_area(
-    "Custom Instructions (Advanced)",
-    value="Focus on analyzing the highlighted regions from the GradCAM visualization. Examine facial inconsistencies, lighting irregularities, and other artifacts visible in the heat map.",
-    help="Add specific instructions for the LLM analysis"
-)
+use_custom_instructions = st.sidebar.toggle("Enable Custom Instructions", value=False, help="Toggle to enable/disable custom instructions")
 
-# About section in sidebar
-st.sidebar.markdown("---")
-st.sidebar.subheader("About")
-st.sidebar.markdown("""
-This analyzer performs multi-stage detection:
-1. **Initial Detection**: CLIP-based classifier
-2. **GradCAM Visualization**: Highlights suspicious regions
-3. **LLM Analysis**: Fine-tuned Llama 3.2 Vision provides detailed explanations
-
-The system looks for:
-- Facial inconsistencies
-- Unnatural movements
-- Lighting issues
-- Texture anomalies
-- Edge artifacts
-- Blending problems
-""")
+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 = ""
 
 # ----- GradCAM Implementation -----
 
@@ -341,15 +326,19 @@ def save_comparison(image, cam, overlay, face_box=None):
 @st.cache_resource
 def load_clip_model():
     with st.spinner("Loading CLIP model for GradCAM..."):
-        model = CLIPModel.from_pretrained("openai/clip-vit-large-patch14")
-        
-        # Apply a simple classification head
-        model.classification_head = nn.Linear(1024, 2)
-        model.classification_head.weight.data.normal_(mean=0.0, std=0.02)
-        model.classification_head.bias.data.zero_()
-        
-        model.eval()
-        return model
+        try:
+            model = CLIPModel.from_pretrained("openai/clip-vit-large-patch14")
+            
+            # Apply a simple classification head
+            model.classification_head = nn.Linear(1024, 2)
+            model.classification_head.weight.data.normal_(mean=0.0, std=0.02)
+            model.classification_head.bias.data.zero_()
+            
+            model.eval()
+            return model
+        except Exception as e:
+            st.error(f"Error loading CLIP model: {str(e)}")
+            return None
 
 def get_target_layer_clip(model):
     """Get the target layer for GradCAM"""
@@ -417,6 +406,105 @@ def process_image_with_gradcam(image, model, device, pred_class):
             comparison = save_comparison(original_image, default_cam, overlay, face_box)
             return default_cam, overlay, comparison, face_box
 
+# ----- BLIP Image Captioning -----
+
+# Function to load BLIP captioning models
+@st.cache_resource
+def load_blip_models():
+    with st.spinner("Loading BLIP captioning models..."):
+        try:
+            # Load original BLIP model for general image captioning
+            original_processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-large")
+            original_model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-large")
+            
+            # Load fine-tuned BLIP model for GradCAM analysis
+            finetuned_processor = BlipProcessor.from_pretrained("saakshigupta/deepfake-blip-large")
+            finetuned_model = BlipForConditionalGeneration.from_pretrained("saakshigupta/deepfake-blip-large")
+            
+            return original_processor, original_model, finetuned_processor, finetuned_model
+        except Exception as e:
+            st.error(f"Error loading BLIP models: {str(e)}")
+            return None, None, None, None
+
+# Function to generate image caption using BLIP's VQA approach for GradCAM
+def generate_gradcam_caption(image, processor, model, max_length=60):
+    """
+    Generate a detailed analysis of GradCAM visualization using the fine-tuned BLIP model
+    """
+    try:
+        # Process image first
+        inputs = processor(image, return_tensors="pt")
+        
+        # Check for available GPU and move model and inputs
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        model = model.to(device)
+        inputs = {k: v.to(device) if hasattr(v, 'to') else v for k, v in inputs.items()}
+        
+        # Generate caption
+        with torch.no_grad():
+            output = model.generate(**inputs, max_length=max_length, num_beams=5)
+            
+        # Decode the output
+        caption = processor.decode(output[0], skip_special_tokens=True)
+        
+        # Extract descriptions using the full text
+        high_match = caption.split("high activation :")[1].split("moderate")[0] if "high activation :" in caption else ""
+        moderate_match = caption.split("moderate activation :")[1].split("low")[0] if "moderate activation :" in caption else ""
+        low_match = caption.split("low activation :")[1] if "low activation :" in caption else ""
+        
+        # Format the output
+        formatted_text = ""
+        if high_match:
+            formatted_text += f"**High activation**:\n{high_match.strip()}\n\n"
+        if moderate_match:
+            formatted_text += f"**Moderate activation**:\n{moderate_match.strip()}\n\n"
+        if low_match:
+            formatted_text += f"**Low activation**:\n{low_match.strip()}"
+            
+        return formatted_text.strip()
+        
+    except Exception as e:
+        st.error(f"Error analyzing GradCAM: {str(e)}")
+        return "Error analyzing GradCAM visualization"
+
+# Function to generate caption for original image
+def generate_image_caption(image, processor, model, max_length=75, num_beams=5):
+    """Generate a caption for the original image using the original BLIP model"""
+    try:
+        # Check for available GPU
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        model = model.to(device)
+        
+        # For original image, use unconditional captioning
+        inputs = processor(image, return_tensors="pt").to(device)
+        
+        # Generate caption
+        with torch.no_grad():
+            output = model.generate(**inputs, max_length=max_length, num_beams=num_beams)
+            
+        # Decode the output
+        caption = processor.decode(output[0], skip_special_tokens=True)
+        
+        # Format into structured description
+        structured_caption = f"""
+**Subject**: The image shows a person in a photograph.
+
+**Appearance**: {caption}
+
+**Background**: The background appears to be a controlled environment.
+
+**Lighting**: The lighting appears to be professional with even illumination.
+
+**Colors**: The image contains natural skin tones and colors typical of photography.
+
+**Notable Elements**: The facial features and expression are the central focus of the image.
+"""
+        return structured_caption.strip()
+            
+    except Exception as e:
+        st.error(f"Error generating caption: {str(e)}")
+        return "Error generating caption"
+
 # ----- Fine-tuned Vision LLM -----
 
 # Function to fix cross-attention masks
@@ -427,7 +515,6 @@ def fix_cross_attention_mask(inputs):
         new_mask = torch.ones((batch_size, seq_len, visual_features, num_tiles),
                             device=inputs['cross_attention_mask'].device)
         inputs['cross_attention_mask'] = new_mask
-        st.success("Fixed cross-attention mask dimensions")
     return inputs
 
 # Load model function
@@ -512,7 +599,7 @@ def analyze_image_with_llm(image, gradcam_overlay, face_box, pred_label, confide
 
 # Main app
 def main():
-    # Create placeholders for model state
+    # Initialize session state variables
     if 'clip_model_loaded' not in st.session_state:
         st.session_state.clip_model_loaded = False
         st.session_state.clip_model = None
@@ -522,10 +609,23 @@ def main():
         st.session_state.llm_model = None
         st.session_state.tokenizer = None
     
+    if 'blip_model_loaded' not in st.session_state:
+        st.session_state.blip_model_loaded = False
+        st.session_state.original_processor = None
+        st.session_state.original_model = None
+        st.session_state.finetuned_processor = None
+        st.session_state.finetuned_model = None
+    
+    # Initialize chat history
+    if 'chat_history' not in st.session_state:
+        st.session_state.chat_history = []
+    
     # Create expanders for each stage
     with st.expander("Stage 1: Model Loading", expanded=True):
-        # Button for loading CLIP model
-        clip_col, llm_col = st.columns(2)
+        st.write("Please load the models using the buttons below:")
+        
+        # Button for loading models
+        clip_col, blip_col, llm_col = st.columns(3)
         
         with clip_col:
             if not st.session_state.clip_model_loaded:
@@ -541,6 +641,23 @@ def main():
             else:
                 st.success("✅ CLIP model loaded and ready!")
         
+        with blip_col:
+            if not st.session_state.blip_model_loaded:
+                if st.button("📥 Load BLIP for Captioning", type="primary"):
+                    # Load BLIP models
+                    original_processor, original_model, finetuned_processor, finetuned_model = load_blip_models()
+                    if all([original_processor, original_model, finetuned_processor, finetuned_model]):
+                        st.session_state.original_processor = original_processor
+                        st.session_state.original_model = original_model
+                        st.session_state.finetuned_processor = finetuned_processor
+                        st.session_state.finetuned_model = finetuned_model
+                        st.session_state.blip_model_loaded = True
+                        st.success("✅ BLIP captioning models loaded successfully!")
+                    else:
+                        st.error("❌ Failed to load BLIP models.")
+            else:
+                st.success("✅ BLIP captioning models loaded and ready!")
+                
         with llm_col:
             if not st.session_state.llm_model_loaded:
                 if st.button("📥 Load Vision LLM for Analysis", type="primary"):
@@ -562,116 +679,221 @@ def main():
         uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
         
         if uploaded_file is not None:
-            # Display the uploaded image
-            image = Image.open(uploaded_file).convert("RGB")
-            st.image(image, caption="Uploaded Image", use_column_width=True)
-            
-            # Detect with CLIP model if loaded
-            if st.session_state.clip_model_loaded:
-                with st.spinner("Analyzing image with CLIP model..."):
-                    # Preprocess image for CLIP
-                    transform = transforms.Compose([
-                        transforms.Resize((224, 224)),
-                        transforms.ToTensor(),
-                        transforms.Normalize(mean=[0.48145466, 0.4578275, 0.40821073], std=[0.26862954, 0.26130258, 0.27577711]),
-                    ])
-                    
-                    # Create a simple dataset for the image
-                    dataset = ImageDataset(image, transform=transform, face_only=True)
-                    tensor, _, _, _, face_box, _ = dataset[0]
-                    tensor = tensor.unsqueeze(0)
-                    
-                    # Get device
-                    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-                    
-                    # Move model and tensor to device
-                    model = st.session_state.clip_model.to(device)
-                    tensor = tensor.to(device)
-                    
-                    # Forward pass
-                    with torch.no_grad():
-                        outputs = model.vision_model(pixel_values=tensor).pooler_output
-                        logits = model.classification_head(outputs)
-                        probs = torch.softmax(logits, dim=1)[0]
-                        pred_class = torch.argmax(probs).item()
-                        confidence = probs[pred_class].item()
-                        pred_label = "Fake" if pred_class == 1 else "Real"
-                    
-                    # Display results
-                    result_col1, result_col2 = st.columns(2)
-                    with result_col1:
-                        st.metric("Prediction", pred_label)
-                    with result_col2:
-                        st.metric("Confidence", f"{confidence:.2%}")
-                    
-                    # GradCAM visualization
-                    st.subheader("GradCAM Visualization")
-                    cam, overlay, comparison, detected_face_box = process_image_with_gradcam(
-                        image, model, device, pred_class
-                    )
-                    
-                    # Display GradCAM results
-                    st.image(comparison, caption="Original | CAM | Overlay", use_column_width=True)
-                    
-                    # Save results in session state for LLM analysis
-                    st.session_state.current_image = image
-                    st.session_state.current_overlay = overlay
-                    st.session_state.current_face_box = detected_face_box
-                    st.session_state.current_pred_label = pred_label
-                    st.session_state.current_confidence = confidence
-                    
-                    st.success("✅ Initial detection and GradCAM visualization complete!")
-            else:
-                st.warning("⚠️ Please load the CLIP model first to perform initial detection.")
+            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
+                        
+                        # Store caption but don't display it yet
+                
+                # Detect with CLIP model if loaded
+                if st.session_state.clip_model_loaded:
+                    with st.spinner("Analyzing image with CLIP model..."):
+                        # Preprocess image for CLIP
+                        transform = transforms.Compose([
+                            transforms.Resize((224, 224)),
+                            transforms.ToTensor(),
+                            transforms.Normalize(mean=[0.48145466, 0.4578275, 0.40821073], std=[0.26862954, 0.26130258, 0.27577711]),
+                        ])
+                        
+                        # Create a simple dataset for the image
+                        dataset = ImageDataset(image, transform=transform, face_only=True)
+                        tensor, _, _, _, face_box, _ = dataset[0]
+                        tensor = tensor.unsqueeze(0)
+                        
+                        # Get device
+                        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+                        
+                        # Move model and tensor to device
+                        model = st.session_state.clip_model.to(device)
+                        tensor = tensor.to(device)
+                        
+                        # Forward pass
+                        with torch.no_grad():
+                            outputs = model.vision_model(pixel_values=tensor).pooler_output
+                            logits = model.classification_head(outputs)
+                            probs = torch.softmax(logits, dim=1)[0]
+                            pred_class = torch.argmax(probs).item()
+                            confidence = probs[pred_class].item()
+                            pred_label = "Fake" if pred_class == 1 else "Real"
+                        
+                        # Display results
+                        with col2:
+                            st.markdown("### Detection Result")
+                            st.markdown(f"**Classification:** {pred_label} (Confidence: {confidence:.2%})")
+                        
+                        # GradCAM visualization
+                        st.subheader("GradCAM Visualization")
+                        cam, overlay, comparison, detected_face_box = process_image_with_gradcam(
+                            image, model, device, pred_class
+                        )
+                        
+                        # Display GradCAM results (controlled size)
+                        st.image(comparison, caption="Original | CAM | Overlay", width=700)
+                        
+                        # Generate caption for GradCAM overlay image if BLIP model is loaded
+                        if st.session_state.blip_model_loaded:
+                            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
+                                
+                                # Store caption but don't display it yet
+                        
+                        # Save results in session state for LLM analysis
+                        st.session_state.current_image = image
+                        st.session_state.current_overlay = overlay
+                        st.session_state.current_face_box = detected_face_box
+                        st.session_state.current_pred_label = pred_label
+                        st.session_state.current_confidence = confidence
+                        
+                        st.success("✅ Initial detection and GradCAM visualization complete!")
+                else:
+                    st.warning("⚠️ Please load the CLIP 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
+    
+    # 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
+            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("---")
     
-    # LLM Analysis section
+    # 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 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("---")
+            
+            # 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}"
+                
+            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 {st.session_state.current_pred_label}. Analyze the key features that led to this classification, focusing on the highlighted areas in the GradCAM visualization. Provide both a technical explanation for experts and a simple explanation for non-technical users."
-            question = st.text_area("Question/Prompt:", value=default_question, height=100)
-            
-            # Analyze button
-            if st.button("🔍 Perform Detailed Analysis", type="primary"):
-                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,
-                    question,
-                    st.session_state.llm_model,
-                    st.session_state.tokenizer,
-                    temperature=temperature,
-                    max_tokens=max_tokens,
-                    custom_instruction=custom_instruction
-                )
-                
-                # Display results
-                st.success("✅ Analysis complete!")
+            
+            # 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.experimental_rerun()
                 
-                # Check if the result contains both technical and non-technical explanations
-                if "Technical" in result and "Non-Technical" in result:
-                    # Split the result into technical and non-technical sections
-                    parts = result.split("Non-Technical")
-                    technical = parts[0]
-                    non_technical = "Non-Technical" + parts[1]
+            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
                     
-                    # Display in two columns
-                    col1, col2 = st.columns(2)
-                    with col1:
-                        st.subheader("Technical Analysis")
-                        st.markdown(technical)
+                    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)
+                    
+                    # Rerun to update the chat history display
+                    st.experimental_rerun()
+                    
+                except Exception as e:
+                    st.error(f"Error during LLM analysis: {str(e)}")
                     
-                    with col2:
-                        st.subheader("Simple Explanation")
-                        st.markdown(non_technical)
-                else:
-                    # Just display the whole result
-                    st.subheader("Analysis Result")
-                    st.markdown(result)
         elif not hasattr(st.session_state, 'current_image'):
             st.warning("⚠️ Please upload an image and complete the initial detection first.")
         else:
@@ -679,7 +901,6 @@ def main():
     
     # Footer
     st.markdown("---")
-    st.caption("Advanced Deepfake Image Analyzer")
 
 if __name__ == "__main__":
     main()