multimodal/build/lib/open_flamingo/eval/evaluate_debug.py

import argparse
import json
from math import ceil
import os
import random
import uuid
from collections import defaultdict
from typing import Callable
import time
import cv2

import more_itertools
import numpy as np
import torch
from coco_metric import compute_cider, postprocess_captioning_generation
from eval_datasets import VQADataset, GQADataset
from tqdm import tqdm
from collections import Counter

from vqa_metric import compute_vqa_accuracy, compute_gqa_accuracy
from open_flamingo.eval.classification import (
    compute_per_sample_probs,
    compute_per_sample_loss,
)
from open_flamingo.eval.imagenet_utils import (
    openai_imagenet_classnames,
    IMAGENET_1K_CLASS_ID_TO_LABEL,
)

from open_flamingo.src.factory import create_model_and_transforms
from PIL import Image
from io import BytesIO
import base64
from open_flamingo.train.distributed import init_distributed_device, world_info_from_env
import string
from lavis.datasets.builders import load_dataset


def get_iou(box1, box2):
    # box1 and box2 should be in the format [x1, y1, x2, y2]
    intersection = max(0, min(box1[2], box2[2]) - max(box1[0], box2[0])) * \
                   max(0, min(box1[3], box2[3]) - max(box1[1], box2[1]))
    area_box1 = (box1[2] - box1[0]) * (box1[3] - box1[1])
    area_box2 = (box2[2] - box2[0]) * (box2[3] - box2[1])
    union = area_box1 + area_box2 - intersection
    iou = intersection / union if union > 0 else 0
    return iou

