app.py

import os
import re
import json
import cv2
import time
import requests
import numpy as np
import gradio as gr
import pandas as pd
from PIL import Image
from io import BytesIO
from openai import OpenAI
from sklearn.metrics.pairwise import cosine_similarity
from sklearn.feature_extraction.text import TfidfVectorizer

import warnings
warnings.filterwarnings("ignore")


client = OpenAI(api_key='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')
STABILITY_KEY = "sk-8Wcz3bcvHo4r0ehXEvcC2mQu4gV3aYWdrdFmtri9Isn1ycDp"

def send_generation_request(host,params,):
    headers = {"Accept": "image/*","Authorization": f"Bearer {STABILITY_KEY}"}

    # Encode parameters
    files = {}
    image = params.pop("image", None)
    mask = params.pop("mask", None)
    if image is not None and image != '':
        files["image"] = open(image, 'rb')
    if mask is not None and mask != '':
        files["mask"] = open(mask, 'rb')
    if len(files)==0:
        files["none"] = ''

    # Send request
    print(f"Sending REST request to {host}...")
    response = requests.post(
        host,
        headers=headers,
        files=files,
        data=params
    )
    if not response.ok:
        raise Exception(f"HTTP {response.status_code}: {response.text}")

    return response

def send_async_generation_request(
    host,
    params,
):
    headers = {
        "Accept": "application/json",
        "Authorization": f"Bearer {STABILITY_KEY}"
    }

    # Encode parameters
    files = {}
    if "image" in params:
        image = params.pop("image")
        files = {"image": open(image, 'rb')}

    # Send request
    print(f"Sending REST request to {host}...")
    response = requests.post(
        host,
        headers=headers,
        files=files,
        data=params
    )
    if not response.ok:
        raise Exception(f"HTTP {response.status_code}: {response.text}")

    # Process async response
    response_dict = json.loads(response.text)
    generation_id = response_dict.get("id", None)
    assert generation_id is not None, "Expected id in response"

    # Loop until result or timeout
    timeout = int(os.getenv("WORKER_TIMEOUT", 500))
    start = time.time()
    status_code = 202
    while status_code == 202:
        response = requests.get(
            f"{host}/result/{generation_id}",
            headers={
                **headers,
                "Accept": "image/*"
            },
        )

        if not response.ok:
            raise Exception(f"HTTP {response.status_code}: {response.text}")
        status_code = response.status_code
        time.sleep(10)
        if time.time() - start > timeout:
            raise Exception(f"Timeout after {timeout} seconds")

    return response

def bckgrnd_removal(Im): 
    image = str(Im)
    output_format = "png"
    host = f"https://api.stability.ai/v2beta/stable-image/edit/remove-background"
    params = {"image" : image,"output_format": output_format}
    response = send_generation_request(host,params)

    # Decode response
    output_image = response.content
    finish_reason = response.headers.get("finish-reason")
    seed = response.headers.get("seed")

    # Check for NSFW classification
    if finish_reason == 'CONTENT_FILTERED':
        raise Warning("Generation failed NSFW classifier")

    # Save and display result
    filename, _ = os.path.splitext(os.path.basename(image))
    bckgrnd_removed_image = f"data\\03_app_generated\\01_bckgrnd_removed_image_{filename}_{seed}.{output_format}"
    with open(bckgrnd_removed_image, "wb") as f:
        f.write(output_image)
    return bckgrnd_removed_image

def custom_tokenizer(text):
    # Using regular expressions to split and filter alphanumeric tokens
    return [token for token in re.split('[ _]', text) if token.isalnum()]

def search_image(prmpt):
    df = pd.read_csv("data\\02_processed\\car_image_file_name.csv")
    filenames = df['Filename'].values

    # Initialize the TfidfVectorizer with the custom tokenizer
    vectorizer = TfidfVectorizer(tokenizer=custom_tokenizer)
    X_filenames = vectorizer.fit_transform(filenames)

    # Define a user query and vectorize it
    user_query = prmpt
    X_query = vectorizer.transform([user_query])

    # Compute the cosine similarity scores
    similarity_scores = cosine_similarity(X_query, X_filenames)

    # Create a DataFrame with filenames and their similarity scores
    results_df = pd.DataFrame({
        'Filename': filenames,
        'Similarity Score': similarity_scores.flatten()  # ensure the similarity scores are in the correct shape
    })

    # Write the results to a CSV file
    results_df.to_csv("data\\02_processed\\similarity_scores.csv", index=False)

    sorted_results_df = results_df.sort_values(by='Similarity Score', ascending=False)

    # Display the top 10 results
    top_6_results = sorted_results_df.head(20)
    top_6_results['Filename']= top_6_results['Filename'].apply(lambda x:"data\\02_processed\\cv_final_renamed_2\\"+x)
    return top_6_results['Filename'].tolist()


