serve/gradio_web_bbox.py

import gradio as gr
import time
import numpy as np
from gradio import processing_utils
from PIL import Image, ImageDraw, ImageFont

from serve.utils import *
from serve.vote_utils import (
    upvote_last_response_ig as upvote_last_response,
    downvote_last_response_ig as downvote_last_response,
    flag_last_response_ig as flag_last_response,
    leftvote_last_response_igm as leftvote_last_response,
    left1vote_last_response_igm as left1vote_last_response,
    rightvote_last_response_igm as rightvote_last_response,
    right1vote_last_response_igm as right1vote_last_response,
    tievote_last_response_igm as tievote_last_response,
    bothbad_vote_last_response_igm as bothbad_vote_last_response,
    share_click_igm as share_click,
    generate_ig,
    generate_ig_museum,
    generate_igm,
    generate_igm_museum,
    generate_igm_annoy,
    generate_b2i_annoy,
    generate_igm_annoy_museum,
    generate_igm_cache_annoy,
    share_js
)
from serve.Ksort import (
    add_foreground,
    reset_level,
    reset_rank,
    revote_windows,
    submit_response_igm,
    submit_response_rank_igm,
    reset_submit,
    clear_rank,
    reset_mode,
    reset_chatbot,
    reset_btn_rank,
    reset_vote_text,
    text_response_rank_igm,
    check_textbox,
)

from functools import partial
from serve.upload import get_random_mscoco_prompt
from serve.constants import SSH_SERVER, SSH_PORT, SSH_USER, SSH_PASSWORD
from serve.upload import get_random_mscoco_prompt, create_ssh_client
from serve.update_skill import create_ssh_skill_client
from model.matchmaker import create_ssh_matchmaker_client

def set_ssh():
    create_ssh_client(SSH_SERVER, SSH_PORT, SSH_USER, SSH_PASSWORD)
    create_ssh_skill_client(SSH_SERVER, SSH_PORT, SSH_USER, SSH_PASSWORD)
    create_ssh_matchmaker_client(SSH_SERVER, SSH_PORT, SSH_USER, SSH_PASSWORD)

def binarize(x):
    return (x != 0).astype('uint8') * 255

def sized_center_crop(img, cropx, cropy):
    y, x = img.shape[:2]
    startx = x // 2 - (cropx // 2)
    starty = y // 2 - (cropy // 2)    
    return img[starty:starty+cropy, startx:startx+cropx]

def sized_center_fill(img, fill, cropx, cropy):
    y, x = img.shape[:2]
    startx = x // 2 - (cropx // 2)
    starty = y // 2 - (cropy // 2)    
    img[starty:starty+cropy, startx:startx+cropx] = fill
    return img

def sized_center_mask(img, cropx, cropy):
    y, x = img.shape[:2]
    startx = x // 2 - (cropx // 2)
    starty = y // 2 - (cropy // 2)    
    center_region = img[starty:starty+cropy, startx:startx+cropx].copy()
    img = (img * 0.2).astype('uint8')
    img[starty:starty+cropy, startx:startx+cropx] = center_region
    return img

def center_crop(img, HW=None, tgt_size=(512, 512)):
    if HW is None:
        H, W = img.shape[:2]
        HW = min(H, W)
    img = sized_center_crop(img, HW, HW)
    img = Image.fromarray(img)
    img = img.resize(tgt_size)
    return np.array(img)



def draw_box(boxes=[], labels=[], img=None):
    if len(boxes) == 0 and img is None:
        return None
    
    # 确保输入的 img 是 PIL.Image 对象
    if isinstance(img, np.ndarray):
        img = Image.fromarray(img)

    if img is None:
        img = Image.new('RGB', (512, 512), (255, 255, 255))

    colors = ["red", "olive", "blue", "green", "orange", "brown", "cyan", "purple"]

    # 创建绘图对象
    draw_obj = ImageDraw.Draw(img)
    font = ImageFont.load_default()
    # 获取字体大小
    font_size = getattr(font, 'size_in_points', 10)  # 如果没有 size_in_points 属性，使用默认值 10

    for bid, box in enumerate(boxes):
        draw_obj.rectangle([box[0], box[1], box[2], box[3]], outline=colors[bid % len(colors)], width=4)
        anno_text = labels[bid]
        draw_obj.rectangle(
            [
                box[0], 
                box[3] - int(font_size * 1.2), 
                box[0] + int((len(anno_text) + 0.8) * font_size * 0.6), 
                box[3]
            ], 
            outline=colors[bid % len(colors)], 
            fill=colors[bid % len(colors)], 
            width=4
        )
        draw_obj.text(
            [box[0] + int(font_size * 0.2), box[3] - int(font_size * 1.2)], 
            anno_text, 
            font=font, 
            fill=(255,255,255)
        )
    return img


