utils/utils.py

from PIL import Image
import numpy as np
import cv2


#first stage sketch preprocess
def conventional_resize(img):
    original_width, original_height = img.size    
    aspect_ratio = original_width / original_height

    conventional_sizes = {
        1: (1024, 1024),  
        4/3: (1152, 896), 
        3/2: (1216, 832),
        16/9: (1344, 768),
        21/9: (1568, 672),
        3/1: (1728, 576),
        1/4: (512, 2048), 
        1/3: (576, 1728),
        9/16: (768, 1344),
        2/3: (832, 1216),
        3/4: (896, 1152)
    }

    closest_aspect_ratio = min(conventional_sizes.keys(), key=lambda x: abs(x - aspect_ratio)) #1.333333333 #
    new_width, new_height = conventional_sizes[closest_aspect_ratio]
    
    resized_img = img.resize((new_width, new_height), Image.LANCZOS)

    return resized_img


def get_substrate(img, color=(255, 255, 255, 255)):
    size=img.size
    substrate = Image.new("RGBA", size, color) 
    return substrate.convert("RGB")


def sketch_process(img):
    substrate=conventional_resize(get_substrate(img))
    resized_img = conventional_resize(img) 
    return substrate, resized_img




#first stage prompt preprocess
def remove_duplicates(base_prompt):    
    prompt_list = base_prompt.split(", ")
    seen = set()
    unique_tags = []
    for tag in prompt_list :
        tag_clean = tag.lower().strip()
        if tag_clean not in seen and tag_clean != "":
            unique_tags.append(tag)
            seen.add(tag_clean)
    return ", ".join(unique_tags)


def remove_color(base_prompt):    
    prompt_list = base_prompt.split(", ")
    color_list = ["pink", "red", "orange", "brown", "yellow", "green", "blue", "purple", "blonde", "colored skin", "white hair"]
    cleaned_tags = [tag for tag in prompt_list if all(color.lower() not in tag.lower() for color in color_list)]
    return ", ".join(cleaned_tags)


def execute_prompt(base_prompt):
    prompt_list = base_prompt.split(", ")
    execute_tags = ["sketch", "transparent background"]    
    filtered_tags = [tag for tag in prompt_list if tag not in execute_tags]    
    return ", ".join(filtered_tags)

def prompt_preprocess(prompt):
    result=execute_prompt(prompt)
    result=remove_duplicates(result)
    result=remove_color(result)
    return result