Update app.py

app.py

@@ -442,18 +442,33 @@ def generate_gradcam_caption(image, processor, model, max_length=60):
         # Decode the output
         caption = processor.decode(output[0], skip_special_tokens=True)
         
-        # Split by activation levels and format
-        parts = caption.split(" activation : ")
-        if len(parts) > 1:  # Make sure we have activation levels to split
-            formatted_text = parts[0].replace("high", "**High activation**:\n") + "\n\n"
-            formatted_text += "**Moderate activation**:\n" + parts[1].split("moderate")[1] + "\n\n"
-            formatted_text += "**Low activation**:\n" + parts[2].split("low")[1]
-            return formatted_text.strip()
+        # 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
         
     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