Spaces:
Running
Running
import streamlit as st | |
from PIL import Image | |
import google.generativeai as genai | |
import io | |
import os | |
# 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-pro-vision') | |
image_parts = [ | |
{"mime_type": uploaded_image.type, "data": uploaded_image.getvalue()} | |
] | |
prompt_parts = [ | |
"You are an AI that can modify the background of an image based on a text description.", | |
"Here is the image:", | |
image_parts, | |
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." | |
] | |
response = model.generate_content(prompt_parts, stream=False) | |
response.resolve() | |
if response and hasattr(response, 'parts') and len(response.parts) > 0: | |
image_bytes = response.parts[-1].data['image/png'] # Assuming the output is PNG | |
return Image.open(io.BytesIO(image_bytes)) | |
else: | |
st.error("Failed to generate the modified image.") | |
return None | |
except Exception as e: | |
st.error(f"An error occurred: {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.") |