detectron2/data/clip_datasets/clip_img_txt_pair_tsv.py

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os
from io import BytesIO
import json
import logging
import base64
import threading
import random
import numpy as np
from typing import Callable, List, Tuple, Union
from PIL import Image
from PIL import ImageFile
ImageFile.LOAD_TRUNCATED_IMAGES = True
import torch
import torch.utils.data as data
from .oscar_tsv import InputExample, convert_example_to_features
from detectron2.structures.tsv_file import TSVFile, CompositeTSVFile
from detectron2.data.clip_datasets.clip_prompt_engineering import get_prompt_templates, prompt_engineering
#import spacy

def pre_fetch(tsv_filename: str):
    logging.info('Pre-loading %s ...' % tsv_filename)
    with open(tsv_filename, 'r'):
        logging.info('Pre-loading %s ended.' % tsv_filename)

class CLIPImgTxtPairTSVDataset(data.Dataset):
    """
        This class is intended for encapsulating Image/Text pair data for contrastive learning described in
        the following paper,
        "Learning Transferable Visual Models From Natural Language Supervision" (a.k.a CLIP)
        Specifically, it is used to accomadate the tsv data format from Azure Cognition Service Group.
    """
    def __init__(self,
                 image_tsv_file: Union[str, List[str]],
                 text_tsv_file: Union[str, List[str]],
                 transforms: Callable = None,
                 tokenizer: Callable = None,
                 seq_len = 0, context_length = 77, target_offset=0, 
                 args = None, 
                 dataset_name = "", 
                 tokenizer_type = "bert",
                 is_train = True,
                 map_file = None, 
                 filtered_datasets = ''):
        self.args = args
        self.is_train = is_train    
        self.dataset_names = dataset_name
        self.tokenizer_type = tokenizer_type
        self.target_offset = target_offset
        self.seq_len = seq_len

        self.transforms = transforms
        self.tokenizer = tokenizer
        self._chunk_sizes = None
        self.context_length = context_length
        
        self.prompt_templates = get_prompt_templates() # [:2]
        self.spacy_nlp = None # spacy.load('en_core_web_sm')
        
        self.class_selector = None
        # self.class_selector = list(self.label2idx.keys()) if self.label2idx else None

        self.label2idx = {}
        self.idx2label = {}        
        self.classnames = {}
        self.dataset_target_offsets = {}; offset = 0

        self.num_classes = sum([len(val) for val in self.classnames.values()])          

        self.filtered_classnames = []

        if isinstance(image_tsv_file, str) and isinstance(text_tsv_file, str):
            # single tsv file
            if (
                os.path.splitext(image_tsv_file)[1].lower() == '.tsv'
                and os.path.splitext(text_tsv_file)[1].lower() == '.tsv'
            ):
                self.image_tsv_file = TSVFile(image_tsv_file, if_generate_lineidx=True)
                self.text_tsv_file = TSVFile(text_tsv_file, if_generate_lineidx=True)
            # multiple tsv files specified in a text file
            elif (
                os.path.splitext(image_tsv_file)[1].lower() == '.txt'
                and os.path.splitext(text_tsv_file)[1].lower() == '.txt'
            ):
                self.image_tsv_file = CompositeTSVFile(image_tsv_file)
                self.text_tsv_file = CompositeTSVFile(text_tsv_file)
                self._chunk_sizes = self.image_tsv_file.get_chunk_size()
            else:
                raise ValueError("Invalid input! Please check the tsv filenames.")
        # multiple tsv files specified in a list
        elif (
            isinstance(image_tsv_file, list)
            and isinstance(text_tsv_file, list)
        ):
            assert len(image_tsv_file) == len(text_tsv_file), \
                "Inconsistent number of Image/Text tsv files!"
            assert len(image_tsv_file) == len(text_tsv_file), \
                "Inconsistent number of Image/Text tsv files!"
            self.image_tsv_path = image_tsv_file
            self.text_tsv_path = text_tsv_file            
            self.image_tsv_file = CompositeTSVFile(image_tsv_file, class_selector=self.class_selector)
            self.text_tsv_file = CompositeTSVFile(text_tsv_file, class_selector=self.class_selector)
            self._chunk_sizes = self.image_tsv_file.get_chunk_size()
            self._accumulated_chunk_sizes = np.cumsum(self._chunk_sizes).tolist()
        else:
            raise ValueError("Invalid input! Please check the tsv filenames.")
        
        assert len(self.image_tsv_file) == len(self.text_tsv_file), \
            "Inconsistent size of Image/Text ({}/{}) data!".format(
                len(self.image_tsv_file), len(self.text_tsv_file)
            )
    
    def get_chunk_sizes(self):
        return self._chunk_sizes

    def get_class_boundaries(self):
        # The samples of each class are organized class-by-class.
        # _class_boundaries stores the lower- and upper-bound of each class.
        return self.image_tsv_file.get_class_boundaries()

    def _load_map(self, map_file: str):
        if not map_file:
            return None

        label2idx = {}
        with open(map_file) as f:
            for line in f:
                items = line.strip().split('\t')
                label2idx[items[0]] = int(items[1])

        return label2idx
    
    def _load_darknet_map(self, map_file):
        if not map_file:
            return None

        label2idx = {}
        with open(map_file) as f:
            linenum = 0
            for l in f:
                item = l.strip()
                label2idx[item] = linenum
                linenum += 1

        return label2idx    

    def _pre_tokenize(self):
        """
        pre-tokenize class names
        """
        input_ids_all = []
        input_masks_all = []
        segment_ids_all = []
        for k in range(len(self.classnames["imagenet"])):
            cur_id = 0; img_id = 0
            scale = 1.0

            v = self.classnames["imagenet"].label_to_name(k)
            if isinstance(v, str):
                vs = [v]
            elif isinstance(v, list):
                vs = v
            t1s = []
            t2s = []
            for v in vs:
                for pt in self.prompt_templates:
                    t1s.append(prompt_engineering(v, template=pt))
                    t2s.append("")
            input_ids = []
            input_masks = []
            segment_ids = []
            is_next_labels = [0] * len(t1s)
            is_img_matchs = [1] * len(t1s)
            img_feat_len = 0
            for t1, t2, is_next_label, is_img_match in zip(t1s, t2s, is_next_labels, is_img_matchs):
                if self.tokenizer_type == "bert":
                    # tokenize
                    tokens_a = self.tokenizer.tokenize(t1)
                    tokens_b = None

                    # combine to one sample
                    cur_example = InputExample(guid=cur_id, tokens_a=tokens_a,
                                            tokens_b=tokens_b, is_next=is_next_label,
                                            img_id=img_id, is_img_match=is_img_match)

                    # transform sample to features
                    cur_features = convert_example_to_features(self.args, cur_example,
                                                            self.seq_len, self.tokenizer,
                                                            img_feat_len)                                                        

                    input_ids.append(torch.tensor(cur_features.input_ids, dtype=torch.long))
                    input_masks.append(torch.tensor(cur_features.input_mask, dtype=torch.long))
                    segment_ids.append(torch.tensor(cur_features.segment_ids, dtype=torch.long))

                elif self.tokenizer_type == "bpe":
                    tokens_a = t1; tokens_b = None
                    # combine to one sample
                    cur_example = InputExample(guid=cur_id, tokens_a=tokens_a,
                                            tokens_b=tokens_b, is_next=is_next_label,
                                            img_id=img_id, is_img_match=is_img_match)

                    # transform sample to features
                    cur_features = convert_example_to_features_bpe(self.args, cur_example,
                                                            self.seq_len, self.tokenizer,
                                                            img_feat_len)                                                        

                    input_ids.append(torch.tensor(cur_features.input_ids, dtype=torch.long))
                    input_masks.append(torch.tensor(cur_features.input_mask, dtype=torch.long))
                    segment_ids.append(torch.tensor(cur_features.segment_ids, dtype=torch.long))

                else:
                    raise NotImplementedError
            input_ids_all.append(torch.stack(input_ids, 0))
            input_masks_all.append(torch.stack(input_masks, 0))
            segment_ids_all.append(torch.stack(segment_ids, 0))

        self.input_ids_all_classes = torch.stack(input_ids_all, 0)
        self.input_mask_all_classes = torch.stack(input_masks_all, 0)
        self.segment_ids_all_classes = torch.stack(segment_ids_all, 0)

    def _online_tokenize(self, text):

        # random select a prompt template
        temp_idx = np.random.randint(len(self.prompt_templates))
        pt = self.prompt_templates[temp_idx]

        names = text.split(";")
        num_names = np.random.randint(len(names)) + 1
        names_sampled = random.sample(names, num_names)
        text = ", ".join(names_sampled)    

        t1 = prompt_engineering(text, template=pt)

        cur_id = 0; img_id = 0; scale = 1.0
        is_next_label = 0; is_img_match = 1
        img_feat_len = 0

        if self.tokenizer_type == "bert":
            # tokenize
            tokens_a = self.tokenizer.tokenize(t1)
            tokens_b = None

            # combine to one sample
            cur_example = InputExample(guid=cur_id, tokens_a=tokens_a,
                                    tokens_b=tokens_b, is_next=is_next_label,
                                    img_id=img_id, is_img_match=is_img_match)

            # transform sample to features
            cur_features = convert_example_to_features(self.args, cur_example,
                                                    self.context_length, self.tokenizer,
                                                    img_feat_len)                                                        


        elif self.tokenizer_type == "bpe":
            tokens_a = t1; tokens_b = None
            # combine to one sample
            cur_example = InputExample(guid=cur_id, tokens_a=tokens_a,
                                    tokens_b=tokens_b, is_next=is_next_label,
                                    img_id=img_id, is_img_match=is_img_match)

            # transform sample to features
            cur_features = convert_example_to_features_bpe(self.args, cur_example,
                                                    self.context_length, self.tokenizer,
                                                    img_feat_len)                                                              

        return torch.tensor(cur_features.input_ids, dtype=torch.long), \
                torch.tensor(cur_features.input_mask, dtype=torch.long), \
                torch.tensor(cur_features.segment_ids, dtype=torch.long)

    def get_dataset_name(self, index):
        """
        get dataset name according to index
        """
        assert index < self._accumulated_chunk_sizes[-1], "index must in the range of accumulated data size"
        for k, boundary in enumerate(self._accumulated_chunk_sizes):
            if index < boundary:
                return self.dataset_names[k], k

    def get_target_offset(self, dataset_name):
        return self.dataset_target_offsets[dataset_name]

    def get_img_label_pair(self, items_image, index):             
        dataset_name, chunk_id = self.get_dataset_name(index)
        target_offset = self.get_target_offset(dataset_name)
        _, target, img = self._decode_data(items_image, dataset_name)   

        if self.transforms:
            img = self.transforms(img)

        if target == -1:
            input_ids, input_mask, segment_ids = \
                self._online_tokenize("uncovered image")                    
        else:
            classname = self.classnames[dataset_name].labels2names[self.idx2label[dataset_name][target]]
            if classname in self.filtered_classnames:
                # we filter these classnames for training
                target = -1
                input_ids, input_mask, segment_ids = \
                    self._online_tokenize("uncovered image")                  
            else:            
                input_ids, input_mask, segment_ids = \
                    self._online_tokenize(classname)
                target += target_offset
        return img, \
            input_ids, \
            input_mask, \
            segment_ids, \
            torch.LongTensor([target]), \
            dataset_name                          
        
    def get_img_txt_pair(self, items_image, items_text, index):
        dataset_name, chunk_id = self.get_dataset_name(index)
        assert items_text[0] == items_image[0], \
            'keys do not match for image ({}) and text ({}) for {} at chunk {}-{}'.format(
                len(items_text[0]), len(items_image[0]), dataset_name, chunk_id, self.image_tsv_path[chunk_id]
            )

        img = self._decode_image(items_image, dataset_name)
        #     print("index {}, chunk id {}, name {}".format(index, chunk_id, self.image_tsv_path[chunk_id]))
        #     raise TypeError("cannot decode current item") 
        img_width, img_height = img.size  # img_height, img_width = np.array(img).shape

        txts = self._decode_text(items_text)
        if self.spacy_nlp is not None:
            np_input_ids, np_input_masks, np_segment_ids = self.create_phrase_text(txts)

        if self.transforms:
            img = self.transforms(img)

        if isinstance(txts, str):
            input_ids, input_masks, segment_ids = \
                convert_txt_to_tokens_bpe(txts, self.tokenizer, self.context_length)
            all_str2id_links = []
        elif isinstance(txts, list):
            input_ids = []
            input_masks = []
            segment_ids = []
            all_str2id_links = []
            for txt in txts:
                input_id, input_mask, segment_id, str2id_links = \
                    convert_txt_to_tokens_bpe(txt, self.tokenizer, self.context_length, return_link=True)
                input_ids += input_id          
                input_masks += input_mask
                segment_ids += segment_id
                all_str2id_links += [str2id_links]
        scale = 1.0
        img_id = 0
        
        if self.spacy_nlp is not None:
            return img, \
                    torch.tensor(input_ids).long().view(-1), \
                    torch.tensor(input_masks).long().view(-1), \
                    torch.tensor(segment_ids).long().view(-1), \
                    torch.LongTensor([1e5]), \
                    dataset_name, \
                    torch.tensor(np_input_ids).long().view(-1), \
                    torch.tensor(np_input_masks).long().view(-1), \
                    torch.tensor(np_segment_ids).long().view(-1)
        else:
            return img, \
                    torch.tensor(input_ids).long().view(-1), \
                    torch.tensor(input_masks).long().view(-1), \
                    torch.tensor(segment_ids).long().view(-1), \
                    torch.LongTensor([1e5]), \
                    (dataset_name, items_text[0], (img_height, img_width), all_str2id_links)  # dataset name, image id, image height&width, links bet string and tokenized texts

    def create_phrase_text(self, txt_list):
        """ Use NLP tool to detect noun phrases in captions, fill each identified phrase into a random prompt to create a sentence,
            and convert each sentence to bpe tokens
        """
        if isinstance(txt_list, str):
            txt_list = [txt_list]
        # detect noun phrase
        noun_phrase = []
        for txt in txt_list:
            doc = self.spacy_nlp(txt.lower())
            this_text = [nc.text for nc in doc.noun_chunks]
            this_text = [nc.replace('a ', '').replace('the ', '') for nc in this_text]
            noun_phrase.extend(this_text)
        noun_phrase = list(set(noun_phrase))
        # fill each phrase into a random prompt
        text_list = []
        pts = random.sample(self.prompt_templates, len(noun_phrase))
        for i, np in enumerate(noun_phrase):
            text_list.append(prompt_engineering(np, pts[i]))
        # convert string into bpe tokens
        input_ids = []
        input_masks = []
        segment_ids = []
        for txt in text_list:
            input_id, input_mask, segment_id = \
                convert_txt_to_tokens_bpe(txt, self.tokenizer, self.context_length)
            input_ids += input_id          
            input_masks += input_mask
            segment_ids += segment_id
        return input_ids, input_masks, segment_ids

    def __getitem__(self, index: Union[int, Tuple[int, int]]):
        if isinstance(index, tuple):
            items_image = self.image_tsv_file[index[0]]
            items_text = self.text_tsv_file[index[0]]
            if index[1] >= 0:
                tsv_filename = self.image_tsv_file.file_list[index[1]]

                # Python threads are not truly parallel. Spawn a new process instead.
                # logging.info('Pre-loading %s ...' % tsv_filename)
                # os.system('cat ' + tsv_filename + ' > /dev/null &')
                x = threading.Thread(
                   target=pre_fetch, args=(tsv_filename,), daemon=True
                )
                x.start()
            curr_index = index[0]
        else:
            items_image = self.image_tsv_file[index]
            items_text = self.text_tsv_file[index]
            curr_index = index
        
        # NOTE: since we duplicate image tsv to text tsv for image-label data,
        # we can determine whether the current instance is an image-label pair or 
        # a image-text pair data based on whether items_image is identical to items_text or not.
        if items_image == items_text:
            return self.get_img_label_pair(items_image, curr_index)
        else:
            return self.get_img_txt_pair(items_image, items_text, curr_index)
            
    def _decode_image(self, items: Tuple[str, str], dataset_name=""):
        key = items[0]
        image = Image.open(BytesIO(base64.b64decode(items[1]))).convert('RGB')
        return image

    def _decode_text(self, items: Tuple[str, Union[str, dict]]):
        key = items[0]
        text = ''
        if isinstance(items[1], str):
            try:
                str_dict = json.loads(items[1])
                # in this dict, it may contain either "tags" or "captions" or both
                keys = [key for key in str_dict.keys()]
                selected_key = random.sample(keys, 1)[0]
                if selected_key == "captions":
                    # if this is a caption, we sample a caption
                    captions = str_dict[selected_key]
                    text = captions[:5]
                    # text = random.sample(captions, 1)[0]
                elif selected_key == "tags":
                    # for tags, we randomly disorder it
                    tags = str_dict[selected_key]
                    tag_words = tags.split(' ')
                    random.shuffle(tag_words)
                    tags_shuffled = " ".join(tag_words)
                    # add prompt template
                    pt = random.sample(self.prompt_templates, 1)[0]
                    text = prompt_engineering(tags_shuffled, pt)
            except:
                text = items[1]
        elif isinstance(items[1], dict):
            assert 'captions' in items[1], '"captions" does not in {}'.format(items[1])
            captions = items[1]['captions']
            if isinstance(captions, list):
                text = random.choice(captions)
            elif isinstance(captions, str):
                text = captions
            else:
                raise ValueError('captions should be str or list')

        return text

    def _decode_data(self, items, dataset_name):
        key = items[0]
        label = self._get_label(items[1], dataset_name)
        try:
            image = Image.open(BytesIO(base64.b64decode(items[2])))
        except:
            return None

        return key, label, image.convert('RGB')

    def _get_label(self, item, dataset_name):
        if not self.label2idx[dataset_name]:
            return int(item)

        if item in self.label2idx[dataset_name]:
            return self.label2idx[dataset_name][item]

        label = json.loads(item)[0]['class']
        if label in self.label2idx[dataset_name]:
            return self.label2idx[dataset_name][label]
        else:
            return -1

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

