Update app.py

app.py

@@ -23,20 +23,41 @@ def load_model():
 processor, model = load_model()
 
 # Streamlit UI
-st.title("🖼️ Image Q&A using PaliGemma")
+st.title("🖼️ Image Understanding with PaliGemma")
 
 uploaded_file = st.file_uploader("Upload an Image", type=["png", "jpg", "jpeg"])
 
 if uploaded_file:
     image = Image.open(uploaded_file).convert("RGB")
-    st.image(image, caption="Uploaded Image", use_column_width=True)
-    
-    question = st.text_input("Ask a question about the image:")
-    if question:
-        # Process the image and question
-        inputs = processor(text=question, images=image, return_tensors="pt")
-        with torch.no_grad():
-            output = model.generate(**inputs)
-        
-        answer = processor.batch_decode(output, skip_special_tokens=True)[0]
-        st.success(f"✅ Answer: {answer}")
+    st.image(image, caption="Uploaded Image", use_container_width=True)
+
+    # User selects the task
+    task = st.selectbox(
+        "Select a task:",
+        ["Generate a caption", "Answer a question", "Detect objects", "Generate segmentation"]
+    )
+
+    # User input for question/prompt
+    prompt = st.text_area("Enter a prompt (e.g., 'Describe the image' or 'What objects are present?')")
+
+    if st.button("Run"):
+        if prompt:
+            inputs = processor(text=prompt, images=image, return_tensors="pt")
+
+            with torch.no_grad():
+                output = model.generate(**inputs)
+
+            raw_output = processor.batch_decode(output, skip_special_tokens=False)[0]
+
+            # Handle different outputs
+            if task == "Generate a caption":
+                answer = raw_output
+            elif task == "Answer a question":
+                answer = raw_output
+            elif task == "Detect objects":
+                answer = f"Object bounding boxes: {raw_output}"
+            elif task == "Generate segmentation":
+                answer = f"Segmentation codes: {raw_output}"
+
+            st.success(f"✅ Result: {answer}")
+