Spaces:
Running
Running
| import streamlit as st | |
| from PIL import Image | |
| import google.generativeai as genai | |
| import io | |
| import os | |
| import base64 | |
| # Configure Google Gemini API (replace with your actual API key) | |
| GOOGLE_API_KEY = os.environ.get("GOOGLE_API_KEY") | |
| if not GOOGLE_API_KEY: | |
| st.error("Please set the GOOGLE_API_KEY environment variable.") | |
| st.stop() | |
| genai.configure(api_key=GOOGLE_API_KEY) | |
| # Function to generate the modified image | |
| def generate_modified_image(uploaded_image, background_description): | |
| try: | |
| model = genai.GenerativeModel('gemini-1.5-flash') | |
| img = Image.open(uploaded_image) | |
| prompt_parts = [ | |
| "You are an AI that can modify the background of an image based on a text description.", | |
| "Here is the image:", | |
| img, # Pass the PIL Image object directly | |
| f"Modify the background of this image to: '{background_description}'. Be creative and make the new background look realistic and integrated with the foreground.", | |
| "Output only the modified image. Ensure the output is a valid image format (like PNG or JPEG) encoded as bytes." | |
| ] | |
| response = model.generate_content(prompt_parts, stream=False) | |
| response.resolve() | |
| if response and hasattr(response, 'parts') and len(response.parts) > 0: | |
| for part in response.parts: | |
| try: | |
| # Try to get the image data directly if it's a blob | |
| if hasattr(part, 'blob'): | |
| image_bytes = part.blob.data | |
| return Image.open(io.BytesIO(image_bytes)) | |
| elif hasattr(part, 'text'): | |
| st.warning(f"Received text response instead of image: {part.text}") | |
| else: | |
| st.warning(f"Unexpected part type in response: {part}") | |
| except Exception as part_err: | |
| st.error(f"Error processing response part: {part_err}") | |
| st.error("No valid image data found in the response.") | |
| return None | |
| else: | |
| st.error("Failed to get a valid response from the model.") | |
| return None | |
| except Exception as e: | |
| st.error(f"An error occurred during generation: {e}") | |
| return None | |
| # Streamlit web app | |
| st.title("Image Background Modifier") | |
| uploaded_file = st.file_uploader("Upload an image...", type=["png", "jpg", "jpeg"]) | |
| background_text = st.text_area("Describe the desired background:", "") | |
| if uploaded_file is not None and background_text: | |
| if st.button("Modify Background"): | |
| with st.spinner("Generating modified image..."): | |
| modified_image = generate_modified_image(uploaded_file, background_text) | |
| if modified_image: | |
| st.image(modified_image, caption="Modified Image", use_column_width=True) | |
| else: | |
| st.info("Please upload an image and describe the desired background.") |