def prmpt_enhancer(rd,prmpt):
    gbf_ip_prmpt = f"Act as expert prompt engineer. I will provide you with a preliminary prompt to generate a background image for a car using inpainting. Step-by-step, identify and integrate unique architectural elements suitable for the setting to enhance its luxury—refer to the provided examples. In consultation with the art director, refine the background to harmonize with the car's color. Finally, enhance the prompt to ensure it incorporates aesthetically pleasing elements for image generation in Stable Diffusion.Write the final prompt in JSON format with the prompt key.##Prompt:```{prmpt}``` \n Examples:car before arab home : car with shadow infront of modern arab villa with perforated facade and trees, warm color, villa cover top of the image, hyperrealistic \n car infront of showroom : Car standing outside 8 Modern artisitic luxury car showroom made of glass with blue sk ,  highly detailed, warm tone , 8k resolution, hyperrealistic \n car inside showroom :  car on Showroom Floor in luxury glass showroom with decorative screen, Vibrant, enticing, commercial, product-focused, eye-catching, professional, highly detailed, warm tone, hyperrealistic \n car on mountain: car with shadow  on snow filled path himalayas, top of image is filled with blue sky and snowy mountain, hyperrealistic'\n car before repair center: car with shadow parked in front of a bustling car repair center with bold digital yellow shop hoarding and blue sky , midday, 8K resolution, hyperrealistic"
    ip_ip_prmpt = f"Act as expert Prompt Engineer. You will be provided with a prompt. Your task is to infuse expressions to human prompt you receive. Only enhance the expression and dress style; do not add details about the car background or the car color. consult with art director and human behaviour analyst to revamp the prompt. \n Example Prompt ```A Arab Traditional family standing before car``` transformed to ```8k, Fashion portrait of a traditional Arab family stands proudly in front of their car, exuding a sense of joyful pride.``` Here are similar transformations and enhancements for other scenarios a. 8K, Fashion portrait of a girl jumping in excitement with a shopping bag; \n b.8K, Fashion portrait of two children dressed in sunflower yellow winter clothing, shivering from the cold as they stand next to their father, who is holding a snowboard, ready for a photo; \n c.8K, Fashion portrait of an excited young woman standing next to her car; \n d. 8K, Fashion portrait of a happy young Arab traditional family caresses the car, children jumping with joy; \n e. 8K, Fashion portrait of a happy young man reclining on a car with one leg on the car;\n f.8K, Fashion portrait of a joyful young man in a suit standing beside a woman in a golden dress. ##Prompt:```{prmpt}```\n Write the final prompt in JSON format with the prompt key."

    if rd == 'Generative Background Fill':
        completion = client.chat.completions.create(
        model="gpt-4",
        messages=[
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": gbf_ip_prmpt}
        ]
        )
        input_text = completion.choices[0].message.content
    else:
        completion = client.chat.completions.create(
        model="gpt-4",
        messages=[
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": ip_ip_prmpt}
        ]
        )
        input_text = completion.choices[0].message.content

    json_match = re.search(r'```json\n({.*?})\n```', input_text, re.DOTALL)
    
    # print(json_match)
    
    if json_match:
        json_str = json_match.group(1)
        # print("----------")
        # print(type(json_str))
        json_data = json.loads(json_str)
        # print("----------")
        # print(type(json_data))
        op_response = json_data["prompt"]
    else:
       op_response =  "No JSON data found"
       
    return op_response


