internvl_g/eval/evaluate_caption.py

import argparse
import itertools
import json
import os
import random
import time
from functools import partial

import torch
import torchvision.transforms as T
from internvl.model.internvl_stage2 import InternVLConfig, InternVLModel
from PIL import Image
from pycocoevalcap.eval import COCOEvalCap
from pycocotools.coco import COCO
from torchvision.transforms.functional import InterpolationMode
from tqdm import tqdm
from transformers import LlamaTokenizer

ds_collections = {
    'flickr30k': {
        'root': 'data/flickr30k/',
        'annotation': 'data/flickr30k/flickr30k_test_karpathy.json',
    },
    'coco': {
        'root': 'data/coco/',
        'annotation': ['data/coco/annotations/coco_karpathy_test.json',
                       'data/coco/annotations/coco_karpathy_test_gt.json'],
    },
    'nocaps': {
        'root': 'data/nocaps/images',
        'annotation': 'data/nocaps/nocaps_val_4500_captions.json',
    },
}


class CaptionDataset(torch.utils.data.Dataset):

    def __init__(self, name, root, annotation, prompt, input_size=224):
        if name == 'coco':
            self.images = json.load(open(annotation))
        else:
            self.images = json.load(open(annotation))['images']
        self.name = name
        self.prompt = prompt
        self.root = root
        self.transform = T.Compose([
            T.Lambda(lambda img: img.convert('RGB') if img.mode != 'RGB' else img),
            T.Resize((input_size, input_size), interpolation=InterpolationMode.BICUBIC),
            T.ToTensor(),
            T.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225))
        ])

    def __len__(self):
        return len(self.images)

    def __getitem__(self, idx):
        if self.name == 'coco':
            filename = self.images[idx]['image']
            image_id = int(filename.split('_')[-1].replace('.jpg', ''))
            image_path = os.path.join(self.root, filename)
        else:
            image_id = self.images[idx]['id']
            if 'file_name' in self.images[idx]:
                image_path = os.path.join(self.root, self.images[idx]['file_name'])
            else:
                image_path = os.path.join(self.root, self.images[idx]['image'])
        image = Image.open(image_path)
        pixel_values = self.transform(image).unsqueeze(0)
        return {
            'image_id': image_id,
            'input_text': self.prompt,
            'pixel_values': pixel_values
        }


def collate_fn(inputs, tokenizer):
    pixel_values = torch.cat([_['pixel_values'] for _ in inputs], dim=0)
    image_ids = [_['image_id'] for _ in inputs]
    input_texts = [_['input_text'] for _ in inputs]
    input_tokens = tokenizer(input_texts, return_tensors='pt')

    return pixel_values, image_ids, input_tokens.input_ids, input_tokens.attention_mask


class InferenceSampler(torch.utils.data.sampler.Sampler):

    def __init__(self, size):
        self._size = int(size)
        assert size > 0
        self._rank = torch.distributed.get_rank()
        self._world_size = torch.distributed.get_world_size()
        self._local_indices = self._get_local_indices(size, self._world_size, self._rank)

    @staticmethod
    def _get_local_indices(total_size, world_size, rank):
        shard_size = total_size // world_size
        left = total_size % world_size
        shard_sizes = [shard_size + int(r < left) for r in range(world_size)]

        begin = sum(shard_sizes[:rank])
        end = min(sum(shard_sizes[:rank + 1]), total_size)
        return range(begin, end)

    def __iter__(self):
        yield from self._local_indices

    def __len__(self):
        return len(self._local_indices)