def convert_txt_to_tokens_bpe(text, tokenizer, context_length, return_link=False):

    sot_token = tokenizer.encoder["<|startoftext|>"]
    eot_token = tokenizer.encoder["<|endoftext|>"]
    if return_link:
        bpe_tokens, str2id_links = tokenizer.encode(text, return_link=return_link)
        str2id_links = [["<|startoftext|>", [sot_token]]] + str2id_links + [["<|endoftext|>", [eot_token]]]
    else:
        bpe_tokens = tokenizer.encode(text, return_link=return_link)
    input_ids = [sot_token] + bpe_tokens + [eot_token]

    if len(input_ids) > context_length:
        input_ids = input_ids[:context_length]
    segment_ids = [0] * len(input_ids)
    lm_label_ids = [-1] * len(input_ids)

    # The mask has 1 for real tokens and 0 for padding tokens. Only real tokens are attended to.
    input_mask = [1] * len(input_ids)

    # Zero-pad up to the sequence length.
    while len(input_ids) < context_length:
        input_ids.append(0)
        input_mask.append(0)
        segment_ids.append(0)
        lm_label_ids.append(-1)

    assert len(input_ids) == context_length
    assert len(input_mask) == context_length
    assert len(segment_ids) == context_length
    assert len(lm_label_ids) == context_length    
    
    if return_link:
        return input_ids, input_mask, segment_ids, str2id_links
    return input_ids, input_mask, segment_ids

