feat: update UI for image captioning

app.py

@@ -1,10 +1,8 @@
 import streamlit as st
 from PIL import Image, ImageOps
+from internal.api import APIClient
 
-import requests
-
-API_URL = 'https://pic-gai.up.railway.app'
-# API_URL = 'http://localhost:8000'
+client = APIClient("http://localhost:3000")
 
 def gallery(column, images):
     groups = []
@@ -18,94 +16,83 @@ def gallery(column, images):
 
 st.title('CollageAI')
 
-# Input field for user prompt
-user_prompt = st.text_area(
-    "Describe the design you'd like to create:",
-    placeholder="For our anniversary, I want to write a card to my partner to celebrate our love and share all the things I adore about them."
-)
-
-uploaded_images = st.file_uploader("Choose photos", accept_multiple_files=True)
-if uploaded_images:
-    images = [Image.open(image) for image in uploaded_images]
-    images = [ImageOps.exif_transpose(image) for image in images]
-    gallery(4, images)
-
-# pick number of photos
-photos_count = len(uploaded_images)
-
-# Submit buttons for templates and stickers
-generate_button = st.button('Generate')
+user_images = None
+user_prompt = None
+uploaded_images = []
+image_captions_dict = {}
+submitted = False
+
+with st.form("user_input_form"):
+    user_images = st.file_uploader("Choose your photos", accept_multiple_files=True)
+    user_prompt = st.text_area(
+        "Describe the design you'd like to create:",
+        placeholder="For our anniversary, I want to write a card to my partner to celebrate our love and share all the things I adore about them."
+    )
+    submitted = st.form_submit_button("Generate")
+
+# Check form
+if submitted:
+    if user_images:
+        with st.container():
+            with st.spinner('Uploading images...'):
+                try:
+                    uploaded_images = client.upload_images(user_images)
+                except Exception as e:
+                    st.error(f"Error uploading images: {e}")
+
+            # Display the photo gallery
+            st.subheader('Your photos:')
+            images = [Image.open(image) for image in user_images]
+            images = [ImageOps.exif_transpose(image) for image in images]
+            gallery(4, images)
+    else:
+        st.warning('Please upload at least one image before submitting.')
 
-if generate_button:
     if user_prompt:
-        # Prepare the params with the user prompt
-        params = {
-            'prompt': user_prompt,
-            'photos_count': photos_count
-        }
-
-        # remove empty params
-        params = {k: v for k, v in params.items() if v is not None}
+        if uploaded_images:
+            # Analysis
+            with st.spinner('Analyzing prompt...'):
+                try:
+                    analysis = client.analyze_prompt(user_prompt, uploaded_images)
+                    keywords = analysis.get("keywords")
+                    captions = analysis.get("captions")
+                    if captions:
+                        st.subheader('Captions of your photos')
+                        st.write(captions)
+
+                    if keywords:
+                        st.subheader('Keywords based on your photos and prompt')
+                        st.write(keywords)
+                    else:
+                        st.warning('No keywords were generated. Please try again with a different prompt.')
 
-        st.markdown("---")
+                except Exception as e:
+                    st.error(f"Error analyzing prompt: {e}")
 
         # Templates
         with st.container():
-            # Define the FastAPI server URL for templates
-            url = f"{API_URL}/api/templates"
-
             with st.spinner('Generating templates...'):
-                # Make a request to the FastAPI server
-                response = requests.get(url, params=params)
-
-                # Display the response in the appropriate output container
-                if response.status_code == 200:
-                    templates = response.json().get('result', [])
-                    image_urls = [template.get('image_medium') for template in templates]
-                    
-                    if image_urls:
-                        st.subheader('Generated templates')
-                        gallery(4, image_urls[:8])
+                try:
+                    template_image_urls = client.suggest_templates(user_prompt, uploaded_images)
+                    if template_image_urls:
+                        st.subheader('Template suggestions')
+                        gallery(4, template_image_urls[:8])
                     else:
                         st.warning('No images were generated. Please try again with a different prompt.')
-                else:
-                    st.error(f"Error: {response.status_code}")
+                except Exception as e:
+                    st.error(f"Error generating templates: {e}")
 
+        # Stickers
         with st.container():
-            # Define the FastAPI server URL for templates
-            url = f"{API_URL}/api/stickers"
-
             with st.spinner('Generating stickers...'):
-                # Make a request to the FastAPI server
-                response = requests.get(url, params=params)
-
-                # Display the response in the appropriate output container
-                if response.status_code == 200:
-                    stickers = response.json().get('result', [])
-                    image_urls = [sticker.get('image_url') for sticker in stickers]
-                    
-                    if image_urls:
-                        st.subheader('Generated stickers')
-                        gallery(4, image_urls[:8])
+                try:
+                    sticker_image_urls = client.suggest_stickers(user_prompt, uploaded_images)
+                    if sticker_image_urls:
+                        st.subheader('Stickers suggestions')
+                        gallery(4, sticker_image_urls[:8])
                     else:
                         st.warning('No images were generated. Please try again with a different prompt.')
-                else:
-                    st.error(f"Error: {response.status_code}")
-        
-        # Keywords
-        with st.container():
-            # Define the FastAPI server URL for keywords
-            url = f"{API_URL}/api/analyze_prompt"
-
-            # Make a request to the FastAPI server
-            response = requests.get(url, params=params)
-
-            # Display the response in the appropriate output container
-            if response.status_code == 200:
-                st.subheader('Keywords based on prompt')
-                keywords = response.json().get('keywords', [])
-                st.write(keywords)
-            else:
-                st.error(f"Error: {response.status_code}")
+                except Exception as e:
+                    st.error(f"Error generating stickers: {e}")
     else:
         st.warning('Please enter a prompt before submitting.')

internal/api.py

@@ -0,0 +1,44 @@
+import requests
+
+class APIClient():
+    def __init__(self, host):
+        self.host = host
+
+    def url_with_path(self, path):
+        return f"{self.host}{path}"
+
+    def upload_images(self, images):
+        uploaded_images = []
+        url = self.url_with_path("/api/upload_images")
+        images_metadata = ([("images", (image.name, image)) for image in images])
+        print("Uploading images...")
+        print(f"url: {url}")
+        print(f"metadata: {images_metadata}")
+        response = requests.post(url, files=images_metadata)
+        uploaded_images = response.json().get('result', [])
+        return uploaded_images
+        
+    def suggest_templates(self, prompt, images):
+        url = self.url_with_path("/api/templates")
+        image_urls = [image.get('image_url') for image in images]
+
+        response = requests.get(url, json={'prompt': prompt, 'image_urls': image_urls})
+        templates = response.json().get('result', [])
+        template_image_urls = [template.get('image_medium') for template in templates]
+        return template_image_urls
+
+    def suggest_stickers(self, prompt, images):
+        url = self.url_with_path("/api/stickers")
+        image_urls = [image.get('image_url') for image in images]
+
+        response = requests.get(url, json={'prompt': prompt, 'image_urls': image_urls})
+        stickers = response.json().get('result', [])
+        sticker_image_urls = [sticker.get('image_url') for sticker in stickers]
+        return sticker_image_urls
+
+    def analyze_prompt(self, prompt, images):
+        url = self.url_with_path("/api/analyze_prompt")
+        image_urls = [image.get('image_url') for image in images]
+
+        response = requests.get(url, json={'prompt': prompt, 'image_urls': image_urls})
+        return response.json()