dataset/.ipynb_checkpoints/caption-checkpoint.py

from torch.utils.data import DataLoader, Dataset, Sampler
from pathlib import Path
from collections import defaultdict
import json
import random
from multiprocessing import Pool
import h5py
import pickle
import math
from tqdm import tqdm
import torch
import numpy as np
from copy import deepcopy
from PIL import Image

from torch.utils.data.distributed import DistributedSampler

# from transformers import T5TokenizerFast, BartTokenizer
# from tokenization import VLT5TokenizerFast

# from vis_encoder import _transform
# from vqa_raw_data import augmentation_transform

from dataset.randaugment import RandomAugment

import torch
from torch import nn
from torchvision import transforms

import os 
import re 
# project_dir = Path(__file__).resolve().parent.parent  # VLT5
# workspace_dir = project_dir.parent
# dataset_dir = workspace_dir.joinpath('datasets/').resolve()
# coco_dir = dataset_dir.joinpath('COCO')
# vg_dir = dataset_dir.joinpath('VG')
# coco_img_dir = coco_dir.joinpath('images/')
# coco_feature_dir = coco_dir.joinpath('features')


class COCOCaptionFineTuneDataset(Dataset):
    def __init__(self, split='karpathy_train', raw_dataset=None, rank=-1, topk=-1, verbose=True, args=None, mode='train', 
        data_dir='/data/mshukor/data'):
        super().__init__()

        self.raw_dataset = raw_dataset
        self.topk = topk
        self.verbose = verbose
        self.args = args

        self.args.BUTD100 = False

        self.mode = mode

        dataset_dir = Path(data_dir)
        coco_dir = dataset_dir.joinpath('COCO')
        vg_dir = dataset_dir.joinpath('VG')
        coco_img_dir = coco_dir.joinpath('images/')
        coco_feature_dir = coco_dir.joinpath('features')


        # Loading datasets to data
        self.source = split
        if self.verbose:
            print('Data source: ', self.source)


        normalize = transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711))
        

        self.train_transform = transforms.Compose([                        
                transforms.RandomResizedCrop(args.image_size,scale=(0.5, 1.0), interpolation=Image.BICUBIC),
                transforms.RandomHorizontalFlip(),
                RandomAugment(2,7,isPIL=True,augs=['Identity','AutoContrast','Equalize','Brightness','Sharpness',
                                                  'ShearX', 'ShearY', 'TranslateX', 'TranslateY', 'Rotate']),     
                transforms.ToTensor(),
                normalize,
            ])  
        self.test_transform = transforms.Compose([
            transforms.Resize((args.image_size,args.image_size),interpolation=Image.BICUBIC),
            transforms.ToTensor(),
            normalize,
            ])  

        # if self.args.tokenizer is None:
        #     self.args.tokenizer = self.args.backbone

        # if 't5' in self.args.tokenizer:
        #     if self.args.use_vision:
        #         self.tokenizer = VLT5TokenizerFast.from_pretrained(
        #             args.backbone,
        #             # max_length=self.args.max_text_length,
        #             do_lower_case=self.args.do_lower_case)
        #     else:
        #         self.tokenizer = T5TokenizerFast.from_pretrained(
        #             args.backbone,
        #             # max_length=self.args.max_text_length,
        #             do_lower_case=self.args.do_lower_case)
        # elif 'bart' in self.args.tokenizer:
        #     self.tokenizer = BartTokenizer.from_pretrained(
        #         args.backbone,
        #         # max_length=self.args.max_text_length,
        #         do_lower_case=self.args.do_lower_case)

        #     additional_special_tokens = [f'<extra_id_{i}>' for i in range(100-1, -1, -1)] + \
        #             [f'<vis_extra_id_{i}>' for i in range(100-1, -1, -1)]
        #     special_tokens_dict = {'additional_special_tokens': additional_special_tokens}
        #     num_added_toks = self.tokenizer.add_special_tokens(special_tokens_dict)

        # if self.args.oscar_tags:
        #     # Load VG Classes
        #     vg_classes = []
        #     with open(vg_dir.joinpath('objects_vocab.txt')) as f:
        #         for obj in f.readlines():
        #             vg_classes.append(obj.split(',')[0].lower().strip())
        #     self.vg_classes = vg_classes

        data_info_path = dataset_dir.joinpath('COCO/dataset_coco.json')
        with open(data_info_path) as f:
            karpathy_data = json.load(f)

        split_rename = {
            'train': 'train',
            'restval': 'train',
            'val': 'val',
            'test': 'test'
        }

        n_images = 0

        data = []
        for datum in karpathy_data['images']:
            re_split = split_rename[datum['split']]
            if re_split != self.source.split('_')[-1]:
                continue

            if re_split == 'train':
                for d in datum['sentences']:
                    # if self.args.BUTD100:
                    #     img_id = str(int(datum['filename'].split('.')[0].split('_')[-1]))
                    # else:
                    img_id = datum['filename'].split('.')[0]
                    new_datum = {
                        'img_id': img_id,
                        'sent': d['raw'].strip(),
                        'targets': [d['raw'].strip() for d in datum['sentences']],
                        'is_train': True,
                    }
                    data.append(new_datum)
            else:
                # if self.args.BUTD100:
                #     img_id = str(
                #         int(datum['filename'].split('.')[0].split('_')[-1]))
                # else:
                img_id = datum['filename'].split('.')[0]
                new_datum = {
                    'img_id': img_id,
                    # 'sent': d['raw'],
                    'targets': [d['raw'].strip() for d in datum['sentences']],
                    'is_train': False,
                }
                data.append(new_datum)

            n_images += 1

        if self.verbose:
            print(f"{self.source} has {n_images} images")
            print(f"Loaded {len(data)} data from", split)

        # self.n_gpus = torch.cuda.device_count()

        # self.rank = rank
        
        if isinstance(self.topk, float) and (0 < self.topk <= 1):
            used_samples = int(self.topk * len(data))
            data = random.sample(data, used_samples)
            if self.verbose:
                print(f"Use only {len(data)} data")

        elif self.topk > 0:
            data = data[:int(self.topk)]
            if self.verbose:
                print(f"Use only {len(data)} data")

        self.data = data

        if self.verbose:
            print("# all sentences:", len(self.data))

        # self.n_boxes = args.n_boxes

        self.image_size = self.args.image_size

        if mode == "train" and self.args.use_data_augmentation:
            self.transform = self.train_transform
        else:
            self.transform = self.test_transform

        self.source_to_h5 = {}

        # if self.args.max_n_boxes == 36:
        self.source_to_h5.update({
            'train2014': coco_img_dir.joinpath(f'train2014'),
            'val2014': coco_img_dir.joinpath(f'val2014'),
        })


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

    def __getitem__(self, idx):

        out_dict = {}
        out_dict['args'] = self.args

        datum = self.data[idx]

        ###### Image ######
        # if self.args.use_vision:
        img_id = datum['img_id']
        out_dict['img_id'] = img_id


        if self.args.BUTD100:
            source = self.source
        else:
            if 'train' in img_id:
                source = 'train2014'
            elif 'val' in img_id:
                source = 'val2014'

        path = self.source_to_h5[source].joinpath(f"{img_id}.jpg")
    
        image = Image.open(path).convert('RGB')   


        out_dict["image"] = self.transform(image)

        # out_dict['n_boxes'] = self.args.n_boxes


        ###### Text #####
        # if self.args.no_prefix:
        #     input_text = ''
        #     input_ids = []

        # else:
        #     if self.args.prefix is None:
        #         prefix = f'{self.args.prompt}'
        #     elif self.args.prefix == 'span':
        #         prefix = "span prediction:"
        #     elif self.args.prefix == 'denoise':
        #         prefix = "denoise text: <mask>"
        #     elif self.args.prefix == 'mask':
        #         if 'bart' in self.args.tokenizer:
        #             prefix = "<mask>"

        #     input_tokens = [prefix]

        #     # if self.args.oscar_tags:
        #     #     prefix = f'describe image with tags:'
        #     #     input_tokens = [prefix]
        #     #     obj_ids = f[f'{img_id}/obj_id'][()]
        #     #     for obj_id in obj_ids:
        #     #         obj = self.vg_classes[obj_id]
        #     #         if obj not in input_tokens:
        #     #             input_tokens.append(obj)
        #     input_text = ' '.join(input_tokens)

        #     # if 't5' in self.args.tokenizer:
        #     #     input_ids = self.tokenizer.encode(
        #     #         input_text,
        #     #         max_length=self.args.max_text_length, truncation=True)
        #     # elif 'bart' in self.args.tokenizer:
        #     #     input_ids = self.tokenizer.encode(
        #     #         input_text,
        #     #         max_length=self.args.max_text_length, truncation=True)
        #     # else:
        #     #     input_ids = self.tokenizer.convert_tokens_to_ids(
        #     #         self.tokenizer.tokenize(input_text)[:self.args.max_text_length - 1] + ['[SEP]'])

        # out_dict['input_text'] = input_text

        # out_dict['input_ids'] = torch.LongTensor(input_ids)
        # out_dict['input_length'] = len(input_ids)
        if datum['is_train']:
            sent = datum['sent'].strip()
            # if 't5' in self.args.tokenizer:
            #     target_ids = self.tokenizer.encode(sent, max_length=self.args.gen_max_length, truncation=True)
            # elif 'bart' in self.args.tokenizer:
            #     target_ids = self.tokenizer.encode(sent, max_length=self.args.gen_max_length, truncation=True)

            # assert len(target_ids) <= self.args.gen_max_length, len(target_ids)
            out_dict['sent'] = sent
            # out_dict['target_ids'] = torch.LongTensor(target_ids)
            # out_dict['target_length'] = len(target_ids)

        if 'targets' in datum:
            out_dict['targets'] = datum['targets']


        return out_dict

    def collate_fn(self, batch):
        batch_entry = {}

        B = len(batch)

        # S_W_L = max(entry['input_length'] for entry in batch)
        # input_ids = torch.ones(B, S_W_L, dtype=torch.long) * self.tokenizer.pad_token_id

        # if self.args.no_prefix:
        #     assert input_ids.size() == (B, 0)

        # if self.args.use_vision:
        #     pass

        if 'target_ids' in batch[0]:
            T_W_L = max(entry['target_length'] for entry in batch)
            target_ids = torch.ones(B, T_W_L, dtype=torch.long) * self.tokenizer.pad_token_id

        # sentences = []

        targets = []
        img_ids = []
        img_paths = []
        input_text = []
        images = []
        sents = []

        for i, entry in enumerate(batch):
            # input_ids[i, :entry['input_length']] = entry['input_ids']

            # if self.args.use_vision:
            # n_boxes = entry['n_boxes']
            images.append(entry['image'])
            img_ids.append(entry['img_id'])
            # img_paths.append(entry['img_path'])

            if 'target_ids' in entry:
                target_ids[i, :entry['target_length']] = entry['target_ids']

            # if 'input_text' in entry:
            #     input_text.append(entry['input_text'])

            # sentences.append(entry['sent'])

            if 'targets' in entry:
                targets.append(entry['targets'])
            if 'sent' in entry:
                sents.append(entry['sent'])
        # batch_entry['input_ids'] = input_ids
        # if 'target_ids' in batch[0]:
        #     word_mask = target_ids != self.tokenizer.pad_token_id
        #     target_ids[~word_mask] = -100
        #     batch_entry['target_ids'] = target_ids

        # if self.args.use_vision:
        batch_entry['images'] = torch.stack(images)
        batch_entry['img_id'] = img_ids
        batch_entry['img_paths'] = img_paths
        if 'sent' in entry:
            batch_entry['sent'] = sents

        # batch_entry['sent'] = sentences

        # batch_entry['input_text'] = input_text

        batch_entry['targets'] = targets

        batch_entry['task'] = 'caption'

        return batch_entry


