wip

.gradio/flagged/dataset1.csv

@@ -0,0 +1,2 @@
+input_image,"Holiday (e.g. Christmas, New Year's, etc.)",Optional Message,First Output,Second Output,Third Output,timestamp
+.gradio/flagged/input_image/aeff2e0970e0cc20c2f3/Linda-Sobolewski-Photography-Family-Session-00002-1024x683.jpg,,,.gradio/flagged/First Output/f0fa5bc42029b76a623f/image.webp,.gradio/flagged/Second Output/1c38f7ceb54aaa9b3c55/image.webp,.gradio/flagged/Third Output/bc5418b161ee6a0f583f/image.webp,2024-12-10 15:59:53.431810

app.py

@@ -4,6 +4,21 @@ import matplotlib.pyplot as plt
 import torch
 from torchvision import transforms
 from transformers import AutoModelForImageSegmentation
+from openai import OpenAI
+import os
+import base64
+import io
+import requests
+import numpy as np
+from scipy import ndimage
+
+IDEOGRAM_API_KEY = os.getenv('IDEOGRAM_API_KEY')
+IDEOGRAM_URL = "https://api.ideogram.ai/edit"
+
+client = OpenAI(api_key=os.getenv('OPENAI_API_KEY'))
+# Constants should be in UPPERCASE
+GPT_MODEL_NAME = "gpt-4o"
+GPT_MAX_TOKENS = 500
 
 model = AutoModelForImageSegmentation.from_pretrained('briaai/RMBG-2.0', trust_remote_code=True)
 torch.set_float32_matmul_precision(['high', 'highest'][0])
@@ -11,7 +26,51 @@ if torch.cuda.is_available():
     model = model.to('cuda')
 model.eval()
 
-def remove_background(input_image, holiday, message):
+GPT_PROMPT = '''
+I work with a tool that knows how to edit backgrounds.
+I want your help with prompt.
+I want to adjust their background to be in a christmas vibes.
+For example, if you see a tree there, cover it in snow,
+add christmas lights to some of the stuff in the background, maybe add a few elements like christmas tree, but take into considration the perspective and the logic of the image.
+'''
+
+def image_to_prompt(image: str) -> tuple[str, str]:
+    base64_image = encode_image(image)
+    
+    messages = [{
+        "role": "user",
+        "content": [
+            {"type": "text", "text": GPT_PROMPT},
+            {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{base64_image}"}}
+        ]
+    }]
+    
+    response = client.chat.completions.create(
+        model=GPT_MODEL_NAME,
+        messages=messages,
+        max_tokens=GPT_MAX_TOKENS
+    )
+    
+    full_response = response.choices[0].message.content
+    return full_response
+
+def encode_image(image: Image.Image) -> str:
+    """Convert a PIL Image to base64 encoded string.
+    
+    Args:
+        image (PIL.Image.Image): The PIL Image to encode
+        
+    Returns:
+        str: Base64 encoded image string
+    """
+    # Create a temporary buffer to save the image
+    buffer = io.BytesIO()
+    # Save the image as PNG to the buffer
+    image.save(buffer, format='PNG')
+    # Get the bytes from the buffer and encode to base64
+    return base64.b64encode(buffer.getvalue()).decode('utf-8')
+    
+def remove_background(input_image):
     image_size = (1024, 1024)
     # Transform the input image
     transform_image = transforms.Compose([
@@ -41,15 +100,72 @@ def remove_background(input_image, holiday, message):
     inverted_mask = Image.eval(mask, lambda x: 255 - x)  # Invert the mask
     only_background_image.putalpha(inverted_mask)
     
-    first_output_image = result_image
-    second_output_image = only_background_image
-    third_output_image = result_image
+    return result_image, only_background_image, mask
+
+def modify_background(image: Image.Image, mask: Image.Image, prompt: str) -> Image.Image:
+    # Convert PIL images to bytes
+    image_buffer = io.BytesIO()
+    image.save(image_buffer, format='PNG')
+    image_bytes = image_buffer.getvalue()
+    
+    mask_buffer = io.BytesIO()
+    mask.save(mask_buffer, format='PNG')
+    mask_bytes = mask_buffer.getvalue()
+    
+    # Create the files dictionary with actual bytes data
+    files = {
+        "image_file": ("image.png", image_bytes, "image/png"),
+        "mask": ("mask.png", mask_bytes, "image/png")  # You might want to send a different mask file
+    }
+    
+    payload = {
+        "prompt": prompt,  # Use the actual prompt parameter
+        "model": "V_2",
+        "magic_prompt_option": "ON",
+        "num_images": 1,
+        "style_type": "REALISTIC"
+    }
+    headers = {"Api-Key": IDEOGRAM_API_KEY}
+
+    response = requests.post(IDEOGRAM_URL, data=payload, files=files, headers=headers)
+    
+    if response.status_code == 200:
+        # Assuming the API returns an image in the response
+        response_data = response.json()
+        # You'll need to handle the response according to Ideogram's API specification
+        # This is a placeholder - adjust according to actual API response format
+        result_image_url = response_data.get('data')[0].get('url')
+        if result_image_url:
+            result_response = requests.get(result_image_url)
+            return Image.open(io.BytesIO(result_response.content))
+    
+    raise Exception(f"Failed to modify background: {response.text}")
+
+def dilate_mask(mask: Image.Image) -> Image.Image:
+    # Convert mask to numpy array
+    mask_array = np.array(mask)
+    
+    # Apply maximum filter using scipy.ndimage
+    dilated_mask = ndimage.maximum_filter(mask_array, size=20)
+    
+    # Convert back to PIL Image
+    return Image.fromarray(dilated_mask.astype(np.uint8))
+
+def run_flow(input_image, holiday, message):
+    prompt = image_to_prompt(input_image)
+    print(prompt)
+    result_image, only_background_image, mask = remove_background(input_image)
+    dilated_mask = dilate_mask(mask)
+    modified_image = modify_background(input_image, dilated_mask, prompt)
+    first_output_image = mask
+    second_output_image = dilated_mask
+    third_output_image = modified_image
     
     return first_output_image, second_output_image, third_output_image
 
 # Replace the demo interface
 demo = gr.Interface(
-    fn=remove_background,
+    fn=run_flow,
     inputs=[
         gr.Image(type="pil"),
         gr.Text(label="Holiday (e.g. Christmas, New Year's, etc.)"),

requirements.txt

@@ -5,4 +5,6 @@ pillow
 kornia
 transformers
 timm
-matplotlib
+matplotlib
+openai
+requests