Update app.py

app.py

@@ -1,15 +1,27 @@
 import streamlit as st
-import requests
 import base64
 import os
+import requests
+from PIL import Image
+from io import BytesIO
+
+# Function to compress and resize the image before base64 encoding
+def compress_and_resize_image(image, max_size=(1024, 1024), quality=85):
+    img = Image.open(image)
+    img.thumbnail(max_size)  # Resize image while maintaining aspect ratio
+    with BytesIO() as byte_io:
+        img.save(byte_io, format="JPEG", quality=quality)  # Save with reduced quality
+        byte_io.seek(0)
+        return byte_io
 
-# Function to convert image to base64
+# Function to convert uploaded image to base64
 def convert_image_to_base64(image):
-    image_bytes = image.read()
+    compressed_image = compress_and_resize_image(image)
+    image_bytes = compressed_image.read()
     encoded_image = base64.b64encode(image_bytes).decode("utf-8")
     return encoded_image
 
-# Function to generate a caption using Nebius API
+# Function to generate caption using Nebius API
 def generate_caption(encoded_image):
     API_URL = "https://api.studio.nebius.ai/v1/chat/completions"
     API_KEY = os.environ.get("NEBIUS_API_KEY")
@@ -20,7 +32,7 @@ def generate_caption(encoded_image):
     }
 
     payload = {
-        "model": "llava-hf/llava-1.5-7b-hf",
+        "model": "Qwen/Qwen2-VL-72B-Instruct",
         "messages": [
             {
                 "role": "system",
@@ -28,20 +40,19 @@ def generate_caption(encoded_image):
             },
             {
                 "role": "user",
-                "content": "write a detailed caption for this image"
+                "content": "Write a caption for this image"
+            },
+            {
+                "role": "user",
+                "content": f"data:image/png;base64,{encoded_image}"  # This is where the image is passed as base64 directly
             }
         ],
-        "image": {
-            "type": "image_url",
-            "image_url": {
-                "url": f"data:image/png;base64,{encoded_image}"
-            }
-        },
-        "temperature": 0.7
+        "temperature": 0
     }
 
+    # Send request to Nebius API
     response = requests.post(API_URL, headers=headers, json=payload)
-    
+
     if response.status_code == 200:
         result = response.json()
         caption = result.get("choices", [{}])[0].get("message", {}).get("content", "No caption generated.")
@@ -50,58 +61,32 @@ def generate_caption(encoded_image):
         st.error(f"API Error {response.status_code}: {response.text}")
         return None
 
-# Streamlit app
+# Streamlit app layout
 def main():
-    st.set_page_config(page_title="Image to Caption Converter", layout="centered", initial_sidebar_state="collapsed")
-    
-    # Gradient background style
-    st.markdown("""
-        <style>
-            body {
-                background: linear-gradient(135deg, #1e3c72, #2a5298);
-                color: white;
-                font-family: 'Arial', sans-serif;
-            }
-            .uploaded-image {
-                max-width: 100%;
-                border: 2px solid #ffffff;
-                border-radius: 10px;
-            }
-            .copy-button {
-                background-color: #ff8800;
-                color: white;
-                border: none;
-                border-radius: 5px;
-                padding: 10px 15px;
-                cursor: pointer;
-            }
-            .copy-button:hover {
-                background-color: #cc6b00;
-            }
-        </style>
-    """, unsafe_allow_html=True)
-
-    st.title("🖼️ Image to Caption Converter")
+    st.set_page_config(page_title="Image Caption Generator", layout="centered", initial_sidebar_state="collapsed")
+    st.title("🖼️ Image to Caption Generator")
 
     uploaded_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
+
     if uploaded_file:
         # Display the uploaded image
         st.image(uploaded_file, caption="Uploaded Image", use_container_width=True)
 
-        # Convert image to base64 and get caption
         if st.button("Generate Caption"):
+            # Convert the uploaded image to base64
             with st.spinner("Generating caption..."):
                 encoded_image = convert_image_to_base64(uploaded_file)
+
+                # Debugging: Ensure the encoded image is valid and not too large
+                st.write(f"Encoded image length: {len(encoded_image)} characters")
+
+                # Get the generated caption from the API
                 caption = generate_caption(encoded_image)
 
                 if caption:
                     st.subheader("Generated Caption:")
                     st.text_area("", caption, height=100, key="caption_area")
-
-                    # Copy button
-                    if st.button("Copy to Clipboard"):
-                        st.code(caption, language="text")
-                        st.success("Caption copied to clipboard!")
+                    st.success("Caption generated successfully!")
 
 if __name__ == "__main__":
     main()