def bckgrnd_gen(Im,bckgrnd_removal_prmpt,negative_prmpt,seed,output_format,ks): 
    image = cv2.imread(Im, cv2.IMREAD_UNCHANGED)
    if image is not None:
        print("Image loaded successfully!")
        print("Shape of the image:", image.shape)
        height, width, _ = image.shape
        print(image.shape)

        if image.shape[2] == 4:
            alpha_channel = image[:, :, 3]
            mask = alpha_channel == 0
            mask = (mask * 255).astype(np.uint8)

            kernel_size = int(ks)
            kernel = np.ones((kernel_size, kernel_size), np.uint8)
            dilated_mask = cv2.dilate(mask, kernel, iterations=7)
            cv2.imwrite('data\\03_app_generated\\02_transparent_mask.png', dilated_mask)

            white_background = np.ones((height, width, 3), dtype=np.uint8) * 255 
            alpha_channel = image[:, :, 3] / 255.0
            alpha_layer = np.stack([alpha_channel]*3, axis=-1)

            foreground = (image[:, :, :3] * alpha_layer).astype(np.uint8)
            background = (white_background * (1 - alpha_layer)).astype(np.uint8)
            blended_image = cv2.add(foreground, background)
            cv2.imwrite('data\\03_app_generated\\02_blended_image.png', blended_image)

        else:
            # mask = np.all(image == [0, 0, 0], axis=-1)
            # white_mask = np.where(mask[..., None], [255, 255, 255], [0, 0, 0])
            # gray_image = cv2.cvtColor(white_mask.astype(np.uint8), cv2.COLOR_BGR2GRAY)
            # cv2.imwrite('transparent_mask.png', gray_image)

            # blended_image = cv2.imread(Im, cv2.IMREAD_UNCHANGED)
            # cv2.imwrite('blended_image.png', blended_image)
            print("No Alpha Layer")

    image = Im
    mask = "data\\03_app_generated\\02_transparent_mask.png"
    prompt = bckgrnd_removal_prmpt
    negative_prompt = negative_prmpt
    seed = str(seed)
    output_format = output_format

    host = f"https://api.stability.ai/v2beta/stable-image/edit/inpaint"

    params = {
        "image" : image,
        "mask" : mask,
        "negative_prompt" : negative_prompt,
        "seed" : seed,
        "mode": "mask",
        "output_format": output_format,
        "prompt" : prompt
    }

    response = send_generation_request(
        host,
        params
    )

    # Decode response
    output_image = response.content
    finish_reason = response.headers.get("finish-reason")
    seed = response.headers.get("seed")

    # Check for NSFW classification
    if finish_reason == 'CONTENT_FILTERED':
        raise Warning("Generation failed NSFW classifier")

    # Save and display result
    filename, _ = os.path.splitext(os.path.basename(image))
    edited = f"data\\03_app_generated\\BCG_{filename}_{seed}.{output_format}"
    with open(edited, "wb") as f:
        f.write(output_image)

    return edited


def character_inpainting(Im,prmpt,negative_prmpt,seed,output_format):
    Im["composite"].save('data\\03_app_generated\\composite_final.png')
    Im["layers"][0] = Im["layers"][0].convert('L')
    Im["layers"][0].save('data\\03_app_generated\\mask_final.png')
    Im["composite"].save('data\\03_app_generated\\background_final.png')

    mask_path = 'data\\03_app_generated\\mask_final.png'
    with Image.open(mask_path) as mask:
        gray_mask = mask.convert('L')  # Convert to grayscale
        gray_mask.save('data\\03_app_generated\\IP_mask_white.png')
	
    image = 'data\\03_app_generated\\background_final.png'
    mask = 'data\\03_app_generated\\IP_mask_white.png'
    
    prompt = prmpt 
    negative_prompt = negative_prmpt
    seed = int(seed) 
    output_format = output_format

    host = f"https://api.stability.ai/v2beta/stable-image/edit/inpaint"

    params = {
        "image" : image,
        "mask" : mask,
        "negative_prompt" : negative_prompt,
        "seed" : seed,
        "mode": "mask",
        "output_format": output_format,
        "prompt" : prompt
    }

    response = send_generation_request(
        host,
        params
    )

    # Decode response
    output_image = response.content
    finish_reason = response.headers.get("finish-reason")
    seed = response.headers.get("seed")

    # Check for NSFW classification
    if finish_reason == 'CONTENT_FILTERED':
        raise Warning("Generation failed NSFW classifier")

    # Save and display result
    filename, _ = os.path.splitext(os.path.basename(image))
    edited = f"edited_{filename}_{seed}.{output_format}"

    with open(edited, "wb") as f:
        f.write(output_image)

    return edited

def list_folder_contents(folder_path):
    gallery_path_list = []
    # Walk through all directories and files in the folder
    for root, dirs, files in os.walk(folder_path):
        for name in files:
            gallery_path_list.append(os.path.join(root, name))
        for name in dirs:
            gallery_path_list.append(os.path.join(root, name))
    return gallery_path_list

# Example usage
folder_path = 'data\\01_original\\gallery_images\\'
folder_contents = list_folder_contents(folder_path)
#print(folder_contents)