def pre_caption(caption,max_words):
    caption = re.sub(
        r"([,.'!?\"()*#:;~])",
        '',
        caption.lower(),
    ).replace('-', ' ').replace('/', ' ').replace('<person>', 'person')

    caption = re.sub(
        r"\s{2,}",
        ' ',
        caption,
    )
    caption = caption.rstrip('\n') 
    caption = caption.strip(' ')

    #truncate caption
    caption_words = caption.split(' ')
    if len(caption_words)>max_words:
        caption = ' '.join(caption_words[:max_words])
            
    return caption

class CCDataset(Dataset):
    def __init__(self, split='CC', raw_dataset=None, rank=-1, topk=-1, verbose=True, args=None, mode='train', 
        data_dir='/data/mshukor/data', config_dir='/data/mshukor/data/cc3m.json', max_words=30):
        super().__init__()

        self.raw_dataset = raw_dataset
        self.topk = topk
        self.verbose = verbose
        self.args = args


        self.mode = mode


        data = []
        ann_files = [config_dir]
        ann_file = []
        for p in ann_files:
            ann_file.append(os.path.join(args.data_json_dir, p))
            
        for f in ann_file:
            tmp =  json.load(open(f,'r'))
            data += tmp
            print('size of', f, len(tmp))
        print(len(data))
        self.max_words = max_words
        for e in data:
            e['image'] = os.path.join(data_dir, ('/').join(e['image'].split('/')[4:]))

        # Loading datasets to data
        self.source = split
        if self.verbose:
            print('Data source: ', self.source)


        normalize = transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711))
        

        self.train_transform = transforms.Compose([                        
                transforms.RandomResizedCrop(args.image_size,scale=(0.5, 1.0), interpolation=Image.BICUBIC),
                transforms.RandomHorizontalFlip(),
                RandomAugment(2,7,isPIL=True,augs=['Identity','AutoContrast','Equalize','Brightness','Sharpness',
                                                  'ShearX', 'ShearY', 'TranslateX', 'TranslateY', 'Rotate']),     
                transforms.ToTensor(),
                normalize,
            ])  
        self.test_transform = transforms.Compose([
            transforms.Resize((args.image_size,args.image_size),interpolation=Image.BICUBIC),
            transforms.ToTensor(),
            normalize,
            ])  

        
        if isinstance(self.topk, float) and (0 < self.topk <= 1):
            used_samples = int(self.topk * len(data))
            data = random.sample(data, used_samples)
            if self.verbose:
                print(f"Use only {len(data)} data")

        elif self.topk > 0:
            data = data[:int(self.topk)]
            if self.verbose:
                print(f"Use only {len(data)} data")

        self.data = data

        if self.verbose:
            print("# all sentences:", len(self.data))


        self.image_size = self.args.image_size

        if mode == "train" and self.args.use_data_augmentation:
            self.transform = self.train_transform
        else:
            self.transform = self.test_transform



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

    def __getitem__(self, idx):

        out_dict = {}
        out_dict['args'] = self.args

        datum = self.data[idx]
        
        if type(datum['caption']) == list:
            caption = pre_caption(random.choice(datum['caption']), self.max_words)
        else:
            caption = pre_caption(datum['caption'], self.max_words)
      
        ###### Image ######
        image = Image.open(datum['image']).convert('RGB')   

        img_id = datum['image'].split('/')[-1].split('.')[0]
        out_dict['img_id'] = img_id


        out_dict["image"] = self.transform(image)

        out_dict['sent'] = caption

        out_dict['targets'] = caption


        return out_dict

    def collate_fn(self, batch):
        batch_entry = {}

        B = len(batch)

 
        if 'target_ids' in batch[0]:
            T_W_L = max(entry['target_length'] for entry in batch)
            target_ids = torch.ones(B, T_W_L, dtype=torch.long) * self.tokenizer.pad_token_id


        targets = []
        img_ids = []
        img_paths = []
        input_text = []
        images = []
        sents = []

        for i, entry in enumerate(batch):

            images.append(entry['image'])
            img_ids.append(entry['img_id'])

            if 'target_ids' in entry:
                target_ids[i, :entry['target_length']] = entry['target_ids']

            if 'targets' in entry:
                targets.append(entry['targets'])
            if 'sent' in entry:
                sents.append(entry['sent'])


        # if self.args.use_vision:
        batch_entry['images'] = torch.stack(images)
        batch_entry['img_id'] = img_ids
        batch_entry['img_paths'] = img_paths
        if 'sent' in entry:
            batch_entry['sent'] = sents

        batch_entry['targets'] = targets

        batch_entry['task'] = 'caption'

        return batch_entry