def draw(input, grounding_texts, new_image_trigger, state):
    
    # 确保输入数据中有必要的键
    if isinstance(input, dict):
        background = input.get('background', None)
        print("background.shape", background.shape)
        layers = input.get('layers', 1)
        print("len(layers)", len(layers))
        composite = input.get('composite', None)
        print("composite.shape", composite.shape)
    else:
        # 如果 input 不是字典，直接使用其作为图像
        background = input
        layers = 1
        composite = None

    # 检查 background 是否有效
    if background is None:
        print("background is None")
        return None
    # background = np.ones((512, 512, 3), dtype='uint8') * 255

    # 默认使用 composite 作为最终图像，如果没有 composite 则使用 background
    if composite is None:
        print("composite is None")
        image = background
    else:
        image = composite

    mask = binarize(image)

    if type(mask) != np.ndarray:
        mask = np.array(mask)

    if mask.sum() == 0:
        state = {}

    # 更新状态，如果没有 boxes 和 masks，则初始化它们
    if 'boxes' not in state:
        state['boxes'] = []

    if 'masks' not in state or len(state['masks']) == 0:
        state['masks'] = []
        last_mask = np.zeros_like(mask)
    else:
        last_mask = state['masks'][-1]

    if type(mask) == np.ndarray and mask.size > 1:
        diff_mask = mask - last_mask
    else:
        diff_mask = np.zeros([])

    # 根据 mask 的变化来计算 box 的位置
    if diff_mask.sum() > 0:
        x1x2 = np.where(diff_mask.max(0) != 0)[0]
        y1y2 = np.where(diff_mask.max(1) != 0)[0]
        y1, y2 = y1y2.min(), y1y2.max()
        x1, x2 = x1x2.min(), x1x2.max()

        if (x2 - x1 > 5) and (y2 - y1 > 5):
            state['masks'].append(mask.copy())
            state['boxes'].append((x1, y1, x2, y2))

    # 处理 grounding_texts
    grounding_texts = [x.strip() for x in grounding_texts.split(';')]
    grounding_texts = [x for x in grounding_texts if len(x) > 0]
    if len(grounding_texts) < len(state['boxes']):
        grounding_texts += [f'Obj. {bid+1}' for bid in range(len(grounding_texts), len(state['boxes']))]

    # 绘制标注框
    box_image = draw_box(state['boxes'], grounding_texts, background)

    if box_image is not None and state.get('inpaint_hw', None):
        inpaint_hw = state['inpaint_hw']
        box_image_resize = np.array(box_image.resize((inpaint_hw, inpaint_hw)))
        original_image = state['original_image'].copy()
        box_image = sized_center_fill(original_image, box_image_resize, inpaint_hw, inpaint_hw)

    return [box_image, new_image_trigger, 1.0, state]