def convert_example_to_features_bpe(args, example, max_seq_length, tokenizer,
                                img_feat_len, context_length=77):
    """
    Convert a raw sample (pair of sentences as tokenized strings) into a proper training sample with
    IDs, LM labels, input_mask, CLS and SEP tokens etc.
    :param args: parameter settings
    :param img_feat_len: lens of actual img features
    :param example: InputExample, containing sentence input as strings and is_next label
    :param max_seq_length: int, maximum length of sequence.
    :param tokenizer: Tokenizer
    :return: InputFeatures, containing all inputs and labels of one sample as IDs (as used for model training)
    """
    # we do not consider tokens_b for now in original CLIP
    text = example.tokens_a
    assert isinstance(text, str)

    sot_token = tokenizer.encoder["<|startoftext|>"]
    eot_token = tokenizer.encoder["<|endoftext|>"]
    input_ids = [sot_token] + tokenizer.encode(text) + [eot_token]

    if len(input_ids) > context_length:
        input_ids = input_ids[:context_length]
    segment_ids = [0] * len(input_ids)
    lm_label_ids = [-1] * len(input_ids)

    # The mask has 1 for real tokens and 0 for padding tokens. Only real tokens are attended to.
    input_mask = [1] * len(input_ids)

    # Zero-pad up to the sequence length.
    while len(input_ids) < context_length:
        input_ids.append(0)
        input_mask.append(0)
        segment_ids.append(0)
        lm_label_ids.append(-1)

    assert len(input_ids) == context_length
    assert len(input_mask) == context_length
    assert len(segment_ids) == context_length
    assert len(lm_label_ids) == context_length

    if example.guid < 1:
        logging.info("*** Example ***")
        logging.info("guid: %s" % example.guid)
        logging.info("input_ids: %s" % " ".join([str(x) for x in input_ids]))
        logging.info("input_mask: %s" % " ".join([str(x) for x in input_mask]))
        logging.info("segment_ids: %s" % " ".join([str(x) for x in segment_ids]))
        logging.info("LM label: %s " % lm_label_ids)
        logging.info("Is next sentence label: %s " % example.is_next)

    features = InputFeatures(input_ids=input_ids,
                             input_mask=input_mask,
                             segment_ids=segment_ids,
                             lm_label_ids=lm_label_ids,
                             is_next=example.is_next,
                             img_feat_len=img_feat_len,
                             is_img_match=example.is_img_match)
    return features

class InputFeatures(object):
    """A single set of features of data."""

    def __init__(self, input_ids, input_mask, segment_ids, is_next,
                 lm_label_ids, img_feat_len, is_img_match):
        self.input_ids = input_ids
        self.input_mask = input_mask
        self.segment_ids = segment_ids
        self.is_next = is_next
        self.lm_label_ids = lm_label_ids

        self.img_feat_len = img_feat_len
        self.is_img_match = is_img_match