Update app.py

app.py

@@ -3,6 +3,8 @@ import torch
 from PIL import Image
 import numpy as np
 from transformers import AutoProcessor, AutoModelForCausalLM
+from io import BytesIO
+import base64
 
 # Initialize Florence model
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -34,13 +36,12 @@ def generate_caption(image):
     parsed_answer = florence_processor.post_process_generation(
         generated_text,
         task="<MORE_DETAILED_CAPTION>",
-        image_size=(image.width, image.height)  # Use PIL image dimensions here
+        image_size=(image.width, image.height)
     )
 
-    prompt = parsed_answer["<MORE_DETAILED_CAPTION>"]
-    return prompt
+    return parsed_answer["<MORE_DETAILED_CAPTION>"]
 
-# Streamlit Interface
+# Streamlit UI
 st.title("Florence 2 Caption Generator")
 st.write("Upload an image to generate a caption:")
 
@@ -50,10 +51,29 @@ uploaded_image = st.file_uploader("Choose an Image", type=["jpg", "jpeg", "png"]
 # If an image is uploaded
 if uploaded_image is not None:
     image = Image.open(uploaded_image)
-    st.image(image, caption="Uploaded Image", use_container_width=True)  # <-- FIXED DEPRECATED PARAMETER
-    
+    st.image(image, caption="Uploaded Image", use_container_width=True)
+
     # Generate caption when button is pressed
     if st.button("Generate Caption"):
         caption = generate_caption(image)
         st.subheader("Generated Caption:")
-        st.write(caption)
+        st.write(caption)
+
+# ✅ API Mode: Handle API Requests
+st.experimental_set_query_params()  # Ensure Streamlit can handle query params
+
+def handle_api_request():
+    """Handle API request by checking URL query parameters."""
+    query_params = st.experimental_get_query_params()
+    
+    if "image" in query_params:
+        image_base64 = query_params["image"][0]  # Get Base64-encoded image
+        image_bytes = BytesIO(base64.b64decode(image_base64))
+        image = Image.open(image_bytes)
+
+        caption = generate_caption(image)
+        st.json({"caption": caption})  # Return JSON response
+
+# Check if API mode is enabled
+if "image" in st.experimental_get_query_params():
+    handle_api_request()