def build_side_by_side_bbox_ui_anony(models):
    notice_markdown = """
    # ⚔️  Control-Ability-Arena (Bbox-to-Image Generation) ⚔️
    ## 📜 Rules
    - Input a prompt for four anonymized models and vote on their outputs.
    - Two voting modes available: Rank Mode and Best Mode. Switch freely between modes. Please note that ties are always allowed. In ranking mode, users can input rankings like 1 3 3 1. Any invalid rankings, such as 1 4 4 1, will be automatically corrected during post-processing.
    - Users are encouraged to make evaluations based on subjective preferences. Evaluation criteria: Alignment (50%) + Aesthetics (50%).
        - Alignment includes: Entity Matching (30%) + Style Matching (20%);
        - Aesthetics includes: Photorealism (30%) + Light and Shadow (10%) + Absence of Artifacts (10%).

    ## 👇 Generating now!
    - Note: Due to the API's image safety checks, errors may occur. If this happens, please re-enter a different prompt.
    - At times, high API concurrency can cause congestion, potentially resulting in a generation time of up to 1.5 minutes per image. Thank you for your patience.
    """
    model_list = models.model_b2i_list

    state = gr.State({})
    state0 = gr.State()
    state1 = gr.State()
    state2 = gr.State()
    state3 = gr.State()

    gen_func = partial(generate_b2i_annoy, models.generate_image_b2i_parallel_anony)
    # gen_cache_func = partial(generate_igm_cache_annoy, models.generate_image_ig_cache_anony)
    
   
    gr.Markdown(notice_markdown, elem_id="notice_markdown")


    with gr.Row():
        sketch_pad_trigger = gr.Number(value=0, visible=False)
        sketch_pad_resize_trigger = gr.Number(value=0, visible=False)
        image_scale = gr.Number(value=0, elem_id="image_scale", visible=False)

    with gr.Row():
        sketch_pad = gr.ImageEditor(
            label="Sketch Pad", 
            type="numpy",
            crop_size="1:1",
            width=512,
            height=512
        )
        out_imagebox = gr.Image(
            type="pil",
            label="Parsed Sketch Pad",
            width=512,
            height=512
        )

    with gr.Row():
        textbox = gr.Textbox(
            show_label=False,
            placeholder="👉 Enter your prompt and press ENTER",
            container=True,
            elem_id="input_box",
        )
        send_btn = gr.Button(value="Send", variant="primary", scale=0, elem_id="btnblue")

    with gr.Row():
        grounding_instruction = gr.Textbox(
            label="Grounding instruction (Separated by semicolon)",
            placeholder="👉 Enter your Grounding instruction (e.g. a cat; a dog; a bird; a fish)",
        )

    with gr.Group(elem_id="share-region-anony"):
        with gr.Accordion("🔍 Expand to see all Arena players", open=False):
            # model_description_md = get_model_description_md(model_list)
            gr.Markdown("", elem_id="model_description_markdown")

        
        with gr.Row():
            with gr.Column():
                chatbot_left = gr.Image(width=512, label = "Model A")
            with gr.Column():
                chatbot_left1 = gr.Image(width=512, label = "Model B")
            with gr.Column():
                chatbot_right = gr.Image(width=512, label = "Model C")
            with gr.Column():
                chatbot_right1 = gr.Image(width=512, label = "Model D")

        with gr.Row():
            with gr.Column():
                model_selector_left = gr.Markdown("", visible=False)
            with gr.Column():
                model_selector_left1 = gr.Markdown("", visible=False)
            with gr.Column():
                model_selector_right = gr.Markdown("", visible=False)
            with gr.Column():
                model_selector_right1 = gr.Markdown("", visible=False)
        with gr.Row():
            slow_warning = gr.Markdown("", elem_id="notice_markdown")

        with gr.Row(elem_classes="row"):
            with gr.Column(scale=1, min_width=10):
                leftvote_btn = gr.Button(
                    value="A is Best", visible=False, interactive=False, elem_id="btncolor1", elem_classes="best-button"
                )
            with gr.Column(scale=1, min_width=10):
                left1vote_btn = gr.Button(
                    value="B is Best", visible=False, interactive=False, elem_id="btncolor1", elem_classes="best-button"
                )
            with gr.Column(scale=1, min_width=10):
                rightvote_btn = gr.Button(
                    value="C is Best", visible=False, interactive=False, elem_id="btncolor1", elem_classes="best-button"
                )
            with gr.Column(scale=1, min_width=10):
                right1vote_btn = gr.Button(
                    value="D is Best", visible=False, interactive=False, elem_id="btncolor1", elem_classes="best-button"
                )
            with gr.Column(scale=1, min_width=10):
                tie_btn = gr.Button(
                    value="🤝  Tie", visible=False, interactive=False, elem_id="btncolor2", elem_classes="best-button"
                )


        with gr.Row():
            with gr.Blocks():
                with gr.Row():
                    with gr.Column(scale=1, min_width=10):
                        A1_btn = gr.Button(
                            value="1", visible=False, interactive=False, elem_id="btncolor1", elem_classes="custom-button"
                        )
                    with gr.Column(scale=1, min_width=10):
                        A2_btn = gr.Button(
                            value="2", visible=False, interactive=False, elem_id="btncolor2", elem_classes="custom-button"
                        )
                    with gr.Column(scale=1, min_width=10):
                        A3_btn = gr.Button(
                            value="3", visible=False, interactive=False, elem_id="btncolor3", elem_classes="custom-button"
                        )
                    with gr.Column(scale=1, min_width=10):
                        A4_btn = gr.Button(
                            value="4", visible=False, interactive=False, elem_id="btncolor4", elem_classes="custom-button"
                        )
            with gr.Blocks(): 
                with gr.Row():
                    with gr.Column(scale=1, min_width=10):
                        B1_btn = gr.Button(
                            value="1", visible=False, interactive=False, elem_id="btncolor1", elem_classes="custom-button"
                        )
                    with gr.Column(scale=1, min_width=10):
                        B2_btn = gr.Button(
                            value="2", visible=False, interactive=False, elem_id="btncolor2", elem_classes="custom-button"
                        )
                    with gr.Column(scale=1, min_width=10):
                        B3_btn = gr.Button(
                            value="3", visible=False, interactive=False, elem_id="btncolor3", elem_classes="custom-button"
                        )
                    with gr.Column(scale=1, min_width=10):
                        B4_btn = gr.Button(
                            value="4", visible=False, interactive=False, elem_id="btncolor4", elem_classes="custom-button"
                        )
            with gr.Blocks(): 
                with gr.Row():
                    with gr.Column(scale=1, min_width=10):
                        C1_btn = gr.Button(
                            value="1", visible=False, interactive=False, elem_id="btncolor1", elem_classes="custom-button"
                        )
                    with gr.Column(scale=1, min_width=10):
                        C2_btn = gr.Button(
                            value="2", visible=False, interactive=False, elem_id="btncolor2", elem_classes="custom-button"
                        )
                    with gr.Column(scale=1, min_width=10):
                        C3_btn = gr.Button(
                            value="3", visible=False, interactive=False, elem_id="btncolor3", elem_classes="custom-button"
                        )
                    with gr.Column(scale=1, min_width=10):
                        C4_btn = gr.Button(
                            value="4", visible=False, interactive=False, elem_id="btncolor4", elem_classes="custom-button"
                        )
            with gr.Blocks(): 
                with gr.Row():
                    with gr.Column(scale=1, min_width=10):
                        D1_btn = gr.Button(
                            value="1", visible=False, interactive=False, elem_id="btncolor1", elem_classes="custom-button"
                        )
                    with gr.Column(scale=1, min_width=10):
                        D2_btn = gr.Button(
                            value="2", visible=False, interactive=False, elem_id="btncolor2", elem_classes="custom-button"
                        )
                    with gr.Column(scale=1, min_width=10):
                        D3_btn = gr.Button(
                            value="3", visible=False, interactive=False, elem_id="btncolor3", elem_classes="custom-button"
                        )
                    with gr.Column(scale=1, min_width=10):
                        D4_btn = gr.Button(
                            value="4", visible=False, interactive=False, elem_id="btncolor4", elem_classes="custom-button"
                        )
    with gr.Row():
        vote_textbox = gr.Textbox(
            show_label=False,
            placeholder="👉 Enter your rank (you can use buttons above, or directly type here, e.g. 1 2 3 4)",
            container=True,
            elem_id="input_box",
            visible=False,
        )
        vote_submit_btn = gr.Button(value="Submit", visible=False, interactive=False, variant="primary", scale=0, elem_id="btnpink", elem_classes="submit-button")
        vote_mode_btn = gr.Button(value="🔄 Mode", visible=False, interactive=False, variant="primary", scale=0, elem_id="btnpink", elem_classes="submit-button")

    with gr.Row():
        clear_btn = gr.Button(value="🎲 New Round", interactive=False)
        # regenerate_btn = gr.Button(value="🔄  Regenerate", interactive=False)
        # share_btn = gr.Button(value="📷  Share")
    with gr.Blocks(): 
        with gr.Row(elem_id="centered-text"): # 
            user_info = gr.Markdown("User information (to appear on the contributor leaderboard)", visible=True, elem_id="centered-text") #, elem_id="centered-text"
            # with gr.Blocks(): 
            # name = gr.Markdown("Name", visible=True)
            user_name = gr.Textbox(show_label=False,placeholder="👉 Enter your name (optional)", elem_classes="custom-width")
            # with gr.Blocks(): 
            # institution = gr.Markdown("Institution", visible=True)
            user_institution = gr.Textbox(show_label=False,placeholder="👉 Enter your affiliation (optional)", elem_classes="custom-width")
    
    sketch_pad.change(
        draw,
        inputs=[sketch_pad, grounding_instruction, sketch_pad_resize_trigger, state],
        outputs=[out_imagebox, sketch_pad_resize_trigger, image_scale, state],
        queue=False,
    )
    grounding_instruction.change(
        draw,
        inputs=[sketch_pad, grounding_instruction, sketch_pad_resize_trigger, state],
        outputs=[out_imagebox, sketch_pad_resize_trigger, image_scale, state],
        queue=False,
    )

    order_btn_list = [textbox, send_btn, clear_btn, grounding_instruction, sketch_pad, out_imagebox]
    vote_order_list = [leftvote_btn, left1vote_btn, rightvote_btn, right1vote_btn, tie_btn, \
                       A1_btn, A2_btn, A3_btn, A4_btn, B1_btn, B2_btn, B3_btn, B4_btn, C1_btn, C2_btn, C3_btn, C4_btn, D1_btn, D2_btn, D3_btn, D4_btn, \
                       vote_textbox, vote_submit_btn, vote_mode_btn]
    
    generate_ig0 = gr.Image(width=512, label = "generate A", visible=False, interactive=False)
    generate_ig1 = gr.Image(width=512, label = "generate B", visible=False, interactive=False)
    generate_ig2 = gr.Image(width=512, label = "generate C", visible=False, interactive=False)
    generate_ig3 = gr.Image(width=512, label = "generate D", visible=False, interactive=False)
    dummy_left_model = gr.State("")
    dummy_left1_model = gr.State("")
    dummy_right_model = gr.State("")
    dummy_right1_model = gr.State("")

    ig_rank = [None, None, None, None]
    bastA_rank = [0, 3, 3, 3]
    bastB_rank = [3, 0, 3, 3]
    bastC_rank = [3, 3, 0, 3]
    bastD_rank = [3, 3, 3, 0]
    tie_rank = [0, 0, 0, 0]
    bad_rank = [3, 3, 3, 3]
    rank = gr.State(ig_rank)
    rankA = gr.State(bastA_rank)
    rankB = gr.State(bastB_rank)
    rankC = gr.State(bastC_rank)
    rankD = gr.State(bastD_rank)
    rankTie = gr.State(tie_rank)
    rankBad = gr.State(bad_rank)
    Top1_text = gr.Textbox(value="Top 1", visible=False, interactive=False)
    Top2_text = gr.Textbox(value="Top 2", visible=False, interactive=False)
    Top3_text = gr.Textbox(value="Top 3", visible=False, interactive=False)
    Top4_text = gr.Textbox(value="Top 4", visible=False, interactive=False)
    window1_text = gr.Textbox(value="Model A", visible=False, interactive=False)
    window2_text = gr.Textbox(value="Model B", visible=False, interactive=False)
    window3_text = gr.Textbox(value="Model C", visible=False, interactive=False)
    window4_text = gr.Textbox(value="Model D", visible=False, interactive=False)
    vote_level = gr.Number(value=0, visible=False, interactive=False)
    # Top1_btn.click(reset_level, inputs=[Top1_text], outputs=[vote_level])
    # Top2_btn.click(reset_level, inputs=[Top2_text], outputs=[vote_level])
    # Top3_btn.click(reset_level, inputs=[Top3_text], outputs=[vote_level])
    # Top4_btn.click(reset_level, inputs=[Top4_text], outputs=[vote_level])
    vote_mode = gr.Textbox(value="Rank", visible=False, interactive=False)
    right_vote_text = gr.Textbox(value="wrong", visible=False, interactive=False)
    cache_mode = gr.Textbox(value="True", visible=False, interactive=False)


    textbox.submit(
        disable_order_buttons,
        inputs=[textbox],
        outputs=order_btn_list
    ).then(
        gen_func,
        inputs=[state0, state1, state2, state3, textbox, grounding_instruction, state, model_selector_left, model_selector_left1, model_selector_right, model_selector_right1],
        outputs=[state0, state1, state2, state3, generate_ig0, generate_ig1, generate_ig2, generate_ig3, chatbot_left, chatbot_left1, chatbot_right, chatbot_right1, \
                 model_selector_left, model_selector_left1, model_selector_right, model_selector_right1],
        api_name="submit_btn_annony"
    ).then(
        enable_vote_mode_buttons,
        inputs=[vote_mode, textbox],
        outputs=vote_order_list
    )   


if __name__ == "__main__":
    with gr.Blocks() as demo:
        build_side_by_side_bbox_ui_anony()
    demo.launch()