Update app.py

app.py

@@ -441,14 +441,25 @@ def generate_gradcam_caption(image, processor, model, max_length=60):
         # Decode the output
         caption = processor.decode(output[0], skip_special_tokens=True)
         
-        # Format into structured analysis
+        # Split the caption into activation levels
+        caption_parts = caption.split("activation")
+        formatted_parts = []
+        
+        for part in caption_parts[1:]:  # Skip the first empty part
+            if ":" in part:
+                level, description = part.split(":", 1)
+                formatted_parts.append(f"**{level.strip()} activation**:{description.strip()}")
+        
+        # Format into structured analysis with better spacing
         structured_output = f"""
-**Heatmap Analysis**: {caption}
+**Heatmap Analysis**:
+
+{chr(10).join(formatted_parts)}
 
-**Key Observations**:
-- The red/yellow regions indicate areas of high importance in the detection process
-- The green/cyan areas show regions of medium importance
-- The blue/dark blue regions represent features with lower impact on the model's decision
+**Legend**:
+🔴 Red/Yellow = High importance regions
+🟡 Green/Cyan = Medium importance regions
+🔵 Blue/Dark = Low importance regions
 """
         return structured_output.strip()
         
@@ -764,14 +775,15 @@ 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):
-            # Display image, captions, and results in organized layout with proper formatting
+            # Display images and analysis in organized layout
             col1, col2 = st.columns([1, 2])
             
             with col1:
-                # Display original image and overlay side by side with controlled size
+                # Display original image
                 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 Overlay", 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