Create utils

utils

@@ -0,0 +1,588 @@
+import torch
+import numpy as np
+from diffusers.pipelines import FluxPipeline
+from src.flux.condition import Condition
+from PIL import Image
+import argparse
+import os
+import json
+import base64
+import io
+import re
+from PIL import Image, ImageFilter
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from scipy.ndimage import binary_dilation
+import cv2
+import openai
+from tenacity import retry, wait_exponential, stop_after_attempt, retry_if_exception_type
+
+
+from src.flux.generate import generate, seed_everything
+
+try:
+    from mmengine.visualization import Visualizer
+except ImportError:
+    Visualizer = None
+    print("Warning: mmengine is not installed, visualization is disabled.")
+
+import re
+
+def encode_image_to_datauri(path, size=(512, 512)):
+    with Image.open(path).convert('RGB') as img:
+        img = img.resize(size, Image.LANCZOS)
+        buffer = io.BytesIO()
+        img.save(buffer, format='PNG')
+    b64 = base64.b64encode(buffer.getvalue()).decode('utf-8')
+    return b64
+    # return f"data:image/png;base64,{b64}"
+
+
+@retry(
+    reraise=True,
+    wait=wait_exponential(min=1, max=60),
+    stop=stop_after_attempt(6),
+    retry=retry_if_exception_type((openai.error.RateLimitError, openai.error.APIError))
+)
+def cot_with_gpt(image_uri, instruction):
+    response = openai.ChatCompletion.create(
+        model="gpt-4o",   
+        messages=[
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": f'''
+                    Now you are an expert in image editing. Based on the given single image, what atomic image editing instructions should be if the user wants to {instruction}? Let's think step by step. 
+                    Atomic instructions include 13 categories as follows:
+                    - Add: e.g.: add a car on the road
+                    - Remove: e.g.: remove the sofa in the image
+                    - Color Change: e.g.: change the color of the shoes to blue
+                    - Material Change: e.g.: change the material of the sign like stone
+                    - Action Change: e.g.: change the action of the boy to raising hands
+                    - Expression Change: e.g.: change the expression to smile
+                    - Replace: e.g.: replace the coffee with an apple
+                    - Background Change: e.g.: change the background into forest
+                    - Appearance Change: e.g.: make the cup have a floral pattern
+                    - Move: e.g.: move the plane to the left
+                    - Resize: e.g.: enlarge the clock
+                    - Tone Transfer: e.g.: change the weather to foggy
+                    - Style Change: e.g.: make the style of the image to cartoon
+                    Respond *only* with a numbered list.  
+                    Each line must begin with the category in square brackets, then the instruction. Please strictly follow the atomic categories.
+                    The operation (what) and the target (to what) are crystal clear.  
+                    Do not split replace to add and remove.
+                    For example:
+                    “1. [Add] add a car on the road\n
+                    2. [Color Change] change the color of the shoes to blue\n
+                    3. [Move] move the lamp to the left\n"
+                    Do not include any extra text, explanations, JSON or markdown—just the list.
+                    '''},
+                    {
+                        "type": "image_url",   
+                        "image_url": {
+                            "url": f"data:image/jpeg;base64,{image_uri}"
+                        }
+                    },
+                ],
+            }
+        ],
+        max_tokens=300,
+    )
+    text = response.choices[0].message.content.strip()
+    print(text)
+
+    categories, instructions = extract_instructions(text)
+    return categories, instructions
+
+
+def extract_instructions(text):
+    categories = []
+    instructions = []
+    
+    pattern = r'^\s*\d+\.\s*\[(.*?)\]\s*(.*?)$'
+
+    for line in text.split('\n'):
+        line = line.strip()
+        if not line:
+            continue
+        
+        match = re.match(pattern, line)
+        if match:
+            category = match.group(1).strip()
+            instruction = match.group(2).strip()
+            
+            if category and instruction:
+                categories.append(category)
+                instructions.append(instruction)
+    
+    return categories, instructions
+
+def extract_last_bbox(result):
+    pattern = r'\[?<span data-type="inline-math" data-value="XCcoW15cJ10rKVwnLFxzKlxbXHMqKFxkKylccyosXHMqKFxkKylccyosXHMqKFxkKylccyosXHMqKFxkKylccypcXQ=="></span>\]?'
+    matches = re.findall(pattern, result)
+    
+    if not matches:
+        simple_pattern = r'\[\s*(\d+)\s*,\s*(\d+)\s*,\s*(\d+)\s*,\s*(\d+)\s*\]'
+        simple_matches = re.findall(simple_pattern, result)
+        if simple_matches:
+            x0, y0, x1, y1 = map(int, simple_matches[-1])
+            return [x0, y0, x1, y1]
+        else:
+            print(f"No bounding boxes found, please try again: {result}")
+            return None
+    
+    last_match = matches[-1]
+    x0, y0, x1, y1 = map(int, last_match[1:])
+    return x0, y0, x1, y1
+
+
+def infer_with_DiT(task, image, instruction, category):
+    seed_everything(3407)
+
+    if task == 'RoI Inpainting':
+        if category == 'Add' or category == 'Replace':
+            lora_path = "weights/add.safetensors"
+            added = extract_object_with_gpt(instruction)
+            instruction_dit = f"add {added} on the black region"
+        elif category == 'Remove' or category == 'Action Change':
+            lora_path = "weights/remove.safetensors"
+            instruction_dit = f"Fill the hole of the image"
+       
+        condition = Condition("scene", image, position_delta=(0, 0))
+    elif task == 'RoI Editing':
+        image = Image.open(image).convert('RGB').resize((512, 512))
+        condition = Condition("scene", image, position_delta=(0, -32))
+        instruction_dit = instruction
+        if category == 'Action Change':
+            lora_path = "weights/action.safetensors"
+        elif category == 'Expression Change':
+            lora_path = "weights/expression.safetensors"
+        elif category == 'Add':
+            lora_path = "weights/addition.safetensors"
+        elif category == 'Material Change':
+            lora_path = "weights/material.safetensors"
+        elif category == 'Color Change':
+            lora_path = "weights/color.safetensors"
+        
+    elif task == 'RoI Compositioning':
+        lora_path = "weights/fusion.safetensors"
+        condition = Condition("scene", image, position_delta=(0, 0))
+        instruction_dit = "inpaint the black-bordered region so that the object's edges blend smoothly with the background"
+
+    elif task == 'Global Transformation':
+        image = Image.open(image).convert('RGB').resize((512, 512))
+        instruction_dit = instruction
+        lora_path = "weights/overall.safetensors"
+
+        condition = Condition("scene", image, position_delta=(0, -32))
+    else:
+        raise ValueError(f"Invalid task: '{task}'")
+    pipe = FluxPipeline.from_pretrained(
+        "black-forest-labs/FLUX.1-dev",
+        torch_dtype=torch.bfloat16
+    )
+
+    pipe = pipe.to("cuda")
+
+    pipe.load_lora_weights(
+        "Cicici1109/IEAP",
+        weight_name=lora_path,
+        adapter_name="scene",
+    )
+    result_img = generate(
+        pipe,
+        prompt=instruction_dit,
+        conditions=[condition],
+        config_path = "train/config/scene_512.yaml",
+        num_inference_steps=28,
+        height=512,
+        width=512,
+    ).images[0]
+    # result_img
+    if task == 'RoI Editing' and category == 'Action Change':
+        text_roi = extract_object_with_gpt(instruction)
+        instruction_loc = f"<image>Please segment {text_roi}."
+        # (model, tokenizer, image_path, instruction, work_dir, dilate):
+        img = result_img
+        print(f"Instruction: {instruction_loc}")
+
+        model, tokenizer = load_model("ByteDance/Sa2VA-8B")
+
+        result = model.predict_forward(
+            image=img,
+            text=instruction_loc,
+            tokenizer=tokenizer,
+        )
+
+        prediction = result['prediction']
+        print(f"Model Output: {prediction}")
+
+        if '[SEG]' in prediction and 'prediction_masks' in result:
+            pred_mask = result['prediction_masks'][0]   
+            pred_mask_np = np.squeeze(np.array(pred_mask))
+    
+            ## obtain region bbox
+            rows = np.any(pred_mask_np, axis=1)
+            cols = np.any(pred_mask_np, axis=0)
+            if not np.any(rows) or not np.any(cols):
+                print("Warning: Mask is empty, cannot compute bounding box")
+                return img
+        
+            y0, y1 = np.where(rows)[0][[0, -1]]
+            x0, x1 = np.where(cols)[0][[0, -1]] 
+
+            changed_instance = crop_masked_region(result_img, pred_mask_np)
+
+            return changed_instance, x0, y1, 1
+
+
+    return result_img
+
+def load_model(model_path):
+    model = AutoModelForCausalLM.from_pretrained(
+        model_path,
+        torch_dtype="auto",
+        device_map="auto",
+        trust_remote_code=True
+    ).eval()
+    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+    return model, tokenizer
+
+def extract_object_with_gpt(instruction):
+    system_prompt = (
+        "You are a helpful assistant that extracts the object or target being edited in an image editing instruction. "
+        "Only return a concise noun phrase describing the object. "
+        "Examples:\n"
+        "- Input: 'Remove the dog' → Output: 'the dog'\n"
+        "- Input: 'Add a hat on the dog' → Output: 'a hat'\n"
+        "- Input: 'Replace the biggest bear with a tiger' → Output: 'the biggest bear'\n"
+        "- Input: 'Change the action of the girl to riding' → Output: 'the girl'\n"
+        "- Input: 'Move the red car on the lake' → Output: 'the red car'\n"
+        "- Input: 'Minify the carrot on the rabbit's hand' → Output: 'the carrot on the rabbit's hand'\n"
+        "- Input: 'Swap the location of the dog and the cat' → Output: 'the dog and the cat'\n"
+        "Now extract the object for this instruction:"
+    )
+
+    try:
+        response = openai.ChatCompletion.create(
+            model="gpt-3.5-turbo",
+            messages=[
+                {"role": "system", "content": system_prompt},
+                {"role": "user", "content": instruction}
+            ],
+            temperature=0.2,
+            max_tokens=20,
+        )
+        object_phrase = response.choices[0].message['content'].strip().strip('"')
+        print(f"Identified object: {object_phrase}")
+        return object_phrase
+    except Exception as e:
+        print(f"GPT extraction failed: {e}")
+        return instruction 
+    
+def extract_region_with_gpt(instruction):
+    system_prompt = (
+        "You are a helpful assistant that extracts target region being edited in an image editing instruction. "
+        "Only return a concise noun phrase describing the target region. "
+        "Examples:\n"
+        "- Input: 'Add a red hat to the man on the left' → Output: 'the man on the left'\n"
+        "- Input: 'Add a cat beside the dog' → Output: 'the dog'\n"
+        "Now extract the target region for this instruction:"
+    )
+
+    try:
+        response = openai.ChatCompletion.create(
+            model="gpt-3.5-turbo",
+            messages=[
+                {"role": "system", "content": system_prompt},
+                {"role": "user", "content": instruction}
+            ],
+            temperature=0.2,
+            max_tokens=20,
+        )
+        object_phrase = response.choices[0].message['content'].strip().strip('"')
+        print(f"Identified object: {object_phrase}")
+        return object_phrase
+    except Exception as e:
+        print(f"GPT extraction failed: {e}")
+        return instruction 
+    
+def get_masked(mask, image):
+    if mask.shape[:2] != image.size[::-1]:  
+        raise ValueError(f"Mask size {mask.shape[:2]} does not match image size {image.size}")
+    
+    image_array = np.array(image)
+    image_array[mask] = [0, 0, 0] 
+
+    return Image.fromarray(image_array)
+
+def bbox_to_mask(x0, y0, x1, y1, image_shape=(512, 512), fill_value=True):
+    height, width = image_shape
+    
+    mask = np.zeros((height, width), dtype=bool)
+    
+    x0 = max(0, int(x0))
+    y0 = max(0, int(y0))
+    x1 = min(width, int(x1))
+    y1 = min(height, int(y1))
+    
+    if x0 >= x1 or y0 >= y1:
+        print("Warning: Invalid bounding box coordinates")
+        return mask
+
+    mask[y0:y1, x0:x1] = fill_value
+    
+    return mask
+
+def combine_bbox(text, x0, y0, x1, y1):
+    bbox = [x0, y0, x1, y1]
+    return [(text, bbox)]
+
+def crop_masked_region(image, pred_mask_np):
+    if not isinstance(image, Image.Image):
+        raise ValueError("The input image is not a PIL Image object")
+    if not isinstance(pred_mask_np, np.ndarray) or pred_mask_np.dtype != bool:
+        raise ValueError("pred_mask_np must be a NumPy array of boolean type")
+    if pred_mask_np.shape[:2] != image.size[::-1]:
+        raise ValueError(f"Mask size {pred_mask_np.shape[:2]} does not match image size {image.size}")
+    
+    image_rgba = image.convert("RGBA")
+    image_array = np.array(image_rgba)
+    
+    rows = np.any(pred_mask_np, axis=1)
+    cols = np.any(pred_mask_np, axis=0)
+    
+    if not np.any(rows) or not np.any(cols):
+        print("Warning: Mask is empty, cannot compute bounding box")
+        return image_rgba
+    
+    y0, y1 = np.where(rows)[0][[0, -1]]
+    x0, x1 = np.where(cols)[0][[0, -1]]
+
+    cropped_image = image_array[y0:y1+1, x0:x1+1].copy()
+    cropped_mask = pred_mask_np[y0:y1+1, x0:x1+1]
+    
+    alpha_channel = np.ones(cropped_mask.shape, dtype=np.uint8) * 255
+    alpha_channel[~cropped_mask] = 0  
+    
+    cropped_image[:, :, 3] = alpha_channel
+
+    return Image.fromarray(cropped_image, mode='RGBA')
+
+def roi_localization(image, instruction, category): # add, remove, replace, action change, move, resize
+    model, tokenizer = load_model("ByteDance/Sa2VA-8B")
+    if category == 'Add':
+        text_roi = extract_region_with_gpt(instruction)
+    else:
+        text_roi = extract_object_with_gpt(instruction)
+    instruction_loc = f"<image>Please segment {text_roi}."
+    img = Image.open(image).convert('RGB').resize((512, 512))
+    print(f"Processing image: {os.path.basename(image)}, Instruction: {instruction_loc}")
+
+    result = model.predict_forward(
+        image=img,
+        text=instruction_loc,
+        tokenizer=tokenizer,
+    )
+
+    prediction = result['prediction']
+    print(f"Model Output: {prediction}")
+
+    if '[SEG]' in prediction and 'prediction_masks' in result:
+        pred_mask = result['prediction_masks'][0]   
+        pred_mask_np = np.squeeze(np.array(pred_mask))
+        if category == 'Add':
+            ## obtain region bbox
+            rows = np.any(pred_mask_np, axis=1)
+            cols = np.any(pred_mask_np, axis=0)
+            if not np.any(rows) or not np.any(cols):
+                print("Warning: Mask is empty, cannot compute bounding box")
+                return img  
+    
+            y0, y1 = np.where(rows)[0][[0, -1]]
+            x0, x1 = np.where(cols)[0][[0, -1]] 
+            
+            ## obtain inpainting bbox
+            bbox = combine_bbox(text_roi, x0, y0, x1, y1) #? multiple?
+            print(bbox)
+            x0, y0, x1, y1 = layout_add(bbox, instruction)
+            mask = bbox_to_mask(x0, y0, x1, y1)
+            ## make it black
+            masked_img = get_masked(mask, img)
+        elif category == 'Move' or category == 'Resize':
+            dilated_original_mask = binary_dilation(pred_mask_np, iterations=3)  
+            masked_img = get_masked(dilated_original_mask, img)
+            ## obtain region bbox
+            rows = np.any(pred_mask_np, axis=1)
+            cols = np.any(pred_mask_np, axis=0)
+            if not np.any(rows) or not np.any(cols):
+                print("Warning: Mask is empty, cannot compute bounding box")
+                return img  
+    
+            y0, y1 = np.where(rows)[0][[0, -1]]
+            x0, x1 = np.where(cols)[0][[0, -1]] 
+            
+            ## obtain inpainting bbox
+            bbox = combine_bbox(text_roi, x0, y0, x1, y1) #? multiple?
+            print(bbox)
+            x0_new, y0_new, x1_new, y1_new, = layout_change(bbox, instruction)
+            scale = (y1_new - y0_new) / (y1 - y0)
+            print(scale)
+            changed_instance = crop_masked_region(img, pred_mask_np)
+
+            return masked_img, changed_instance, x0_new, y1_new, scale
+        else:
+            dilated_original_mask = binary_dilation(pred_mask_np, iterations=3)  
+            masked_img = get_masked(dilated_original_mask, img)
+                
+        return masked_img
+
+    else:
+        print("No valid mask found in the prediction.")
+        return None
+
+def fusion(background, foreground, x, y, scale):
+    background = background.convert("RGBA")
+    bg_width, bg_height = background.size
+
+    fg_width, fg_height = foreground.size
+    new_size = (int(fg_width * scale), int(fg_height * scale))
+    foreground_resized = foreground.resize(new_size, Image.Resampling.LANCZOS)
+
+    left = x
+    top = y - new_size[1]
+
+    canvas = Image.new('RGBA', (bg_width, bg_height), (0, 0, 0, 0))
+    canvas.paste(foreground_resized, (left, top), foreground_resized)
+    masked_foreground = process_edge(canvas, left, top, new_size)
+    result = Image.alpha_composite(background, masked_foreground)
+    
+    return result
+
+def process_edge(canvas, left, top, size):
+    width, height = size
+    
+    region = canvas.crop((left, top, left + width, top + height))
+    alpha = region.getchannel('A')
+    
+    dilated_alpha = alpha.filter(ImageFilter.MaxFilter(5)) 
+    eroded_alpha = alpha.filter(ImageFilter.MinFilter(3)) 
+    
+    edge_mask = Image.new('L', (width, height), 0)
+    edge_pixels = edge_mask.load()
+    dilated_pixels = dilated_alpha.load()
+    eroded_pixels = eroded_alpha.load()
+    
+    for y in range(height):
+        for x in range(width):
+            if dilated_pixels[x, y] > 0 and eroded_pixels[x, y] == 0:
+                edge_pixels[x, y] = 255  
+    
+    black_edge = Image.new('RGBA', (width, height), (0, 0, 0, 0))
+    black_edge.putalpha(edge_mask)
+    
+    canvas.paste(black_edge, (left, top), black_edge)
+    
+    return canvas
+
+def combine_text_and_bbox(text_roi, x0, y0, x1, y1):
+    return [(text_roi, [x0, y0, x1, y1])]
+
+@retry(
+    reraise=True,
+    wait=wait_exponential(min=1, max=60),
+    stop=stop_after_attempt(6),
+    retry=retry_if_exception_type((openai.error.RateLimitError, openai.error.APIError))
+)
+def layout_add(bbox, instruction):
+    response = openai.ChatCompletion.create(
+        model="gpt-4o",  
+        messages=[
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": f'''
+                    You are an intelligent bounding box editor. I will provide you with the current bounding boxes and an add editing instruction. 
+                    Your task is to determine the new bounding box of the added object. Let's think step by step.
+                    The images are of size 512x512. The top-left corner has coordinate [0, 0]. The bottom-right corner has coordinnate [512, 512]. 
+                    The bounding boxes should not go beyond the image boundaries. The new box must be large enough to reasonably encompass the added object in a visually appropriate way, allowing for partial overlap with existing objects when it comes to accessories like hat, necklace. etc.
+                    Each bounding box should be in the format of (object name,[top-left x coordinate, top-left y coordinate, bottom-right x coordinate, bottom-right y coordinate]).
+                    Only return the bounding box of the newly added object. Do not include the existing bounding boxes.
+                    Please consider the semantic information of the layout, preserve semantic relations.
+                    If needed, you can make reasonable guesses. Please refer to the examples below:
+                    Input bounding boxes: [('a green car', [21, 281, 232, 440])]
+                    Editing instruction: Add a bird on the green car.
+                    Output bounding boxes: [('a bird', [80, 150, 180, 281])]
+                    Input bounding boxes: [('stool', [300, 350, 380, 450])]
+                    Editing instruction: Add a cat to the left of the stool.
+                    Output bounding boxes: [('a cat', [180, 250, 300, 450])]
+                     
+                    Here are some examples to illustrate appropriate overlapping for better visual effects:
+                    Input bounding boxes: [('the white cat', [200, 300, 320, 420])]
+                    Editing instruction: Add a hat on the white cat.
+                    Output bounding boxes: [('a hat', [200, 150, 320, 330])]
+                    Now, the current bounding boxes is {bbox}, the instruction is {instruction}.
+                    '''},
+                ],
+            }
+        ],
+        max_tokens=1000,
+    )
+
+    result = response.choices[0].message.content.strip()
+
+    print(result)
+    bbox = extract_last_bbox(result)
+    return bbox
+
+@retry(
+    reraise=True,
+    wait=wait_exponential(min=1, max=60),
+    stop=stop_after_attempt(6),
+    retry=retry_if_exception_type((openai.error.RateLimitError, openai.error.APIError))
+)
+def layout_change(bbox, instruction):
+    response = openai.ChatCompletion.create(
+        model="gpt-4o",   
+        messages=[
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": f'''
+                    You are an intelligent bounding box editor. I will provide you with the current bounding boxes and the editing instruction. 
+                    Your task is to generate the new bounding boxes after editing. 
+                    The images are of size 512x512. The top-left corner has coordinate [0, 0]. The bottom-right corner has coordinnate [512, 512]. 
+                    The bounding boxes should not overlap or go beyond the image boundaries. 
+                    Each bounding box should be in the format of (object name, [top-left x coordinate, top-left y coordinate, bottom-right x coordinate, bottom-right y coordinate]). 
+                    Do not add new objects or delete any object provided in the bounding boxes. Do not change the size or the shape of any object unless the instruction requires so.
+                    Please consider the semantic information of the layout. 
+                    When resizing, keep the bottom-left corner fixed by default. When swaping locations, change according to the center point. 
+                    If needed, you can make reasonable guesses. Please refer to the examples below:
+                     
+                    Input bounding boxes: [('a car', [21, 281, 232, 440])]
+                    Editing instruction: Move the car to the right.
+                    Output bounding boxes: [('a car', [121, 281, 332, 440])]
+                     
+                    Input bounding boxes: [("bed", [50, 300, 450, 450]), ("pillow", [200, 200, 300, 230])]
+                    Editing instruction: Move the pillow to the left side of the bed.
+                    Output bounding boxes: [("bed", [50, 300, 450, 450]), ("pillow", [70, 270, 170, 300])]
+                     
+                    Input bounding boxes: [("dog", [150, 250, 250, 300])]
+                    Editing instruction: Enlarge the dog.
+                    Output bounding boxes: [("dog", [150, 225, 300, 300])]
+                     
+                    Input bounding boxes: [("chair", [100, 350, 200, 450]), ("lamp", [300, 200, 360, 300])]
+                    Editing instruction: Swap the location of the chair and the lamp.
+                    Output bounding boxes: [("chair", [280, 200, 380, 300]), ("lamp", [120, 350, 180, 450])]
+
+
+                    Now, the current bounding boxes is {bbox}, the instruction is {instruction}. Let's think step by step, and output the edited layout.
+                    '''},
+                ],
+            }
+        ],
+        max_tokens=1000,
+    )
+    result = response.choices[0].message.content.strip()
+    print(result)
+    bbox = extract_last_bbox(result)
+    return bbox