def expand2square(pil_img, background_color):
    width, height = pil_img.size
    if width == height:
        return pil_img
    elif width > height:
        result = Image.new(pil_img.mode, (width, width), background_color)
        result.paste(pil_img, (0, (width - height) // 2))
        return result
    else:
        result = Image.new(pil_img.mode, (height, height), background_color)
        result.paste(pil_img, ((height - width) // 2, 0))
        return result

parser = argparse.ArgumentParser()
parser.add_argument("--lm_path", type=str, default="facebook/opt-1.3b")
parser.add_argument("--lm_tokenizer_path", type=str, default="facebook/opt-30b")
parser.add_argument("--vision_encoder_path", default="ViT-L-14", type=str)
parser.add_argument("--vision_encoder_pretrained", default="openai", type=str)
parser.add_argument("--checkpoint_path", type=str, required=True)
parser.add_argument(
    "--results_file", type=str, default=None, help="JSON file to save results"
)

# Trial arguments
parser.add_argument("--shots", nargs="+", default=[0, 4, 8, 16, 32], type=int)
parser.add_argument(
    "--num_trials",
    type=int,
    default=1,
    help="Number of trials to run for each shot using different demonstrations",
)
parser.add_argument(
    "--trial_seeds",
    nargs="+",
    default=[0],
    help="Seeds to use for each trial for picking demonstrations and eval sets",
)
parser.add_argument(
    "--num_samples", type=int, default=5000, help="Number of samples to evaluate on"
)

parser.add_argument("--batch_size", type=int, default=8)

# Per-dataset evaluation flags
parser.add_argument(
    "--eval_coco",
    action="store_true",
    default=False,
    help="Whether to evaluate on COCO.",
)
parser.add_argument(
    "--eval_vqav2",
    action="store_true",
    default=False,
    help="Whether to evaluate on VQAV2.",
)
parser.add_argument(
    "--eval_ok_vqa",
    action="store_true",
    default=False,
    help="Whether to evaluate on OK-VQA.",
)
parser.add_argument(
    "--eval_imagenet",
    action="store_true",
    default=False,
    help="Whether to evaluate on ImageNet.",
)

parser.add_argument(
    "--eval_flickr30",
    action="store_true",
    default=False,
    help="Whether to evaluate on Flickr30.",
)

parser.add_argument(
    "--eval_refcoco",
    action="store_true",
    default=False,
    help="Whether to evaluate on RefCOCO.",
)

# Dataset arguments

## Flickr30 Dataset
parser.add_argument(
    "--flickr_image_dir_path",
    type=str,
    help="Path to the flickr30/flickr30k_images directory.",
    default=None,
)
parser.add_argument(
    "--flickr_annotations_json_path",
    type=str,
    help="Path to the dataset_flickr30k_coco_style.json file.",
    default=None,
)

## COCO Dataset
parser.add_argument(
    "--coco_image_dir_path",
    type=str,
    help="Path to the flickr30/flickr30k_images directory.",
    default=None,
)
parser.add_argument(
    "--coco_annotations_json_path",
    type=str,
    default=None,
)

## VQAV2 Dataset
parser.add_argument(
    "--vqav2_image_dir_path",
    type=str,
    default=None,
)
parser.add_argument(
    "--vqav2_questions_json_path",
    type=str,
    default=None,
)
parser.add_argument(
    "--vqav2_annotations_json_path",
    type=str,
    default=None,
)

## OK-VQA Dataset
parser.add_argument(
    "--ok_vqa_image_dir_path",
    type=str,
    help="Path to the vqav2/train2014 directory.",
    default=None,
)
parser.add_argument(
    "--ok_vqa_questions_json_path",
    type=str,
    help="Path to the v2_OpenEnded_mscoco_train2014_questions.json file.",
    default=None,
)
parser.add_argument(
    "--ok_vqa_annotations_json_path",
    type=str,
    help="Path to the v2_mscoco_train2014_annotations.json file.",
    default=None,
)

## Imagenet dataset
parser.add_argument("--imagenet_root", type=str, default="/tmp")

## RefCOCO dataset
parser.add_argument("--refcoco_tsvfile", type=str, default=None)

parser.add_argument(
    "--location_token_num",
    default=1000,
    type=int,
)
# distributed training
parser.add_argument(
    "--dist-url",
    default="env://",
    type=str,
    help="url used to set up distributed training",
)
parser.add_argument(
    "--dist-backend", default="nccl", type=str, help="distributed backend"
)
parser.add_argument(
    "--horovod",
    default=False,
    action="store_true",
    help="Use horovod for distributed training.",
)
parser.add_argument(
    "--no-set-device-rank",
    default=False,
    action="store_true",
    help="Don't set device index from local rank (when CUDA_VISIBLE_DEVICES restricted to one per proc).",
)
parser.add_argument(
    "--dist",
    default=False,
    action="store_true",
)
parser.add_argument(
    "--lora",
    default=False,
    action="store_true",
)
parser.add_argument(
    "--lora_r",
    default=16,
    type=int,
    required=False,
)
parser.add_argument(
    "--legacy",
    default=False,
    action="store_true",
)
parser.add_argument(
    "--special",
    default=False,
    action="store_true",
)
parser.add_argument(
    "--id",
    default=0,
    type=int,
    required=False,
)

parser.add_argument(
    "--eval_gqa",
    default=False,
    action="store_true",
)
parser.add_argument(
    "--use_sam",
    default=None,
    type=str,
    required=False,
)
parser.add_argument(
    "--add_visual_token",
    default=False,
    action="store_true",
)
parser.add_argument(
    "--use_format_v2",
    default=False,
    action="store_true",
)


class OKVQAPostProcess():
    def __init__(self):
        self._lemmatizer = None

    def _lemmatize(self, answers):
        def apply(answer):
            doc = self.lemmatizer(answer)

            words = []
            for token in doc:
                if token.pos_ in ["NOUN", "VERB"]:
                    words.append(token.lemma_)
                else:
                    words.append(token.text)
            answer = " ".join(words)

            return answer

        return [apply(answer) for answer in answers]

    @property
    def lemmatizer(self):
        if self._lemmatizer is None:
            try:
                import spacy

                self._lemmatizer = spacy.load("en_core_web_sm")
            except ImportError:
                logging.error(
                    """
                    Please install spacy and en_core_web_sm model to apply lemmatization.
                    python -m spacy download en_core_web_sm
                    OR
                    import spacy.cli
                    spacy.cli.download("en_core_web_sm")
                    """
                )
                exit(1)

        return self._lemmatizer
        

def main():
    args = parser.parse_args()
    if args.dist:
        args.local_rank, args.rank, args.world_size = world_info_from_env()
        print(f"local_rank: {args.local_rank} rank: {args.rank} world_size: {args.world_size}")
        device_id = init_distributed_device(args)
    else:
        args.rank = 0
        args.world_size = 1
        print(f"rank: {args.rank} world_size: {args.world_size}")
    
    if "sam" in args.checkpoint_path:
        args.use_sam = "vit_l"

    args.add_visual_token = True
    if "lora" in args.checkpoint_path:
        args.lora = True


    args.add_pe = False
    args.add_box = False
    args.relation = False
    if "debug" in args.checkpoint_path:
        # args.add_pe = True
        args.add_box = True
    if "box" in args.checkpoint_path:
        args.add_box = True
    if "pe" in args.checkpoint_path:
        args.add_pe = True
    if "rel" in args.checkpoint_path:
        args.relation = True
        args.add_pe = False
    if "previsual" in args.checkpoint_path:
        args.use_format_v2 = True
        args.relation = False



    # load model
    flamingo, image_processor, tokenizer, vis_embed_size = create_model_and_transforms(
        args.vision_encoder_path,
        args.vision_encoder_pretrained,
        args.lm_path,
        args.lm_tokenizer_path,
        location_token_num=args.location_token_num,
        lora=args.lora,
        lora_r=16,
        use_sam=args.use_sam,
        add_visual_token=args.add_visual_token,
        use_format_v2=args.use_format_v2,
        add_box=args.add_box,
        add_pe=args.add_pe,
        add_relation=args.relation,
    )
    flamingo.use_format_v2 = args.use_format_v2
    if args.special:
        flamingo.special = True
    else:
        flamingo.special = False
    if args.legacy:
        flamingo.legacy = True
        print("use legacy evaluation")
    flamingo.step_num = int(args.checkpoint_path.split("/")[-1].split(".")[0].split("_")[-1])
    flamingo.expr_name = args.checkpoint_path.split("/")[-2]
    if args.rank == 0:
        print("legacy", True if hasattr(flamingo, "legacy") else False)
        print("step:", flamingo.step_num)
        print("expr:", flamingo.expr_name)
        print("use format v2:", flamingo.use_format_v2)
        print(args)
    checkpoint = torch.load(args.checkpoint_path, map_location="cpu")
    model_state_dict = {}
    for key in checkpoint["model_state_dict"].keys():
        model_state_dict[key.replace("module.", "")] = checkpoint["model_state_dict"][key]
    if "vision_encoder.logit_scale"in model_state_dict:
        # previous checkpoint has some unnecessary weights
        del model_state_dict["vision_encoder.logit_scale"]
        del model_state_dict["vision_encoder.visual.proj"]
        del model_state_dict["vision_encoder.visual.ln_post.weight"]
        del model_state_dict["vision_encoder.visual.ln_post.bias"]
    flamingo.load_state_dict(model_state_dict, strict=True)
    results = defaultdict(list)
    if args.eval_coco:
        print("Evaluating on COCO...")
        for shot in args.shots:
            scores = []
            for seed, trial in zip(args.trial_seeds, range(args.num_trials)):
                cider_score = evaluate_coco_flickr(
                    model=flamingo,
                    tokenizer=tokenizer,
                    image_processor=image_processor,
                    batch_size=args.batch_size,
                    image_dir_path=args.coco_image_dir_path,
                    annotations_json_path=args.coco_annotations_json_path,
                    device=args.device,
                    seed=seed,
                    vis_embed_size=vis_embed_size,
                    rank=args.rank,
                    world_size=args.world_size,
                    id=args.id,
                )
                print(f"Shots {shot} Trial {trial} CIDEr score: {cider_score}")
                scores.append(cider_score)
            print(f"Shots {shot} Mean CIDEr score: {np.mean(scores)}")
            results["coco"].append(
                {"shots": shot, "trials": scores, "mean": np.mean(scores)}
            )

    if args.eval_ok_vqa:
        print("Evaluating on OK-VQA...")
        for shot in args.shots:
            scores = []
            for seed, trial in zip(args.trial_seeds, range(args.num_trials)):
                ok_vqa_score = evaluate_vqa(
                    model=flamingo,
                    tokenizer=tokenizer,
                    image_processor=image_processor,
                    batch_size=args.batch_size,
                    image_dir_path=args.ok_vqa_image_dir_path,
                    questions_json_path=args.ok_vqa_questions_json_path,
                    annotations_json_path=args.ok_vqa_annotations_json_path,
                    vqa_dataset="ok_vqa",
                    vis_embed_size=vis_embed_size,
                    rank=args.rank,
                    world_size=args.world_size,
                    id=args.id,
                )
            results["ok_vqa"].append(
                {"shots": shot, "score": ok_vqa_score}
            )

    if args.eval_vqav2:
        print("Evaluating on VQAv2...")
        for shot in args.shots:
            scores = []
            for seed, trial in zip(args.trial_seeds, range(args.num_trials)):
                vqa_score = evaluate_vqa(
                    model=flamingo,
                    tokenizer=tokenizer,
                    image_processor=image_processor,
                    batch_size=args.batch_size,
                    image_dir_path=args.vqav2_image_dir_path,
                    questions_json_path=args.vqav2_questions_json_path,
                    annotations_json_path=args.vqav2_annotations_json_path,
                    vqa_dataset="vqa",
                    vis_embed_size=vis_embed_size,
                    rank=args.rank,
                    world_size=args.world_size,
                    id=args.id,
                )
            results["vqav2"].append(
                {"shots": shot, "score": vqa_score}
            )

    if args.eval_gqa:
        print("Evaluating on GQA...")
        for shot in args.shots:
            scores = []
            for seed, trial in zip(args.trial_seeds, range(args.num_trials)):
                vqa_score = evaluate_vqa(
                    model=flamingo,
                    tokenizer=tokenizer,
                    image_processor=image_processor,
                    batch_size=args.batch_size,
                    vqa_dataset="gqa",
                    vis_embed_size=vis_embed_size,
                    rank=args.rank,
                    world_size=args.world_size,
                    id=args.id,
                )
            results["gqa"].append(
                {"shots": shot, "score": vqa_score}
            )

    if args.eval_imagenet:
        print("Evaluating on ImageNet...")
        for shot in args.shots:
            scores = []
            for seed, trial in zip(args.trial_seeds, range(args.num_trials)):
                imagenet_score = evaluate_imagenet(
                    model=flamingo,
                    tokenizer=tokenizer,
                    image_processor=image_processor,
                    batch_size=args.batch_size,
                    num_samples=args.num_samples,
                    num_shots=shot,
                    device=args.device,
                    seed=seed,
                    imagenet_root=args.imagenet_root,
                )
                print(
                    f"Shots {shot} Trial {trial} " f"ImageNet score: {imagenet_score}"
                )
                scores.append(imagenet_score)
            print(f"Shots {shot} Mean ImageNet score: {np.mean(scores)}")
            results["imagenet"].append(
                {"shots": shot, "trials": scores, "mean": np.mean(scores)}
            )

    if args.eval_refcoco:
        print("Evaluating on RefCOCO...")
        refcoco_score = evaluate_refcoco(
            model=flamingo,
            tokenizer=tokenizer,
            image_processor=image_processor,
            batch_size=args.batch_size,
            device=args.device,
            tsvfile=args.refcoco_tsvfile,
            vis_embed_size=vis_embed_size,
            rank=args.rank,
            world_size=args.world_size,
            id=args.id,
        )
        results["refcoco"].append(
            {"score": refcoco_score}
        )

def prepare_batch_images(batch, image_processor):
    batch_images = None
    for b in batch:
        b_image = image_processor(b["image"]).unsqueeze(0).unsqueeze(1).unsqueeze(0)
        if batch_images is None:
            batch_images = b_image
        else:
            batch_images = torch.cat([batch_images, b_image], dim=0)
    return batch_images

def get_outputs(
    model,
    batch_images,
    attention_mask,
    max_generation_length,
    min_generation_length,
    num_beams,
    length_penalty,
    input_ids,
    image_start_index_list=None,
    image_nums=None,
    bad_words_ids=None,
):
    with torch.inference_mode() and torch.cuda.amp.autocast(dtype=torch.float16):
        outputs = model.generate(
            batch_images,
            input_ids,
            attention_mask=attention_mask,
            max_new_tokens=max_generation_length,
            min_length=min_generation_length,
            num_beams=num_beams,
            length_penalty=length_penalty,
            image_start_index_list=image_start_index_list,
            image_nums=image_nums,
            bad_words_ids=bad_words_ids,
        )

    outputs = outputs[:, len(input_ids[0]) :]
    return outputs


def evaluate_coco_flickr(
    model,
    tokenizer,
    image_processor,
    batch_size,
    image_dir_path,
    annotations_json_path,
    seed=42,
    max_generation_length=20,
    num_beams=1,
    length_penalty=-2.0,
    device=-1,
    is_flickr=False,
    vis_embed_size=None,
    rank=0,
    world_size=1,
    id=0,
):
    """Evaluate a model on COCO dataset.

    Args:
        model (nn.Module): model to evaluate
        tokenizer (transformers.PreTrainedTokenizer): tokenizer for the model
        image_processor : image processor for the model
        batch_size (int): batch size
        image_dir_path (str, optional): path to the directory containing the images.
        annotations_json_path (str, optional): path to the json file containing the annotations.
        seed (int, optional): seed for random number generator. Defaults to 42.
        max_generation_length (int, optional): maximum length of the generated caption. Defaults to 10.
        num_beams (int, optional): number of beams to use for beam search. Defaults to 3.
        length_penalty (float, optional): length penalty for beam search. Defaults to -2.0.
        num_samples (int, optional): number of samples to evaluate on. Defaults to 5000.
        query_set_size (int, optional): number of samples to use for query set. Defaults to 2048.
        num_shots (int, optional): number of in-context samples to use. Defaults to 8.
        device (int, optional): device to use. Defaults to -1.
        num_workers (int, optional): number of workers to use for dataloader. Defaults to 4.
        is_flickr (bool): defines if that data is COCO or Flickr. Defaults to False (COCO).

    Returns:
        float: CIDEr score

    """
    # eval_dataset = COCOFlickrDataset(
    #     image_dir_path=image_dir_path,
    #     annotations_path=annotations_json_path,
    #     is_flickr=is_flickr,
    # )
    coco_dataset = load_dataset("coco_caption")
    eval_dataset = coco_dataset["test"]


    model.eval().cuda()
    predictions = defaultdict()
    lang_encoder_name = model.lang_encoder.__class__.__name__.lower()
    # if "peft" in lang_encoder_name:
        # lang_encoder_name = model.lang_encoder.base_model.model.__class__.__name__.lower()
    try:
        media_token_id = tokenizer("<|#image#|>", add_special_tokens=False)["input_ids"][-1]
        endofmedia_token_id = tokenizer("<|#endofimage#|>", add_special_tokens=False)["input_ids"][-1]
        pad_token_id = tokenizer(tokenizer.pad_token, add_special_tokens=False)["input_ids"][-1]
        bos_token_id = tokenizer(tokenizer.bos_token, add_special_tokens=False)["input_ids"][-1]
    except:
        pass

    def get_prompt(sample):
        return f"<|#image#|>{tokenizer.pad_token*vis_embed_size}<|#endofimage#|>"

    tokenizer.padding_side = "left"
    cnt = 0
    if world_size > 1:
        torch.distributed.barrier()
    desc = "Running inference Flickr30" if is_flickr else "Running inference COCO"
    for ii, batch in enumerate(more_itertools.chunked(
        tqdm(eval_dataset, desc=desc, disable=(rank != 0)), batch_size
    )):
        if ii % world_size != rank:
            continue
        cnt += len(batch)
        batch_images = prepare_batch_images(
            batch=batch,
            image_processor=image_processor,
        ).cuda()
        batch_text = [get_prompt(s) for s in batch]
        encodings = tokenizer(
            batch_text,
            padding="longest",
            truncation=True,
            return_tensors="pt",
            max_length=2000,
        )
        input_ids = encodings["input_ids"].cuda()
        attention_mask = encodings["attention_mask"].cuda()
        skip_special_tokens = False
        if hasattr(model, "legacy") and model.legacy and "opt" in lang_encoder_name:
            if rank == 0:
                tqdm.write("use legacy model")
            skip_special_tokens = True
            for i in range(len(input_ids)):
                media_token_index = (input_ids[i] == media_token_id).nonzero()[0,0]
                endofmedia_token_index = (input_ids[i] == endofmedia_token_id).nonzero()[0,0]
                input_ids[i, media_token_index - 1] = media_token_id
                input_ids[i, media_token_index] = pad_token_id
                input_ids[i, endofmedia_token_index - 1] = endofmedia_token_id
                input_ids[i, endofmedia_token_index] = bos_token_id
        image_start_index_list = ((input_ids == media_token_id).nonzero(as_tuple=True)[-1] + 1).tolist()
        image_start_index_list = [[x] for x in image_start_index_list]
        image_nums = [1] * len(input_ids)
        if "llama" in lang_encoder_name:
            attention_mask[input_ids == 0] = 0
        outputs = get_outputs(
            model=model,
            batch_images=batch_images,
            attention_mask=attention_mask,
            max_generation_length=30,
            min_generation_length=8,
            num_beams=5,
            length_penalty=0,
            input_ids=input_ids,
            image_start_index_list=image_start_index_list,
            image_nums=image_nums,
        )
        new_predictions = [
            postprocess_captioning_generation(out).replace('"', "")
            for out in tokenizer.batch_decode(outputs, skip_special_tokens=True)
        ]
        # if rank == 0:
        #     tqdm.write(f"{batch_images.shape} {batch[0]} pred: {new_predictions[0]}")

        for i, sample in enumerate(batch):
            predictions[int(sample["image_id"])] = {
                "caption": new_predictions[i],
            }
    results_path = (
        f"flickrresults_{lang_encoder_name}_{rank}_{id}.json"
        if is_flickr
        else f"cocoresults_{lang_encoder_name}_{rank}_{id}.json"
    )
    with open(results_path, "w") as f:
        f.write(
            json.dumps(
                [
                    {"image_id": k, "caption": predictions[k]["caption"]}
                    for k in predictions
                ],
                indent=2,
            )
        )
    print("save to", results_path)
    del predictions
    time.sleep(10)
    if world_size > 1:
        torch.distributed.barrier()
    if rank == 0:
        print(f"evaluate on rank {rank}. world size is {world_size}")
        predictions = []
        for rank_i in range(world_size):
            part_results_path = (
                f"flickrresults_{lang_encoder_name}_{rank_i}_{id}.json"
                if is_flickr
                else f"cocoresults_{lang_encoder_name}_{rank_i}_{id}.json"
            )
            print("load", part_results_path)
            predictions.extend(json.load(open(part_results_path)))
            os.remove(part_results_path)
        print("num:", len(predictions))
        results_path = (
            f"flickrresults_{lang_encoder_name}.json"
            if is_flickr
            else f"cocoresults_{lang_encoder_name}.json"
        )
        json.dump(predictions, open(results_path, "w"), indent=2)

        metrics = compute_cider(
            result_path=results_path,
            annotations_path="/gpfs/u/home/LMCG/LMCGljnn/scratch/.cache/lavis/coco_gt/coco_karpathy_test_gt.json",
        )
        os.makedirs("eval_results", exist_ok=True)
        acc = metrics["CIDEr"]
        with open(os.path.join("eval_results", f"cococap_{model.expr_name}_{model.step_num}_{int(time.time())}_{acc}"), "w") as f:
            f.write(json.dumps(predictions, indent=2))

        # delete the temporary file
        os.remove(results_path)
    else:
        metrics = {}
        metrics["CIDEr"] = 0.0

    return metrics["CIDEr"]


def evaluate_vqa(
    model,
    tokenizer,
    image_processor,
    batch_size,
    image_dir_path=None,
    questions_json_path=None,
    annotations_json_path=None,
    vqa_dataset="vqa",
    vis_embed_size=None,
    rank=0,
    world_size=1,
    id=0,
):
    """
    Evaluate a model on VQA datasets. Currently supports VQA v2.0.

    Args:
        model (nn.Module): model to evaluate
        tokenizer (transformers.PreTrainedTokenizer): tokenizer for the model
        image_processor : image processor for the model
        batch_size (int): batch size
        image_dir_path (str): path to image directory
        questions_json_path (str): path to questions json file
        annotations_json_path (str): path to annotations json file
        seed (int, optional): random seed. Defaults to 42.
        max_generation_length (int, optional): max generation length. Defaults to 5.
        num_beams (int, optional): number of beams to use for beam search. Defaults to 3.
        length_penalty (float, optional): length penalty for beam search. Defaults to -2.0.
        num_samples (int, optional): number of samples to evaluate on. Defaults to 5000 samples.
        query_set_size (int, optional): size of the query set. Defaults to 2048.
        num_shots (int, optional): number of shots to use. Defaults to 8.
        device (int, optional): device to use. Defaults to -1 (cpu).
        num_workers (int, optional): number of workers to use. Defaults to 4.
        vqa_dataset (string): type of vqa dataset: currently supports vqa, ok_vqa. Defaults to vqa.
    Returns:
        float: accuracy score
    """
    if world_size > 1:
        torch.distributed.barrier()
    if vqa_dataset == "gqa":
        eval_dataset = GQADataset()
    else:
        eval_dataset = VQADataset(
            image_dir_path=image_dir_path,
            question_path=questions_json_path,
            annotations_path=annotations_json_path,
            vqa_dataset=vqa_dataset,
        )
    postprocessor = OKVQAPostProcess()
    try:
        media_token_id = tokenizer("<|#image#|>", add_special_tokens=False)["input_ids"][-1]
        endofmedia_token_id = tokenizer("<|#endofimage#|>", add_special_tokens=False)["input_ids"][-1]
        pad_token_id = tokenizer(tokenizer.pad_token, add_special_tokens=False)["input_ids"][-1]
        bos_token_id = tokenizer(tokenizer.bos_token, add_special_tokens=False)["input_ids"][-1]
    except:
        pass
    def get_prompt(sample):
        return f"<|#image#|>{tokenizer.pad_token*vis_embed_size}<|#endofimage#|>Question: {sample['question'].strip()} Short answer:"
        # return f"<|#image#|>{tokenizer.pad_token*vis_embed_size}<|#endofimage#|>"

    model.eval().cuda()
    lang_encoder_name = model.lang_encoder.__class__.__name__.lower()
    if "peft" in lang_encoder_name:
        lang_encoder_name = model.lang_encoder.base_model.model.__class__.__name__.lower()
    predictions = []
    tokenizer.padding_side = "left"
    if world_size > 1:
        torch.distributed.barrier()
    for ii, batch in enumerate(more_itertools.chunked(
        tqdm(eval_dataset, desc="Running inference", disable=(rank != 0)), batch_size
    )):
        if ii % world_size != rank:
            continue
        batch_images = prepare_batch_images(
            batch=batch,
            image_processor=image_processor,
        ).cuda()
        batch_text = [get_prompt(s) for s in batch]
        encodings = tokenizer(
            batch_text,
            return_tensors="pt",
            padding="longest",
            truncation=True,
            max_length=2000,
        )
        input_ids = encodings["input_ids"].cuda()
        attention_mask = encodings["attention_mask"].cuda()
        skip_special_tokens = True
        if hasattr(model, "legacy") and model.legacy and "opt" in lang_encoder_name:
            if rank == 0:
                tqdm.write("use legacy model")
            for i in range(len(input_ids)):
                media_token_index = (input_ids[i] == media_token_id).nonzero()[0,0]
                endofmedia_token_index = (input_ids[i] == endofmedia_token_id).nonzero()[0,0]
                input_ids[i, media_token_index - 1] = media_token_id
                input_ids[i, media_token_index] = pad_token_id
                input_ids[i, endofmedia_token_index - 1] = endofmedia_token_id
                input_ids[i, endofmedia_token_index] = bos_token_id
        image_start_index_list = ((input_ids == media_token_id).nonzero(as_tuple=True)[-1] + 1).tolist()
        image_start_index_list = [[x] for x in image_start_index_list]
        image_nums = [1] * len(input_ids)
        if "llama" in lang_encoder_name:
            attention_mask[input_ids == 0] = 0
        outputs = get_outputs(
            model=model,
            batch_images=batch_images,
            attention_mask=attention_mask,
            max_generation_length=10,
            min_generation_length=1,
            num_beams=5,
            length_penalty=0,
            input_ids=input_ids,
            image_start_index_list=image_start_index_list,
            image_nums=image_nums,
        )
        # postprocess begin
        new_predictions = [
            out.strip().lower().strip(string.punctuation+" ") for out in tokenizer.batch_decode(outputs, skip_special_tokens=skip_special_tokens)
        ]
        if vqa_dataset == "ok_vqa":
            new_predictions = postprocessor._lemmatize(new_predictions)
        if model.special:
            for i in range(len(new_predictions)):
                for answer, _ in Counter(batch[i]['answers']).most_common():
                    if answer in new_predictions[i]:
                        new_predictions[i] = answer
                        break
                    if "cant" in new_predictions[i] and "no" == answer:
                        new_predictions[i] = answer
                        break
                    if "can" in new_predictions[i] and "not" not in new_predictions[i] and "cant" not in new_predictions[i] and "yes" == answer:
                        new_predictions[i] = answer
                        break

        # if rank == 0:
        #     tqdm.write(f"{image_nums} {image_start_index_list}")
        #     for i in range(1):
        #         tqdm.write(f"ID: {batch[i]['question_id']} | gt QA: {batch[i]['question']} {Counter(batch[i]['answers']).most_common()}")
        #         tqdm.write("prompt: " + tokenizer.decode(input_ids[i]))
        #         tqdm.write("model output: " + new_predictions[i])

        predictions.extend(
            [
                {"answer": p, "question_id": sample["question_id"], "_question": sample["question"], "answers": sample["answers"]}
                for p, sample in zip(new_predictions, batch)
            ]
        )
    with open(f"{vqa_dataset}_{lang_encoder_name}_results_part{rank}_{id}.json", "w") as f:
        f.write(json.dumps(predictions))
    print("save to", f"{vqa_dataset}_{lang_encoder_name}_results_part{rank}_{id}.json")

    time.sleep(10)
    if world_size > 1:
        torch.distributed.barrier()
    if rank == 0:
        print(f"evaluate on rank {rank}. world size is {world_size}")
        predictions = []
        for rank_i in range(world_size):
            print("load", f"{vqa_dataset}_{lang_encoder_name}_results_part{rank_i}_{id}.json")
            predictions.extend(json.load(open(f"{vqa_dataset}_{lang_encoder_name}_results_part{rank_i}_{id}.json")))
            os.remove(f"{vqa_dataset}_{lang_encoder_name}_results_part{rank_i}_{id}.json")
        print("num:", len(predictions))
        # save the predictions to a temporary file
        random_uuid = str(uuid.uuid4())
        with open(f"{vqa_dataset}results_{random_uuid}.json", "w") as f:
            f.write(json.dumps(predictions, indent=4))

        if vqa_dataset == "gqa":
            acc = compute_gqa_accuracy(predictions)
        else:
            acc = compute_vqa_accuracy(
                f"{vqa_dataset}results_{random_uuid}.json",
                questions_json_path,
                annotations_json_path,
                vqa_dataset=vqa_dataset,
            )
        print(vqa_dataset, "score:", acc, "| save to", f"{vqa_dataset}results_{random_uuid}.json")
        os.makedirs("eval_results", exist_ok=True)
        with open(os.path.join("eval_results", f"{vqa_dataset}_{model.expr_name}_{model.step_num}_{int(time.time())}_{acc}"), "w") as f:
            f.write(json.dumps(predictions, indent=2))

        # delete the temporary file
        os.remove(f"{vqa_dataset}results_{random_uuid}.json")
    else:
        time.sleep(5)
        acc = 0.0
    if world_size > 1:
        torch.distributed.barrier()
    return acc


def evaluate_refcoco(
    model,
    tokenizer,
    image_processor,
    batch_size,
    tsvfile,
    max_generation_length=20,
    num_beams=3,
    length_penalty=-2.0,
    device=-1,
    vis_embed_size=None,
    rank=0,
    world_size=1,
    id=0,
):
    model.eval().cuda()
    loc_token_ids = []
    for i in range(1000):
        loc_token_ids.append(int(tokenizer(f"<loc_{i}>", add_special_tokens=False)["input_ids"][-1]))
    media_token_id = tokenizer("<|#image#|>", add_special_tokens=False)["input_ids"][-1]
    total = 0
    correct = 0
    ious = []
    if "refcocog" in tsvfile:
        dataset_name = "refcocog"
    elif "refcocoplus" in tsvfile:
        dataset_name = "refcocoplus"
    else:
        dataset_name = "refcoco"
    with open(tsvfile, "r") as f:
        lines = f.readlines()
        pbar = tqdm(lines, disable=(rank != 0))
        for ii, line in enumerate(pbar):
            if ii % world_size != rank:
                continue
            total += 1
            line = line.rstrip()
            uniq_id, image_id, text, region_coord, image = line.split("\t")

            # image = Image.open(BytesIO(base64.urlsafe_b64decode(image))).convert("RGB")
            # image = Image.open("/gpfs/u/home/LMCG/LMCGljnn/scratch/code/multimodal/temp/cat.png").convert("RGB")
            # image2 = Image.open("yolo.png").convert("RGB")
            # image1 = image1.resize((224, 224))
            # image2 = image2.resize((224, 224))
            # images = [image1, image2]

            # gt_box = np.array(list(map(float, region_coord.split(","))))
            # width = image.width
            # height = image.height
            # gt_box /= np.array([width, height, width, height])
            # batch_images = [image_processor(image).unsqueeze(0).unsqueeze(1).unsqueeze(0) for image in images]
            # batch_images = torch.cat(batch_images, dim=0)
            # image = Image.open("yolo_test.png").convert("RGB")
            image = Image.open("example.png").convert("RGB")
            image = image.resize((224, 224))
            batch_images = image_processor(image).unsqueeze(0).unsqueeze(1).unsqueeze(0)
            # prompt = [f"<|#image#|>{tokenizer.pad_token*vis_embed_size}<|#endofimage#|>{text.rstrip('.')}<|#visual#|>"]
            prompt = [f"{tokenizer.bos_token}<|#image#|>{tokenizer.pad_token*vis_embed_size}<|#endofimage#|><|#object#|><|#previsual#|><|#prebox#|><|#endofattr#|>man<|#endofobject#|><|#visual#|><|#box#|><|#endofattr#|> is sitting on<|#object#|><|#previsual#|>"]
            # prompt = [f"{tokenizer.bos_token}<|#image#|>{tokenizer.pad_token*vis_embed_size}<|#endofimage#|><|#object#|><|#previsual#|>man<|#endofobject#|><|#visual#|><|#box#|><|#endofattr#|> is sitting on<|#object#|><|#previsual#|>"]
            # prompt = [f"<|#image#|>{tokenizer.pad_token*vis_embed_size}<|#endofimage#|>"]
            # prompt = [f"<|#image#|>{tokenizer.pad_token*vis_embed_size}<|#endofimage#|>a man<|#visual#|> is doing a trick on a skateboard<|#visual#|>"]


            encodings = tokenizer(
                prompt,
                padding="longest",
                truncation=True,
                return_tensors="pt",
                max_length=2000,
            )
            input_ids = encodings["input_ids"]
            attention_mask = encodings["attention_mask"]
            image_start_index_list = ((input_ids == media_token_id).nonzero(as_tuple=True)[-1] + 1).tolist()
            image_start_index_list = [image_start_index_list]
            image_nums = [1]
            vision_x = batch_images.cuda()
            lang_x = input_ids.cuda()
            attention_mask = attention_mask.cuda()
            print(image_start_index_list, image_nums)

            model.debug_id = 0
            # outputs = get_outputs(
            #     model=model,
            #     batch_images=vision_x,
            #     attention_mask=attention_mask,
            #     max_generation_length=20,
            #     min_generation_length=8,
            #     num_beams=5,
            #     length_penalty=0,
            #     input_ids=lang_x,
            #     image_start_index_list=image_start_index_list,
            #     image_nums=image_nums,
            # )
            # print(tokenizer.decode(outputs[0]))
            # exit()

            prebox = [93, 20, 155, 172] # man
            # prebox = [32, 82, 89, 213] # dog
            # prebox = [34, 49, 166, 164] # bike
            with torch.inference_mode() and torch.cuda.amp.autocast(dtype=torch.float16):
                outputs = model(
                    vision_x=vision_x,
                    lang_x=lang_x,
                    attention_mask=attention_mask,
                    labels=None,
                    image_nums=image_nums,
                    image_start_index_list=image_start_index_list,
                    added_bbox_list=[torch.tensor(prebox).cuda().unsqueeze(0) / 224],
                    add_box=True,
                    debug_mode=True,
                )
            
            boxes = outputs["boxes"]
            scores = outputs["scores"]
            box = boxes[scores.argmax()]
            open_cv_image = np.array(image)
            # Convert RGB to BGR 
            open_cv_image = open_cv_image[:, :, ::-1].copy() 
            open_cv_image = cv2.rectangle(open_cv_image, box[:2].astype(int), box[2:].astype(int), (255, 0, 0), 2)
            open_cv_image = cv2.rectangle(open_cv_image, prebox[:2], prebox[2:], (0, 0, 255), 2)
            cv2.imwrite(f"output2.jpg", open_cv_image)
            print(box)
            print(prebox)
            exit()

            # force_words = ["man", "table"]
            # force_words_ids = tokenizer(force_words, add_special_tokens=False).input_ids


            # sequences, hidden_states_for_each_step = get_outputs(
            #     model=model,
            #     batch_images=vision_x,
            #     attention_mask=attention_mask,
            #     max_generation_length=20,
            #     min_generation_length=8,
            #     num_beams=5,
            #     length_penalty=0,
            #     input_ids=lang_x,
            #     image_start_index_list=image_start_index_list,
            #     image_nums=image_nums,
            #     force_words_ids=force_words_ids,
            # )
            # sequence = sequences[0]
            # print(tokenizer.decode(sequence))
            # for i, token in enumerate(sequence):
            #     if token == model.visual_token_id:
            #         print(tokenizer.decode(sequence[:i+1]))
            #         if hasattr(model, "debug_id"):
            #             model.debug_id += 1
            #         else:
            #             model.debug_id = 0
            #         this_lang_x = torch.hstack([lang_x[0], sequence[:i+1]]).unsqueeze(0)
            #         this_attention_mask = torch.ones_like(this_lang_x).cuda()
            #         with torch.inference_mode() and torch.cuda.amp.autocast(dtype=torch.float16) and torch.no_grad():
            #             _ = model(
            #                 vision_x=vision_x,
            #                 lang_x=this_lang_x,
            #                 attention_mask=this_attention_mask,
            #                 labels=None,
            #                 image_nums=image_nums,
            #                 image_start_index_list=image_start_index_list,
            #                 added_bbox_list=None,
            #             )
            # exit()

    with open(f"{dataset_name}_results_part{rank}_{id}.json", "w") as f:
        f.write(json.dumps([total, correct]))
    if world_size > 1:
        torch.distributed.barrier()
    if rank == 0:
        total = 0
        correct = 0
        print(f"evaluate on rank {rank}. world size is {world_size}")
        for rank_i in range(world_size):
            [total_part, correct_part] = json.load(open(f"{dataset_name}_results_part{rank_i}_{id}.json"))
            os.remove(f"{dataset_name}_results_part{rank_i}_{id}.json")
            total += total_part
            correct += correct_part
        score = correct / total
        print("score:", score)
        with open(os.path.join("eval_results", f"{dataset_name}_{model.expr_name}_{model.step_num}_{int(time.time())}_{score}"), "w") as f:
            pass
    else:
        score = 0.0
    if world_size > 1:
        torch.distributed.barrier()
    return score


if __name__ == "__main__":
    main()