Update app.py

app.py

@@ -0,0 +1,170 @@
+import streamlit as st
+import os
+import google.generativeai as genai
+from PIL import Image
+from dotenv import load_dotenv
+from modules.model_loader import load_gemini_model, generate_caption_with_gemini, generate_detailed_description
+
+# Load environment variables
+load_dotenv()
+
+# Fetch API key securely
+api_key = os.getenv("GEMINI_API_KEY")
+# Configure Gemini with your API Key
+genai.configure(api_key=api_key)
+
+# Setup a session state counter
+if "queries_done" not in st.session_state:
+    st.session_state.queries_done = 0
+
+if "user_api_key" not in st.session_state:
+    st.session_state.user_api_key = None
+
+# Add a state variable for the dynamic prompt
+if "dynamic_prompt" not in st.session_state:
+    st.session_state.dynamic_prompt = ""
+
+# Function to reset the dynamic prompt
+def reset_dynamic_prompt():
+    st.session_state.dynamic_prompt = ""
+
+# Streamlit App
+st.title("🖼️ Dynamic Image Caption Generator (Gemini Pro Vision)")
+st.write("Upload an image and generate creative descriptions based on your needs.")
+
+# If more than 2 queries, ask for user's API key
+if st.session_state.queries_done >= 2 and not st.session_state.user_api_key:
+    user_api_key = st.text_input("Enter your Gemini API Key to continue:", type="password")
+    if user_api_key:
+        st.session_state.user_api_key = user_api_key
+        genai.configure(api_key=user_api_key)
+    else:
+        st.warning("Please enter your Gemini API Key to continue using the app.")
+        st.stop()
+else:
+    # For first two requests, use your own key
+    if st.session_state.queries_done < 2:
+        genai.configure(api_key=api_key)
+    elif st.session_state.user_api_key:
+        genai.configure(api_key=st.session_state.user_api_key)
+
+# Upload image
+uploaded_image = st.file_uploader("Upload an Image", type=["jpg", "jpeg", "png"])
+
+if uploaded_image is not None:
+    image = Image.open(uploaded_image).convert("RGB")
+    # Use use_container_width instead of use_column_width
+    st.image(image, caption="Uploaded Image", use_container_width=True)
+    
+    # Create tabs for different ways to interact
+    tab1, tab2 = st.tabs(["Quick Presets", "Ask Anything"])
+    
+    with tab1:
+        # Prompt selection options
+        prompt_options = {
+            "basic": "Basic Description",
+            "chain_of_thought": "Detailed Analysis",
+            "story": "Creative Story",
+            "emotional": "Emotional Analysis",
+            "object": "Object Detection",
+            "context": "Contextual Description",
+            "action": "Action Description"
+        }
+
+        # Prompt selection
+        prompt_type = st.selectbox(
+            "Select description type",
+            options=list(prompt_options.keys()),
+            format_func=lambda x: prompt_options[x]
+        )
+        
+        if st.button("Generate Description", key="preset_button"):
+            with st.spinner("Processing image..."):
+                # Increment the counter BEFORE processing
+                st.session_state.queries_done += 1
+                
+                # Load model
+                gemini_model = load_gemini_model()
+                
+                # Generate caption
+                with st.spinner("Generating basic caption..."):
+                    caption = generate_caption_with_gemini(image, gemini_model)
+                    
+                    # Show caption
+                    st.subheader("✨ Basic Caption")
+                    st.write(caption)
+                
+                # Generate detailed description based on selected prompt type
+                with st.spinner(f"Generating {prompt_options[prompt_type]}..."):
+                    detailed_description = generate_detailed_description(
+                        image, 
+                        gemini_model, 
+                        prompt_type
+                    )
+                    
+                    # Show the result with appropriate header
+                    st.subheader(f"✨ {prompt_options[prompt_type]}")
+                    st.write(detailed_description)
+                    
+                    # Add download button for the generated text
+                    st.download_button(
+                        label="Download Description",
+                        data=detailed_description,
+                        file_name="image_description.txt",
+                        mime="text/plain",
+                        key="download1"
+                    )
+                
+                # Show usage count
+                st.info(f"You have used {st.session_state.queries_done} out of 2 free queries.")
+    
+    with tab2:
+        # Dynamic prompt box for custom questions
+        st.subheader("Ask anything about this image")
+        dynamic_prompt = st.text_area(
+            "Enter your question or what you'd like to know about the image",
+            value=st.session_state.dynamic_prompt,
+            placeholder="Examples:\n- Write a poem about this image\n- Explain what's happening in this scene\n- What emotion does this image evoke?\n- Describe this as if you were a detective",
+            key="dynamic_prompt_input"
+        )
+        
+        if st.button("Get Answer", key="dynamic_button"):
+            if dynamic_prompt:
+                # Save the current prompt before processing
+                current_prompt = dynamic_prompt
+                
+                # Clear the prompt for next use
+                reset_dynamic_prompt()
+                
+                # Increment the counter BEFORE processing
+                st.session_state.queries_done += 1
+                
+                with st.spinner("Processing your request..."):
+                    # Load model
+                    gemini_model = load_gemini_model()
+                    
+                    # Generate response directly from the image and custom prompt
+                    response = generate_detailed_description(
+                        image,
+                        gemini_model,
+                        "custom",
+                        current_prompt
+                    )
+                    
+                    # Display the response
+                    st.subheader("✨ Response")
+                    st.write(response)
+                    
+                    # Add download button
+                    st.download_button(
+                        label="Download Response",
+                        data=response,
+                        file_name="image_response.txt",
+                        mime="text/plain",
+                        key="download2"
+                    )
+                
+                # Show usage count
+                st.info(f"You have used {st.session_state.queries_done} out of 2 free queries.")
+            else:
+                st.warning("Please enter a question or prompt first")