app.py

import os
import base64
from io import BytesIO
import numpy as np
from PIL import Image
import gradio as gr
import replicate
import requests
import google.generativeai as genai

genai.configure(api_key=os.getenv("GENAI_API_KEY"))


def remove_background(image):
    if isinstance(image, np.ndarray):
        try:
            img = Image.fromarray(image)
            buffered = BytesIO()
            img.save(buffered, format="PNG")
            img_base64 = base64.b64encode(buffered.getvalue()).decode('utf-8')

            output = replicate.run(
                "cjwbw/rembg:fb8af171cfa1616ddcf1242c093f9c46bcada5ad4cf6f2fbe8b81b330ec5c003",
                input={"image": f"data:image/png;base64,{img_base64}"}
            )

            response = requests.get(output)
            return Image.open(BytesIO(response.content))
        except Exception as e:
            print(f"Error removing background: {str(e)}")
            return None


def load_background_images(folder_path):
    if not os.path.exists(folder_path):
        return []
    return [Image.open(os.path.join(folder_path, file))
            for file in os.listdir(folder_path)
            if file.lower().endswith(('.png', '.jpg', '.jpeg'))]


def replace_background(input_image_path, background_image):
    if input_image_path is None or background_image is None:
        return None

    try:
        input_img = Image.open(input_image_path)
    except Exception as e:
        print(f"Error loading input image: {e}")
        return None

    input_img_np = np.array(input_img)

    foreground = remove_background(input_img_np)
    if foreground is None:
        return None

    width, height = foreground.size

    if isinstance(background_image, np.ndarray):
        background = Image.fromarray(background_image)
    else:
        try:
            background = Image.open(background_image) if isinstance(background_image, str) else background_image
        except Exception as e:
            print(f"Error loading background image: {e}")
            return None

    background = background.resize((width, height))

    background.paste(foreground, (0, 0), foreground)
    return background


def image_upload_for_gemini(image_path):
    sample_file = genai.upload_file(path=image_path, display_name="Jetpack drawing")
    return sample_file


def generate_title_description(image):
    if image is None:
        return "Please upload an image first."

    try:
        sample_file = image_upload_for_gemini(image)

        # Set up model and prompt
        model = genai.GenerativeModel(model_name="gemini-1.5-flash")
        prompt = """Analyze the provided prodct image and generate an optimized e-commerce product listing with:

SEO-friendly, attention-grabbing product title (3-8 words).
Engaging, shopper-focused product description (80-120 words).
Formatting:

Title: [Insert Product Title]
Description: [Insert Product Description]
Product Description Requirements:

Materials & specifications: Mention key details like metal type, gemstones, chain length, and clasp style.
Unique design features: Emphasize 2-3 standout elements (e.g., intricate patterns, luxurious finish, or versatile design).
Styling suggestions: Suggest outfits or occasions it pairs well with.
Craftsmanship & quality: Mention its artisanal or premium build.
Emotional appeal: Convey how it makes the wearer feel (e.g., confident, elegant, cherished).
Optional care instructions: Briefly include how to maintain its beauty.
Tags: Provide 5-10 keywords for SEO, including material, style, occasion, and audience"""

        response = model.generate_content([sample_file, prompt])
        return response.text
    except Exception as e:
        return f"Error generating description: {str(e)}"


backgrounds = {
    'Studio': load_background_images('resources/bg_images/studio'),
    'Color': load_background_images('resources/bg_images/color'),
    'Wall': load_background_images('resources/bg_images/wall'),
    'Nature': load_background_images('resources/bg_images/nature'),
    'House': load_background_images('resources/bg_images/house'),
    'Wood': load_background_images('resources/bg_images/wood')
}

with gr.Blocks(theme=gr.themes.Default(primary_hue=gr.themes.colors.red)) as demo:
    gr.Markdown("## Product Image Editor and Description Generator")

    with gr.Tabs():
        with gr.Tab("Photo Shoot"):

            with gr.Row():
                input_image = gr.Image(label="Input Image", type='filepath')
                background_image = gr.Image(label="Selected Background", visible=False)
                output_image = gr.Image(label="Result", type='pil')

            for category, images in backgrounds.items():
                if images:
                    with gr.Row():
                        gr.Examples(
                            examples=images,
                            inputs=background_image,
                            label=f"{category} Backgrounds"
                        )

            with gr.Row():
                submit_btn = gr.Button("Generate")
                clear_btn = gr.ClearButton(
                    components=[input_image, output_image, background_image],
                    value="Reset"
                )

            submit_btn.click(
                fn=replace_background,
                inputs=[input_image, background_image],
                outputs=output_image
            )

        with gr.Tab("Product Description"):
            with gr.Row():
                description_input = gr.Image(label="Product Image", type='filepath')
                description_output = gr.Textbox(
                    label="Generated Title and Description",
                    lines=10,
                    placeholder="Upload an image to generate product details..."
                )

            generate_desc_btn = gr.Button("Generate Title & Description")
            generate_desc_btn.click(
                fn=generate_title_description,
                inputs=description_input,
                outputs=description_output
            )

if __name__ == "__main__":
    demo.launch(share=True, debug=True)