app_old.py

from io import BytesIO
import string
import gradio as gr
import requests
from caas import CaptionAnything
import torch
import json
import sys
import argparse
from caas import parse_augment
import os

# download sam checkpoint if not downloaded
def download_checkpoint(url, folder, filename):
    os.makedirs(folder, exist_ok=True)
    filepath = os.path.join(folder, filename)

    if not os.path.exists(filepath):
        response = requests.get(url, stream=True)
        with open(filepath, "wb") as f:
            for chunk in response.iter_content(chunk_size=8192):
                if chunk:
                    f.write(chunk)

    return filepath
checkpoint_url = "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth"
folder = "segmenter"
filename = "sam_vit_h_4b8939.pth"

title = """<h1 align="center">Caption-Anything</h1>"""
description = """Gradio demo for Caption Anything, image to dense captioning generation with various language styles. To use it, simply upload your image, or click one of the examples to load them.
<br> <strong>Code</strong>: GitHub repo: <a href='https://github.com/ttengwang/Caption-Anything' target='_blank'></a>
"""

examples = [
    ["test_img/img2.jpg", "[[1000, 700, 1]]"]
]

args = parse_augment()

def get_prompt(chat_input, click_state):    
    points = click_state[0]
    labels = click_state[1]
    inputs = json.loads(chat_input)
    for input in inputs:
        points.append(input[:2])
        labels.append(input[2])
    
    prompt = {
        "prompt_type":["click"],
        "input_point":points,
        "input_label":labels,
        "multimask_output":"True",
    }
    return prompt
    
def inference_seg_cap(image_input, chat_input, language, sentiment, factuality, length, state, click_state):
    controls = {'length': length,
             'sentiment': sentiment,
             'factuality': factuality,
             'language': language}
    prompt = get_prompt(chat_input, click_state)
    print('prompt: ', prompt, 'controls: ', controls)
    out = model.inference(image_input, prompt, controls)
    state = state + [(None, "Image point: {}, Input label: {}".format(prompt["input_point"], prompt["input_label"]))]
    for k, v in out['generated_captions'].items():
        state = state + [(f'{k}: {v}', None)]
    click_state[2].append(out['generated_captions']['raw_caption'])
    image_output_mask = out['mask_save_path']
    image_output_crop = out['crop_save_path']
    return state, state, click_state, image_output_mask, image_output_crop


def upload_callback(image_input, state):
    state = state + [('Image size: ' + str(image_input.size), None)]
    return state

# get coordinate in format [[x,y,positive/negative]]
def get_select_coords(image_input, point_prompt, language, sentiment, factuality, length, state, click_state, evt: gr.SelectData):
        print("point_prompt: ", point_prompt)
        if point_prompt == 'Positive Point':
            coordinate = "[[{}, {}, 1]]".format(str(evt.index[0]), str(evt.index[1]))
        else:
            coordinate = "[[{}, {}, 0]]".format(str(evt.index[0]), str(evt.index[1]))
        return (coordinate,) + inference_seg_cap(image_input, coordinate, language, sentiment, factuality, length, state, click_state)

def chat_with_points(chat_input, click_state, state):
    points, labels, captions = click_state
    point_chat_prompt = "I want you act as a chat bot in terms of image. I will give you some points (w, h) in the image and tell you what happed on the point in natural language. Note that (0, 0) refers to the top-left corner of the image, w refers to the width and h refers the height. You should chat with me based on the fact in the image instead of imagination. Now I tell you the points with their visual description:\n{points_with_caps}\n. Now begin chatting! Human: {chat_input}\nAI: "
    # "The image is of width {width} and height {height}." 
    
    prev_visual_context = ""
    pos_points = [f"{points[i][0]}, {points[i][1]}" for i in range(len(points)) if labels[i] == 1]
    prev_visual_context = ', '.join(pos_points) + captions[-1] + '\n'
    chat_prompt = point_chat_prompt.format(**{"points_with_caps": prev_visual_context, "chat_input": chat_input})
    response = model.text_refiner.llm(chat_prompt)
    state = state + [(chat_input, response)]
    return state, state

def init_openai_api_key(api_key):
    os.environ['OPENAI_API_KEY'] = api_key
    global model
    model = CaptionAnything(args)

css='''
#image_upload{min-height:200px}
#image_upload [data-testid="image"], #image_upload [data-testid="image"] > div{min-height: 200px}
'''