with gr.Blocks() as demo:
    gr.Markdown(
        """
        # blend.ai (Fractal Hackathon)
        # Team: Moonshot
        Enhance your marketing visuals effortlessly with Blend.ai, our state-of-the-art inpainting tool designed for seamless image edits. Transform your images into captivating masterpieces with precision and ease!
        """)

    with gr.Tab("Background Remover"):
        gr.Interface(
            fn=bckgrnd_removal,
            inputs=gr.Text(label="Image Path"),
            outputs=gr.Image(format="png", show_download_button=True),
            allow_flagging=False,
            title="Background Remover",
            description="Provide the path to a PNG image of your product. Our tool will automatically remove the background, making it perfect for professional presentations and marketing material."
        )

    with gr.Tab("Image Search"):
        gr.Interface(
            fn=search_image,
            inputs=gr.Text(label="Prompt"),
            outputs=gr.Gallery(),
            allow_flagging=False,
            title="Image Search",
            description="Search and retrieve images based on your text prompts. Ideal for finding specific product shots or design inspirations. Example prompts are included for quick testing.",
            examples=['toyota corolla blue cross 30 degree closeup shot', 'Toyota 2024 Land Cruiser 270 black degree normal shot', 'nissan two tone monarch orange left normal shot']
        )

    with gr.Tab("Prompt Enhancer"):
        gr.Interface(
            fn=prmpt_enhancer,
            inputs=[gr.Radio(["Generative Background Fill", "Inpaint Characters"], label="Section", info="Select the section for which you need to enhance prompt"), gr.Text(label="Prompt")],
            outputs=[gr.Text(label="Enhanced Prompt")],
            allow_flagging='never',
            title="Prompt Enhancer",
            description="Enhance your prompts to improve clarity and specificity. Select the type of enhancement and enter your initial prompt."
        )

    with gr.Tab("Generative Background Fill"):
        gr.Interface(
            fn=bckgrnd_gen,
            inputs=[gr.Text(label="Image Path"), gr.TextArea(label="Prompt"), gr.Text(label="Negative Prompt", value="noisy, blurry, unattractive, sloppy, unprofessional, low quality"), gr.Text(label="Seed", value="0"), gr.Text(label="Output Format", value='png'), gr.Slider(minimum=0, maximum=10, value=2, step=1, label="Kernel Size", show_label=True, info="Adjust Kernel size")],
            outputs=[gr.Image(format="png", show_download_button=True),gr.Image(format="png", show_download_button=True)],
            allow_flagging='never',
            title="Generative Background Fill",
            description="Upload an RGBA PNG image to add a custom background. Provide a positive prompt for desired elements, a negative prompt to exclude elements, and customize the output with our sliders."
        )

    with gr.Tab("Inpaint Characters"):
        gr.Interface(
            fn=character_inpainting,
            inputs=[gr.ImageMask(type='pil', brush=gr.Brush(colors=["#ffffff"], color_mode="fixed"),sources = 'upload'), gr.TextArea(label="Prompt"), gr.Text(label="Negative Prompt", value="noisy,blurry,unattractive,sloppy,unprofessional,2D,text logos,signatures,out of frame,jpeg artifacts,ugly,poorly drawn,extra limbs,extra hands,extra feet,backwards limbs,extra fingers,extra toes,unrealistic, incorrect, bad anatomy,cut off body pieces,strange body positions, impossible body positioning, Mismatched eyes, cross eyed, crooked face, crooked lips, unclear, undefined, mutations, deformities,off center,poor_composition, duplicate faces, plastic,fake,negativity,blurry,blurred,doll,unclear,bad anatomy,lowres,worstquality,watermark,bad proportions,long neck,deformed,mutated,mutation,disfigured,poorly drawn face,skin blemishes,skin spots,acnes,missing limb,malformed limbs,floating limbs,disconnected limbs,extra limb,extra arms,poorly drawn hands,malformed hands,mutated hands and fingers,bad hands,missing fingers,fused fingers,too many fingers,extra legs,bad feet,cross-eyed,low quality,bad-hands-5,BadNegAnatomyV1-neg,EasyNegative,FastNegativeV2,bad-picture-chill-75v"), gr.Text(label="Seed", value="0"), gr.Text(label="Output Format", value='png')],
            outputs=["image",gr.Image(format="png", show_download_button=True)],
            title="Inpaint Characters",
            description="Start by uploading an image you've enhanced in the Generative Background Fill section. Use the masking tool to designate the specific area where you wish to add new characters or elements. Then, provide a detailed prompt for the inpainting process. This tool is designed to seamlessly blend new components into your existing image, significantly enriching its visual narrative and overall aesthetic appeal. It's ideal for creating compelling scenes or adding intricate details that align with your creative vision.",
            allow_flagging='never'
        )

    with gr.Tab("Gallery"):
        gr.Gallery(value=folder_contents, columns=[4], rows=[10], object_fit="contain", height="auto")

demo.launch(share=True)