def get_loader(args, split='train', mode='train',
               batch_size=32, workers=4, distributed=False, gpu=0,
               topk=-1, data_dir='/data/mshukor/data', local_rank=None, world_size=None, verbose=False, config_dir=None):

    # if 'mscoco' in split:
    # verbose = (gpu == 0)

    if 'CC' in split:
        dataset = CCDataset(split, data_dir=data_dir, mode=mode, topk=topk, args=args, verbose=verbose, rank=gpu, config_dir=config_dir)
    else:
        dataset = COCOCaptionFineTuneDataset(
            split,
            # raw_dataset=_dset,
            rank=gpu,
            topk=topk,
            verbose=verbose,
            args=args,
            mode=mode, data_dir=data_dir)


    if distributed and mode == 'train':
        train_sampler = DistributedSampler(dataset, num_replicas=world_size, rank=local_rank)
        # train_sampler = DistributedSampler(dataset)
        # train_sampler = RandomNonreplacmentSampler(dataset, dataset.n_iter)
    else:
        train_sampler = None
    if mode == 'train':
        loader = DataLoader(
            dataset, batch_size=batch_size, shuffle=(train_sampler is None),
            num_workers=workers, pin_memory=True, sampler=train_sampler,
            collate_fn=dataset.collate_fn)
    else:
        loader = DataLoader(
            dataset,
            batch_size=batch_size, shuffle=False,
            num_workers=workers, pin_memory=True,
            sampler=None,
            collate_fn=dataset.collate_fn,
            drop_last=False)

    if verbose:
        loader.evaluator = COCOCaptionEvaluator()

    loader.task = 'caption'

    return loader



class COCOCaptionEvaluator:
    def __init__(self):
        import language_evaluation
        self.evaluator = language_evaluation.CocoEvaluator(verbose=False)


    def evaluate(self, predicts, answers):

        results = self.evaluator.run_evaluation(predicts, answers)

        return results