def evaluate_qllama_model():
    prompts = ['English caption:']
    print('prompts:', prompts)

    config = InternVLConfig.from_pretrained(args.checkpoint)
    model = InternVLModel.from_pretrained(args.checkpoint, config=config).eval()
    model = model.to(torch.float16).cuda()
    tokenizer = LlamaTokenizer.from_pretrained(args.checkpoint)
    tokenizer.add_eos_token = False

    random.seed(args.seed)
    summaries = []
    for prompt in prompts:
        for ds_name in args.datasets:
            annotation = ds_collections[ds_name]['annotation']
            if type(annotation) == list:
                annotation = annotation[0]
            if model.config.force_image_size is not None:
                image_size = model.config.force_image_size
            else:
                image_size = model.config.vision_config.image_size
            dataset = CaptionDataset(
                name=ds_name,
                root=ds_collections[ds_name]['root'],
                annotation=annotation,
                prompt=prompt,
                input_size=image_size,
            )
            dataloader = torch.utils.data.DataLoader(
                dataset=dataset,
                sampler=InferenceSampler(len(dataset)),
                batch_size=args.batch_size,
                num_workers=args.num_workers,
                pin_memory=True,
                drop_last=False,
                collate_fn=partial(collate_fn, tokenizer=tokenizer),
            )

            image_ids, captions = [], []
            for _, (pixel_values, ids, input_ids, attention_mask) in tqdm(enumerate(dataloader)):
                pred = model.generate(
                    pixel_values=pixel_values.cuda().to(torch.float16),
                    input_ids=input_ids.cuda(),
                    attention_mask=attention_mask.cuda(),
                    do_sample=False,
                    num_beams=args.num_beams,
                    max_new_tokens=30,
                    min_new_tokens=8,
                    use_cache=True
                )
                image_ids.extend(ids)
                caption = [tokenizer.decode(_.cpu(), skip_special_tokens=True).strip() for _ in pred]
                captions.extend(caption)
                print(caption)

            torch.distributed.barrier()

            world_size = torch.distributed.get_world_size()
            merged_ids = [None for _ in range(world_size)]
            merged_captions = [None for _ in range(world_size)]
            torch.distributed.all_gather_object(merged_ids, image_ids)
            torch.distributed.all_gather_object(merged_captions, captions)

            merged_ids = [_ for _ in itertools.chain.from_iterable(merged_ids)]
            merged_captions = [_ for _ in itertools.chain.from_iterable(merged_captions)]
            average_length = sum(len(x.split()) for x in merged_captions) / len(merged_captions)
            print(f'Average length: {average_length}')

            if torch.distributed.get_rank() == 0:
                print(f'Evaluating {ds_name} ...')

                results = []
                for image_id, caption in zip(merged_ids, merged_captions):
                    results.append({
                        'image_id': int(image_id),
                        'caption': caption,
                    })
                time_prefix = time.strftime('%y%m%d%H%M%S', time.localtime())
                results_file = f'{ds_name}_{time_prefix}.json'
                results_file = os.path.join(args.out_dir, results_file)
                json.dump(results, open(results_file, 'w'))

                annotation = ds_collections[ds_name]['annotation']
                if type(annotation) == list:
                    annotation = annotation[-1]
                coco = COCO(annotation)
                coco_result = coco.loadRes(results_file)
                coco_eval = COCOEvalCap(coco, coco_result)
                coco_eval.evaluate()

                summary = coco_eval.eval.items()
                print([ds_name, prompt, average_length, summary])
                summaries.append([ds_name, prompt, average_length, summary])

            torch.distributed.barrier()

    for summary in summaries:
        print(summary)


if __name__ == '__main__':

    parser = argparse.ArgumentParser()
    parser.add_argument('--checkpoint', type=str, default='')
    parser.add_argument('--datasets', type=str, default='coco,flickr30k,nocaps')
    parser.add_argument('--batch-size', type=int, default=1)
    parser.add_argument('--num-workers', type=int, default=1)
    parser.add_argument('--num-beams', type=int, default=5)
    parser.add_argument('--out-dir', type=str, default='results')
    parser.add_argument('--seed', type=int, default=0)
    args = parser.parse_args()

    os.makedirs(args.out_dir, exist_ok=True)

    args.datasets = args.datasets.split(',')
    print('datasets:', args.datasets)
    assert args.batch_size == 1, 'Only batch size 1 is supported'

    torch.distributed.init_process_group(
        backend='nccl',
        world_size=int(os.getenv('WORLD_SIZE', '1')),
        rank=int(os.getenv('RANK', '0')),
    )

    torch.cuda.set_device(int(os.getenv('LOCAL_RANK', 0)))

    evaluate_qllama_model()