File size: 2,560 Bytes
2649e2f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c7602de
 
2649e2f
2325a30
2649e2f
 
 
 
2325a30
2649e2f
 
 
 
 
 
 
 
2325a30
 
 
 
 
 
 
 
 
2649e2f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
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:
        # Update the model name here
        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."
        ]

        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:
                if hasattr(part.data, 'image/png'):
                    image_bytes = part.data['image/png']
                    return Image.open(io.BytesIO(image_bytes))
                elif hasattr(part.data, 'image/jpeg'):
                    image_bytes = part.data['image/jpeg']
                    return Image.open(io.BytesIO(image_bytes))
            st.error("Generated content did not contain an image.")
            return None
        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.")