LLaVA-Plus-Codebase/playground/llava-plus-data/grounding/generate.py

import os
import json
import time
from tqdm import tqdm
import fire

import openai

import concurrent.futures
import random
import json
import time
from collections import Counter
from functools import partial

from pycocotools.coco import COCO

import requests
from PIL import Image
import base64
import json
import time
from io import BytesIO

import torchvision.transforms.functional as F


# vars
controller_address = "http://localhost:21001"
model_name = 'grounding_dino'

def get_openai_api():
    # Set to your API key
    return {
        'api_type': '',
        'api_version': '2023-03-15-preview',
        'engine': "",
        'api_key': "",
        'api_base': '',
    }



def ask_gpt(messages, max_retries=35, temperature=0.2, top_p=0.9, max_tokens=512):
    if isinstance(messages, str):
        messages = [{"role": "user", "content": messages}]

    openai_kwargs = get_openai_api()

    for i in range(max_retries):
        try:
            response = openai.ChatCompletion.create(
                **openai_kwargs,
                messages=messages,
                temperature=temperature,
                max_tokens=max_tokens,
                top_p=top_p,
                frequency_penalty=0,
                presence_penalty=0,
                stop=None)
            if os.getenv('DEBUG_PRINT'):
                print(response['choices'][0]['message']['content'])
            return response['choices'][0]['message']['content']
        except Exception as e:
            if type(e) in [openai.error.InvalidRequestError, KeyError]:
                print(type(e), e)
                return None
            print(type(e), e)
            time.sleep(2)
            continue


def R(x):
    if isinstance(x, list):
        return [R(i) for i in x]
    elif isinstance(x, dict):
        return {k: R(v) for k, v in x.items()}
    elif isinstance(x, float):
        return round(x, 2)
    
def load_image(image_path):
    img = Image.open(image_path).convert('RGB')
    # import ipdb; ipdb.set_trace()
    # resize if needed
    w, h = img.size
    if max(h, w) > 800:
        if h > w:
            new_h = 800
            new_w = int(w * 800 / h)
        else:
            new_w = 800
            new_h = int(h * 800 / w)
        # import ipdb; ipdb.set_trace()
        img = F.resize(img, (new_h, new_w))
    return img

def encode(image: Image):
    buffered = BytesIO()
    image.save(buffered, format="JPEG")
    img_b64_str = base64.b64encode(buffered.getvalue()).decode()
    buffered.close()
    return img_b64_str

def get_worker_addr(controller_addr, model_name):
    # get worker_addr
    # ret = requests.post(controller_addr + "/refresh_all_workers")
    # ret = requests.post(controller_addr + "/list_models")
    # models = ret.json()["models"]
    # models.sort()
    # print(f"Models: {models}")

    ret = requests.post(
        controller_addr + "/get_worker_address", json={"model": model_name}
    )
    worker_addr = ret.json()["address"]
    del ret
    # print(f"worker_addr: {worker_addr}")
    return worker_addr

