Update app.py

app.py

@@ -442,34 +442,19 @@ def generate_gradcam_caption(image, processor, model, max_length=60):
         # Decode the output
         caption = processor.decode(output[0], skip_special_tokens=True)
         
-        # Clean and format the text
-        # First, split by activation levels
-        text_parts = caption.lower().split("activation")
+        # 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 ""
         
-        # 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
+        # Format the output
         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}"
+        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()