Update app.py

app.py

@@ -1,7 +1,8 @@
 import os
 import streamlit as st
-from PIL import Image
+from PIL import Image, ImageDraw
 import torch
+import re
 from transformers import PaliGemmaProcessor, PaliGemmaForConditionalGeneration
 
 # Get Hugging Face API key from environment variables
@@ -46,7 +47,7 @@ if uploaded_file:
     # User input for task selection
     task = st.selectbox(
         "Select a task:",
-        ["Generate a caption", "Answer a question", "Detect objects", "Generate segmentation"]
+        ["Generate a caption", "Answer a question", "Detect objects", "Segment"]
     )
 
     # User prompt
@@ -62,4 +63,22 @@ if uploaded_file:
                 generation = generation[0][input_len:]  # Remove input tokens from output
                 answer = processor.decode(generation, skip_special_tokens=True)
 
-            st.success(f"✅ Result: {answer}")
+            # Function to highlight <segXXX> values in yellow
+            def highlight_segments(text):
+                return re.sub(r"(<seg\d+>)", r'<span style="background-color:yellow; padding:3px; border-radius:3px;">\1</span>', text)
+
+            # Display output based on task selection
+            if task == "Segment":
+                highlighted_answer = highlight_segments(answer)
+                st.markdown(f'<p style="font-size:16px;">✅ Result: {highlighted_answer}</p>', unsafe_allow_html=True)
+            else:
+                st.success(f"✅ Result: {answer}")
+
+            # Draw red bounding box for object detection
+            if task == "Detect objects":
+                draw = ImageDraw.Draw(image)
+                # Placeholder: Replace with actual detection logic
+                boxes = [(50, 50, 200, 200), (250, 100, 400, 300)]  # Example boxes
+                for box in boxes:
+                    draw.rectangle(box, outline="red", width=3)
+                st.image(image, caption="Detected Objects", use_container_width=True)