def generate_worker(captions_strs, objects_strs, examples, sample, image_dir):
    # 1. captions_strs + objects_strs -> questions
    # 2. questions -> grounding dino input
    # 3. grounding dino input -> grounding dino output
    # 4. captions_strs + objects_strs + questions + grounding dino output -> answer

    # 1. captions_strs + objects_strs -> questions
    messages = [
        {'role': 'system', 'content': """You are an AI visual assistant that can analyze a single image. You receive five sentences, each describing the same image you are observing. In addition, specific object locations within the image are given, along with detailed coordinates. These coordinates are in the form of bounding boxes, represented as (x1, y1, x2, y2) with floating numbers ranging from 0 to 1. These values correspond to the top left x, top left y, bottom right x, and bottom right y.
Generate a question that users may be interested to ask about the image. The question should ask the AI to detect some objects in the image. The question should be answerable by the given sentences and the given object locations.
The question should ask the AI to detect some objects in the image."""},
    {"role": "user", "content": examples[0]['captions']+'\n'+examples[0]['objects']},
    {"role": "assistant", "content": examples[0]['question']},
    {"role": "user", "content": examples[1]['captions']+'\n'+examples[1]['objects']},
    {"role": "assistant", "content": examples[1]['question']},
    {"role": "user", "content": captions_strs + '\n' + objects_strs}
    ]
    question = ask_gpt(messages, temperature=0.9, top_p=0.95)
    if question is None:
        print("question is None, return None")
        return None

    # 2. questions -> grounding dino input
    messages = [
        {'role': 'system', 'content': """You are an AI visual assistant that can help to extract information from an a sentence. 
You will be given a question about detecting something in an image. Please extract the main object name from the question. Using '.' to concat multiple object names."""},
    {"role": "user", "content": examples[0]['question']},
    {"role": "assistant", "content": examples[0]['grounding_dino_input']},
    {"role": "user", "content": examples[1]['question']},
    {"role": "assistant", "content": examples[1]['grounding_dino_input']},
    {"role": "user", "content": "Please detect the green car in the image."},
    {"role": "assistant", "content": "green car"},
    {"role": "user", "content": question}
    ]
    grounding_dino_input = ask_gpt(messages, temperature=0.9, top_p=0.95)
    if grounding_dino_input is None:
        print("grounding_dino_input is None, return None")
        return None

    # 3. grounding dino input -> grounding dino output
    # get grounding dino output
    
    worker_addr = get_worker_addr(controller_address, model_name)
    headers = {"User-Agent": "GSAM Client"}
    # img_path = os.path.join(args.image_dir, image['image_id'])
    img_path = os.path.join(image_dir, sample['file_name'])
    img = load_image(img_path)
    img_arg = encode(img)
    ret = requests.post(
            worker_addr + "/worker_generate",
            json={
                "image": img_arg,
                "caption": grounding_dino_input,
                "box_threshold": 0.3,
                "text_threshold": 0.25,
            },
            headers=headers,
        ).json()
    if os.getenv('DEBUG_PRINT'):
        print(ret)
    ret.pop("size")
    grounding_dino_output = ret

    # 4. captions_strs + objects_strs + questions + grounding dino output -> answer
    q_temp = "caption: {cap}\ngrounding dino input: {gdin}\ngrounding dino output: {gdout}\nquestion: {q}\n"
    messages = [
        {'role': 'system', 'content': """You are an AI visual assistant that can analyze a single image. 
You receive five sentences, each describing the same image you are observing. 
Then you receive the output of the grounding dino model, with its corresponding input of grounding dino. The output is a list of objects detected in the image, with their corresponding bounding boxes. These coordinates are in the form of bounding boxes, represented as (x1, y1, x2, y2) with floating numbers ranging from 0 to 1. These values correspond to the top left x, top left y, bottom right x, and bottom right y.
Then you receive the question asked by the user. 
Answer the question based on the given information with your best. 
Do not reveal the input information of the image. DO NOT say that you are given the captions and the objects in the image, JUST answer the question as if you are seeing the image for the first time."""},
        {'role': 'user', 'content': q_temp.format(cap=examples[0]['captions'], gdin=examples[0]['grounding_dino_input'], gdout=examples[0]['grounding_dino_output'], q=examples[0]['question'])},
        {'role': 'assistant', 'content': examples[0]['answer']},
        {'role': 'user', 'content': q_temp.format(cap=examples[1]['captions'], gdin=examples[1]['grounding_dino_input'], gdout=examples[1]['grounding_dino_output'], q=examples[1]['question'])},
        {'role': 'assistant', 'content': examples[1]['answer']},
        {'role': 'user', 'content': q_temp.format(cap=captions_strs, gdin=grounding_dino_input, gdout=grounding_dino_output, q=question)},
    ]
    answer = ask_gpt(messages, temperature=0.9, top_p=0.95)
    if answer is None:
        print("answer is None, return None")
        return None

    # return
    return {
        "unique_id": str(time.time()) + '_' + str(sample['id']),
        "image_id": sample['id'],
        "image_file_name": sample['file_name'],
        "image_path": os.path.join(image_dir, sample['file_name']),
        "captions": captions_strs,
        "objects": objects_strs,
        "question": question,
        "grounding_dino_input": grounding_dino_input,
        "grounding_dino_output": grounding_dino_output,
        "answer": answer,
    }


