Update app.py

app.py

@@ -442,33 +442,39 @@ def generate_gradcam_caption(image, processor, model, max_length=60):
         # Decode the output
         caption = processor.decode(output[0], skip_special_tokens=True)
         
-        # Print debug info
-        st.write("Debug - Raw caption:", caption)
-        
-        # Split and format the text differently
-        if "high activation" in caption.lower():
-            # Extract the parts for each activation level
-            high_part = caption[caption.lower().find("high activation"):].split("moderate")[0].strip()
-            moderate_part = "moderate" + caption.split("moderate")[1].split("low")[0].strip()
-            low_part = "low" + caption.split("low")[1].strip()
-            
-            formatted_text = f"""**High activation**:
-{high_part.replace('high activation :', '').strip()}
-
-**Moderate activation**:
-{moderate_part.replace('moderate activation :', '').strip()}
-
-**Low activation**:
-{low_part.replace('low activation :', '').strip()}"""
-            
-            return formatted_text
-        
-        # If we can't parse it properly, return the original caption
-        return caption
+        # Clean and format the text
+        # First, split by activation levels
+        text_parts = caption.lower().split("activation")
+        
+        # Initialize empty strings for each activation level
+        high_text = ""
+        moderate_text = ""
+        low_text = ""
+        
+        # Process each part
+        for part in text_parts:
+            if not part.strip():  # Skip empty parts
+                continue
+            if "high" in part:
+                high_text = part.split(":", 1)[1].strip() if ":" in part else part.strip()
+            elif "moderate" in part:
+                moderate_text = part.split(":", 1)[1].strip() if ":" in part else part.strip()
+            elif "low" in part:
+                low_text = part.split(":", 1)[1].strip() if ":" in part else part.strip()
+        
+        # Format the output with only the unique content for each level
+        formatted_text = ""
+        if high_text:
+            formatted_text += f"**High activation**:\n{high_text}\n\n"
+        if moderate_text:
+            formatted_text += f"**Moderate activation**:\n{moderate_text}\n\n"
+        if low_text:
+            formatted_text += f"**Low activation**:\n{low_text}"
+            
+        return formatted_text.strip()
         
     except Exception as e:
         st.error(f"Error analyzing GradCAM: {str(e)}")
-        st.write("Debug - Full error:", str(e))
         return "Error analyzing GradCAM visualization"
 
 # Function to generate caption for original image