Update app.py

app.py

@@ -360,39 +360,65 @@ def generate_gradcam_caption(image, processor, model, max_length=60):
     Generate a detailed analysis of GradCAM visualization using the fine-tuned BLIP model
     """
     try:
+        st.write("Debug: Starting GradCAM caption generation")
         # 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"
+        st.write(f"Debug: Using device: {device}")
         model = model.to(device)
         inputs = {k: v.to(device) if hasattr(v, 'to') else v for k, v in inputs.items()}
         
         # Generate caption
+        st.write("Debug: Generating 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)
+        st.write(f"Debug: Raw caption: {caption}")
         
-        # 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()}"
+        # Try to parse the caption based on different possible formats
+        try:
+            # Original format with "high activation:" etc.
+            formatted_text = ""
+            if "high activation :" in caption:
+                high_match = caption.split("high activation :")[1].split("moderate")[0]
+                formatted_text += f"**High activation**:\n{high_match.strip()}\n\n"
+            
+            if "moderate activation :" in caption:
+                moderate_match = caption.split("moderate activation :")[1].split("low")[0]
+                formatted_text += f"**Moderate activation**:\n{moderate_match.strip()}\n\n"
             
+            if "low activation :" in caption:
+                low_match = caption.split("low activation :")[1]
+                formatted_text += f"**Low activation**:\n{low_match.strip()}"
+            
+            # If nothing was extracted using the original format, try alternative formats
+            if not formatted_text.strip():
+                st.write("Debug: Trying alternative format parsing")
+                
+                # Check for newer format that might be in the Xception model
+                if ":" in caption:
+                    parts = caption.split(":")
+                    if len(parts) > 1:
+                        formatted_text = f"**GradCAM Analysis**:\n{parts[1].strip()}"
+                else:
+                    # As a fallback, just use the entire caption
+                    formatted_text = f"**GradCAM Analysis**:\n{caption.strip()}"
+        except Exception as parsing_error:
+            st.write(f"Debug: Error parsing caption format: {str(parsing_error)}")
+            # Use the entire caption as is
+            formatted_text = f"**GradCAM Analysis**:\n{caption.strip()}"
+            
+        st.write(f"Debug: Formatted caption complete. Length: {len(formatted_text)}")
         return formatted_text.strip()
         
     except Exception as e:
         st.error(f"Error analyzing GradCAM: {str(e)}")
+        import traceback
+        st.error(traceback.format_exc())
         return "Error analyzing GradCAM visualization"
 
 # Function to generate caption for original image
@@ -928,12 +954,18 @@ def main():
                             # 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..."):
+                                    st.write("Debug: Starting GradCAM analysis")
                                     gradcam_caption = generate_gradcam_caption(
                                         overlay,
                                         st.session_state.finetuned_processor,
                                         st.session_state.finetuned_model
                                     )
                                     st.session_state.gradcam_caption = gradcam_caption
+                                    st.write(f"Debug: Saved GradCAM caption to session state, length: {len(gradcam_caption) if gradcam_caption else 0}")
+                                    
+                                    # Display the caption directly here as well for immediate feedback
+                                    st.markdown("### GradCAM Analysis (Direct)")
+                                    st.markdown(gradcam_caption)
                         except Exception as e:
                             st.error(f"Error generating GradCAM: {str(e)}")
                             import traceback
@@ -957,6 +989,11 @@ def main():
     # 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):
+            st.write(f"Debug: Image caption exists: {hasattr(st.session_state, 'image_caption')}")
+            st.write(f"Debug: GradCAM caption exists: {hasattr(st.session_state, 'gradcam_caption')}")
+            if hasattr(st.session_state, 'gradcam_caption'):
+                st.write(f"Debug: GradCAM caption length: {len(st.session_state.gradcam_caption)}")
+            
             # Display images and analysis in organized layout
             col1, col2 = st.columns([1, 2])
             
@@ -979,6 +1016,8 @@ def main():
                     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):