with gr.Blocks(css=css) as iface:
    state = gr.State([])
    click_state = gr.State([[],[],[]])
    caption_state = gr.State([[]])
    gr.Markdown(title)
    gr.Markdown(description)

    with gr.Column():
        openai_api_key = gr.Textbox(
            placeholder="Input your openAI API key and press Enter",
            show_label=False,
            lines=1,
            type="password",
        )
        openai_api_key.submit(init_openai_api_key, inputs=[openai_api_key])
        
        with gr.Row():
            with gr.Column(scale=0.7):
                image_input = gr.Image(type="pil", interactive=True, label="Image", elem_id="image_upload").style(height=260,scale=1.0)

                with gr.Row(scale=0.7):
                    point_prompt = gr.Radio(
                        choices=["Positive Point",  "Negative Point"],
                        value="Positive Point",
                        label="Points",
                        interactive=True,
                    )
                
                # with gr.Row():
                language = gr.Radio(
                    choices=["English", "Chinese", "French", "Spanish", "Arabic", "Portuguese","Cantonese"],
                    value="English",
                    label="Language",
                    interactive=True,
                )
                sentiment = gr.Radio(
                    choices=["Positive", "Natural", "Negative"],
                    value="Natural",
                    label="Sentiment",
                    interactive=True,
                )
                factuality = gr.Radio(
                    choices=["Factual", "Imagination"],
                    value="Factual",
                    label="Factuality",
                    interactive=True,
                )
                length = gr.Slider(
                    minimum=5,
                    maximum=100,
                    value=10,
                    step=1,
                    interactive=True,
                    label="Length",
                )

            with gr.Column(scale=1.5):
                with gr.Row():
                    image_output_mask= gr.Image(type="pil", interactive=False, label="Mask").style(height=260,scale=1.0)
                    image_output_crop= gr.Image(type="pil", interactive=False, label="Cropped Image by Mask", show_progress=False).style(height=260,scale=1.0)
                chatbot = gr.Chatbot(label="Chat Output",).style(height=450,scale=0.5)
        
        with gr.Row():
            with gr.Column(scale=0.7):
                prompt_input = gr.Textbox(lines=1, label="Input Prompt (A list of points like : [[100, 200, 1], [200,300,0]])")
                prompt_input.submit(
                    inference_seg_cap,
                    [
                        image_input,
                        prompt_input,
                        language,
                        sentiment,
                        factuality,
                        length,
                        state,
                        click_state
                    ],
                    [chatbot, state, click_state, image_output_mask, image_output_crop],
                    show_progress=False
                )
                
                image_input.upload(
                    upload_callback,
                    [image_input, state],
                    [chatbot]
                    )

                with gr.Row():
                    clear_button = gr.Button(value="Clear Click", interactive=True)
                    clear_button.click(
                        lambda: ("", [[], [], []], None, None),
                        [],
                        [prompt_input, click_state, image_output_mask, image_output_crop],
                        queue=False,
                        show_progress=False
                    )
                    
                    clear_button = gr.Button(value="Clear", interactive=True)
                    clear_button.click(
                        lambda: ("", [], [], [[], [], []], None, None),
                        [],
                        [prompt_input, chatbot, state, click_state, image_output_mask, image_output_crop],
                        queue=False,
                        show_progress=False
                    )
                    
                    submit_button = gr.Button(
                        value="Submit", interactive=True, variant="primary"
                    )
                    submit_button.click(
                        inference_seg_cap,
                        [
                            image_input,
                            prompt_input,
                            language,
                            sentiment,
                            factuality,
                            length,
                            state,
                            click_state
                        ],
                        [chatbot, state, click_state, image_output_mask, image_output_crop],
                        show_progress=False
                    )
                    
                # select coordinate
                image_input.select(
                    get_select_coords, 
                    inputs=[image_input,point_prompt,language,sentiment,factuality,length,state,click_state], 
                    outputs=[prompt_input, chatbot, state, click_state, image_output_mask, image_output_crop],
                    show_progress=False
                    )

                image_input.change(
                    lambda: ("", [], [[], [], []]),
                    [],
                    [chatbot, state, click_state],
                    queue=False,
                )
                
            with gr.Column(scale=1.5):
                chat_input = gr.Textbox(lines=1, label="Chat Input")
                chat_input.submit(chat_with_points, [chat_input, click_state, state], [chatbot, state])
                
                    
    examples = gr.Examples(
        examples=examples,
        inputs=[image_input, prompt_input],
    )

iface.queue(concurrency_count=1, api_open=False, max_size=10)
iface.launch(server_name="0.0.0.0", enable_queue=True, server_port=args.port, share=args.gradio_share)