def generate_data(
    output_file,
    sample_json,
    overwrite=False,
    num_workers=1,
    num_examples=1000,
    coco_caption_path="/comp_robot/liushilong/data/coco/annotations/captions_{split}2014.json",
    coco_object_path="/comp_robot/liushilong/data/coco/annotations/instances_{split}2014.json",
    image_dir="/comp_robot/liushilong/data/coco/{split}2014",
    split='train',
    seed=23123,
    debug=False,
    reference_json=None,
):
    # load existing data
    if not overwrite and os.path.exists(output_file):
        print("Loading existing data...")
        with open(output_file) as f:
            existing_examples = [json.loads(line) for line in f]
        print("Existing data loaded.")
        if len(existing_examples) >= num_examples:
            print("Enough examples, skip generating.")
            return
        print("Generating {} examples...".format(num_examples - len(existing_examples)))
        num_examples = num_examples - len(existing_examples)
        seed = seed + len(existing_examples)


    # load coco annos
    coco_cap = COCO(coco_caption_path.format(split=split))
    coco_obj = COCO(coco_object_path.format(split=split))
    image_dir = image_dir.format(split=split)

    # load coco images
    coco_images = coco_cap.loadImgs(coco_cap.getImgIds())
    coco_categories = coco_obj.loadCats(coco_obj.getCatIds())

    # filter annos with reference json
    # load reference json
    if reference_json is not None:
        if seed != 20520:
            seed = 20520
            Warning("seed is not 20520, set seed to 20520!")
        with open(reference_json) as f:
            reference_examples = json.load(f)
        print("Loaded reference json, {} examples".format(len(reference_examples)))
        reference_ids = list(set([int(item['id']) for item in reference_examples]))
        coco_images = [item for item in coco_images if int(item['id']) in reference_ids]
        print("Filtered coco images with reference json, {} -> {}".format(len(reference_ids), len(coco_images)))
        # import ipdb; ipdb.set_trace()


    # random select 1000 images
    random.seed(seed)
    random.shuffle(coco_images)
    coco_images = coco_images[:num_examples]

    # load sample json
    with open(sample_json) as f:
        examples = json.load(f)

    # generate data
    print("Start generating data...")
    with concurrent.futures.ThreadPoolExecutor(max_workers=num_workers) as executor:
        futures = {}
        for sample_idx, sample in enumerate(coco_images):
            # load samples
            captions = coco_cap.loadAnns(coco_cap.getAnnIds(sample['id']))
            objects = coco_obj.loadAnns(coco_obj.getAnnIds(sample['id']))
            width, height = sample['width'], sample['height']
            for obj in objects:
                obj['bbox'] = [obj['bbox'][0] / width, obj['bbox'][1] / height, obj['bbox'][2] / width, obj['bbox'][3] / height]
                # xywh -> xyxy
                obj['bbox'][2] += obj['bbox'][0]
                obj['bbox'][3] += obj['bbox'][1]
                obj['bbox'] = R(obj['bbox'])

            captions_strs = "\n".join([cap['caption'].strip() for cap in captions])
            objects_strs = "\n".join([coco_obj.loadCats(obj['category_id'])[0]['name'] + ": " + str(obj['bbox']) for obj in objects])

            if debug:
                generate_worker(captions_strs, objects_strs, examples, sample, image_dir)
                continue


            futures[executor.submit(generate_worker, captions_strs, objects_strs, examples, sample, image_dir)] = sample_idx

        os.makedirs(os.path.dirname(output_file), exist_ok=True)
        writer = open(output_file, 'a')
        for future in tqdm(concurrent.futures.as_completed(futures), total=len(futures)):
            result = future.result()
            if result is None:
                time.sleep(0.1)
                continue
            writer.write(json.dumps(result) + '\n')
            writer.flush()
        writer.close()



def main(task, **kwargs):
    globals()[task](**kwargs)


if __name__ == "__main__":
    fire.Fire(main)