maskrcnn_benchmark/data/datasets/caption.py

import torch
import torch.distributed as dist
import time
from torchvision.ops import nms
import random
import numpy as np
from PIL import Image, ImageDraw
import pdb
from maskrcnn_benchmark.structures.bounding_box import BoxList
from .modulated_coco import ConvertCocoPolysToMask
from .tsv import ODTSVDataset, TSVYamlDataset
from .od_to_grounding import sanity_check_target_after_processing
from maskrcnn_benchmark.data.datasets._caption_aug import CaptionAugmentation
from collections import defaultdict

class CaptionTSV(TSVYamlDataset):
    def __init__(

        self,

        yaml_file,

        transforms,

        return_tokens,

        return_masks,

        tokenizer,

        caption_min_box=1,

        replace_clean_label=False,

        further_screen=False,

        caption_conf=0.5,

        caption_nms=-1,

        pack_random_caption_number=0,

        inference_caption=False,

        sample_negative_for_grounding_data=-1,

        random_pack_prob=-1.0,

        no_random_pack_probability=0.0,

        safeguard_positive_caption=True,

        mlm_obj_for_only_positive=False,

        caption_format_version="v1",

        local_debug=False,

        max_query_len=256,

        cc_caption_augmentation_version=None,

        caption_vocab_file=None,

        **kwargs

    ):
        super(CaptionTSV, self).__init__(yaml_file, None, replace_clean_label)
        self.yaml_file = yaml_file
        self._transforms = transforms
        self.max_query_len = 225
        self.prepare = ConvertCocoPolysToMask(
            return_masks=return_masks, return_tokens=return_tokens, tokenizer=tokenizer, max_query_len=max_query_len
        )
        self.tokenizer = tokenizer
        self.caption_min_box = caption_min_box
        self.replace_clean_label = replace_clean_label
        self.further_screen = further_screen
        self.pack_random_caption_number = pack_random_caption_number
        self.caption_format_version = caption_format_version

        self.caption_conf = caption_conf
        self.caption_nms = caption_nms
        self.inference_caption = inference_caption
        self.sample_negative_for_grounding_data = sample_negative_for_grounding_data
        self.random_pack_prob = random_pack_prob
        self.no_random_pack_probability = no_random_pack_probability
        self.safeguard_positive_caption = safeguard_positive_caption
        self.mlm_obj_for_only_positive = mlm_obj_for_only_positive
        try:
            self.rank = dist.get_rank()
        except:
            self.rank = 0
        self.caption_augmentation_version = cc_caption_augmentation_version
        if self.caption_augmentation_version is not None:
            self.caption_augmentation = CaptionAugmentation(
                self.caption_augmentation_version,
                tokenizer,
                caption_vocab_file=caption_vocab_file
            )

    def __len__(self):
        return super(CaptionTSV, self).__len__()

    def pack_caption(self, positive_caption, negative_captions, original_tokens_positive):
        if len(negative_captions) == 0:
            return positive_caption, original_tokens_positive, [(0, len(positive_caption))]
        if self.safeguard_positive_caption:
            length_of_each_caption = []
            for caption in negative_captions + [positive_caption]:
                tokenized = self.tokenizer(caption, return_tensors="pt")
                length_of_each_caption.append(tokenized.input_ids.size(-1))
            max_length = self.max_query_len - length_of_each_caption[-1]
            indexes = list(range(len(negative_captions)))
            random.shuffle(indexes)
            new_caption_list = [positive_caption]
            for i in indexes:
                if length_of_each_caption[i] < max_length:
                    new_caption_list.append(negative_captions[i])
                    max_length -= length_of_each_caption[i]
        else:
            new_caption_list = [positive_caption] + negative_captions
        random.shuffle(new_caption_list)

        new_caption = ""

        for i in new_caption_list:
            if i == positive_caption:
                start_position = len(new_caption)
            new_caption += i
            if not i.endswith("."):
                new_caption += "."
            new_caption += " "

        # shift the token positions the boxes are aligned to
        for index, i in enumerate(original_tokens_positive):
            original_tokens_positive[index] = [tuple(j) for j in i]
        for i in original_tokens_positive:
            for index, j in enumerate(i):
                i[index] = (j[0] + start_position, j[1] + start_position)

        return new_caption, original_tokens_positive, [(start_position, start_position + len(positive_caption))]

    def __get_negative_captions__(self, idx, negative_size=7):
        negative_captions = []
        for i in range(negative_size):
            img, anno, _, scale = super(CaptionTSV, self).__getitem__(np.random.choice(len(self)))
            caption = anno["caption"]
            negative_captions.append(caption)

        return negative_captions

    def target_transpose_in(self, anno):
            # for the target from "caption", we need to transpose to box format
            new_target = []
            for box in range(len(anno["bboxes"])):
                new_box = {}
                new_box["tokens_positive"] = anno["tokens_positive"][box]
                new_box["nouns"] = anno["all_nounds_in_vocab"][box]
                new_box["bbox"] = anno["bboxes"][box]
                new_target.append(new_box)
            return new_target

    def target_transpose_out(self, target):
        # for the target from "caption", we need to transpose to box format
        new_target = defaultdict(list)

        for box in target:
            new_target["bboxes"].append(box["bbox"])
            new_target["tokens_positive"].append(box["tokens_positive"])
            if "spans_positive" in box:
                new_target["spans_positive"].append(box["spans_positive"])
        return new_target

    def __getitem__(self, idx):
        try:
            img, anno, _, scale = super(CaptionTSV, self).__getitem__(idx)
            if self.inference_caption:
                caption = None
                if isinstance(anno, list):
                    caption = anno[0]["caption"]  # inference mode for bing
                    anno = []
                elif len(anno) == 1:
                    caption = anno["caption"]  # inference mode for googlecc
                    anno = []
                else:
                    caption = " ".join(anno["captions"])
                    anno = []
            else:
                """

                An example

                {'img_h': 1154, 'img_w': 1600, 'caption': 'xxx', 'tokens_positive': [[[47, 50], [51, 53], [54, 59]], [[32, 35], [36, 41]], [[32, 35], [36, 41]], [[0, 3], [3, 6], [6, 10], [11, 16], [17, 19], [20, 23]], [[32, 35], [36, 41]], [[32, 35], [36, 41]]], 'bboxes': [[7.344961166381836, 10.479412078857422, 1592.2679443359375, 1090.0028076171875], [950.32861328125, 346.572021484375, 1333.2373046875, 679.3215942382812], [927.44140625, 342.7712707519531, 1389.833984375, 719.5758666992188], [90.48786163330078, 363.67572021484375, 1381.8631591796875, 1078.687744140625], [122.84217071533203, 422.6786193847656, 507.845703125, 667.2651977539062], [80.62384033203125, 416.500244140625, 563.1666259765625, 734.603271484375]], 'scores': [0.7966700196266174, 0.8952182531356812, 0.8186006546020508, 0.9995516538619995, 0.8021856546401978, 0.8923134803771973]}

                """
                if len(anno["bboxes"]) < self.caption_min_box:  # Retry triggered!
                    return self[np.random.choice(len(self))]
                if self.caption_format_version == "v2":
                    anno = self.convert_anno_from_v2_to_v1(anno)

                if self.further_screen:
                    conf = self.caption_conf
                    nms_thre = self.caption_nms

                    bboxes = torch.as_tensor(anno["bboxes"]).float()
                    scores = torch.as_tensor(anno["scores"])
                    tokens_positive = anno["tokens_positive"]
                    if "all_nounds_in_vocab" in anno:
                        all_nounds_in_vocab = anno["all_nounds_in_vocab"]
                    else:
                        all_nounds_in_vocab = []
                    # print("\n\n\n\n tokens_positive in original data", tokens_positive)

                    keep = scores > conf
                    scores = scores[keep]
                    bboxes = bboxes[keep]
                    tokens_positive = [i for index, i in enumerate(tokens_positive) if keep[index]]
                    all_nounds_in_vocab = [i for index, i in enumerate(all_nounds_in_vocab) if keep[index]]

                    assert len(tokens_positive) == len(bboxes) == len(scores)

                    if len(bboxes) < self.caption_min_box:  # Retry triggered!
                        return self[np.random.choice(len(self))]

                    if nms_thre > 0:
                        keep = nms(boxes=bboxes, scores=scores, iou_threshold=nms_thre)
                        scores = scores[keep]
                        bboxes = bboxes[keep]
                        tokens_positive = [tokens_positive[i] for i in keep]
                        assert len(tokens_positive) == len(bboxes) == len(scores)

                    # Write back
                    anno["bboxes"] = bboxes.tolist()
                    anno["scores"] = scores.tolist()
                    anno["tokens_positive"] = tokens_positive
                    anno["all_nounds_in_vocab"] = all_nounds_in_vocab

                if len(anno["bboxes"]) < self.caption_min_box:  # Retry triggered!
                    return self[np.random.choice(len(self))]

                if self.caption_augmentation_version is not None:
                    caption, new_anno, spans = self.caption_augmentation(
                        anno["caption"],
                        self.target_transpose_in(anno),
                        gpt3_outputs = anno.get("gpt3_outputs", None))
                    anno.update(self.target_transpose_out(new_anno))
                    anno["caption"] = caption
                    do_neg_aug = False
                else:
                    do_neg_aug = True
                    spans = None

                boxes = torch.as_tensor(anno["bboxes"])
                caption = anno["caption"]
                target = BoxList(boxes, (anno["img_w"], anno["img_h"]), mode="xyxy")
                target = target.clip_to_image(remove_empty=True)
                if spans is not None:
                    target.add_field("spans", spans) # add spans to target
                #pdb.set_trace()
                # print("original caption", caption)
                empty_everything = False
                if self.sample_negative_for_grounding_data != -1:
                    if random.random() < self.sample_negative_for_grounding_data:
                        empty_everything = True

                if empty_everything:
                    caption = self.__get_negative_captions__(idx, negative_size=1)[0]

                if self.pack_random_caption_number != 0 and do_neg_aug:
                    if self.random_pack_prob != -1.0:
                        if random.random() < self.no_random_pack_probability:
                            negative_pack_number = 0
                        elif random.random() < self.random_pack_prob:
                            negative_pack_number = self.pack_random_caption_number
                        else:
                            negative_pack_number = np.random.choice(self.pack_random_caption_number)
                    else:
                        negative_pack_number = self.pack_random_caption_number

                    negative_captions = self.__get_negative_captions__(idx, negative_size=negative_pack_number)

                    caption, anno["tokens_positive"], greenlight_span_for_masked_lm_objective = self.pack_caption(
                        caption, negative_captions, anno["tokens_positive"]
                    )
                else:
                    greenlight_span_for_masked_lm_objective = [(0, len(caption))]

                if not self.mlm_obj_for_only_positive:
                    greenlight_span_for_masked_lm_objective = [(0, len(caption))]

                new_anno = []
                areas = target.area()
                for i in range(len(target)):
                    new_anno_i = {}
                    new_anno_i["area"] = areas[i]
                    new_anno_i["iscrowd"] = 0
                    new_anno_i["image_id"] = idx
                    new_anno_i["category_id"] = 1  # following vg and others
                    new_anno_i["id"] = None
                    new_anno_i["bbox"] = target.bbox[i].numpy().tolist()
                    new_anno_i["tokens_positive"] = anno["tokens_positive"][i]
                    if "spans_positive" in anno:
                        new_anno_i["spans_positive"] = anno["spans_positive"][i]
                    new_anno.append(new_anno_i)

                # except:
                #     return self[np.random.choice(len(self))]

                anno = new_anno
                if empty_everything:
                    anno = []

            annotations = {"image_id": idx, "annotations": anno, "caption": caption}
            annotations["greenlight_span_for_masked_lm_objective"] = greenlight_span_for_masked_lm_objective
            if "spans" in target.extra_fields:
                annotations["spans"] = target.extra_fields["spans"]
                if not isinstance(annotations["spans"], list):
                    annotations["spans"] = annotations["spans"].tolist()
            img, annotations = self.prepare(img, annotations, box_format="xyxy")
            if self._transforms is not None:
                img, target = self._transforms(img, target)

            # add additional property
            for ann in annotations:
                target.add_field(ann, annotations[ann])
        except:
            print("Outter Retry triggered!!")
            return self[np.random.choice(len(self))]

        return img, target, idx

    def convert_anno_from_v2_to_v1(self, anno):
        flatterned_bboxes = []
        flatterned_tokens_positive = []
        flatterned_bboxes_scores = []
        flatterned_nouns = []
        for i in range(len(anno["bboxes"])):
            # i is the index for entity
            for j in range(len(anno["bboxes"][i])):
                # j is the index for each box
                flatterned_bboxes.append(anno["bboxes"][i][j])
                flatterned_tokens_positive.append(
                    anno["tokens_positive"][i]
                )  # Assume this box corresponds to all the token_spans for this entity
                if "all_nounds_in_vocab" in anno:
                    flatterned_nouns.append(anno["all_nounds_in_vocab"][i])
                flatterned_bboxes_scores.append(anno["scores"][i][j])
        anno["bboxes"] = flatterned_bboxes
        anno["tokens_positive"] = flatterned_tokens_positive
        anno["scores"] = flatterned_bboxes_scores
        if "all_nounds_in_vocab" in anno:
            anno["all_nounds_in_vocab"] = flatterned_nouns
        return anno

    def get_raw_image(self, idx):
        image, *_ = super(CaptionTSV, self).__getitem__(idx)
        return image

    def get_img_id(self, idx):
        line_no = self.get_line_no(idx)
        if self.label_tsv is not None:
            row = self.label_tsv.seek(line_no)
            img_id = row[0]
            return img_id