data/refcoco_dataset.py

# ------------------------------------------------------------------------
# Modified from OFA (https://github.com/OFA-Sys/OFA)
# Copyright 2022 The OFA-Sys Team. 
# All rights reserved.
# This source code is licensed under the Apache 2.0 license 
# found in the LICENSE file in the root directory.
# ------------------------------------------------------------------------
# Modifications Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: Apache-2.0

from io import BytesIO

import logging
import warnings

import numpy as np
import torch
import base64
import utils.transforms as T
import math
from PIL import Image, ImageFile

from data import data_utils
from data.base_dataset import BaseDataset
from bert.tokenization_bert import BertTokenizer
from data.poly_utils import string_to_polygons, downsample_polygons, polygons_to_string, points_to_token_string
import cv2

ImageFile.LOAD_TRUNCATED_IMAGES = True
ImageFile.MAX_IMAGE_PIXELS = None
Image.MAX_IMAGE_PIXELS = None

logger = logging.getLogger(__name__)
warnings.filterwarnings("ignore", "(Possibly )?corrupt EXIF data", UserWarning)

IMAGENET_DEFAULT_MEAN = (0.485, 0.456, 0.406)
IMAGENET_DEFAULT_STD = (0.229, 0.224, 0.225)


class RefcocoDataset(BaseDataset):
    def __init__(
        self,
        split,
        dataset,
        bpe,
        src_dict,
        tgt_dict=None,
        max_src_length=80,
        max_tgt_length=30,
        patch_image_size=512,
        imagenet_default_mean_and_std=False,
        num_bins=1000,
        max_image_size=512
    ):
        super().__init__(split, dataset, bpe, src_dict, tgt_dict)
        self.max_src_length = max_src_length
        self.max_tgt_length = max_tgt_length
        self.patch_image_size = patch_image_size
        self.num_bins = num_bins

        if imagenet_default_mean_and_std:
            mean = IMAGENET_DEFAULT_MEAN
            std = IMAGENET_DEFAULT_STD
        else:
            mean = [0.5, 0.5, 0.5]
            std = [0.5, 0.5, 0.5]

        # for positioning
        self.positioning_transform = T.Compose([
            T.RandomResize([patch_image_size], max_size=patch_image_size),
            T.ToTensor(),
            T.Normalize(mean=mean, std=std, max_image_size=max_image_size)
        ])
        self.tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')

    def __getitem__(self, index):
        data = self.dataset[index]
        if len(data) == 7:
            uniq_id, base64_str, seg64_str, text, poly_original, region_coord, poly_interpolated = data
            train = True
        else:
            uniq_id, base64_str, seg64_str, text, poly, region_coord = data
            train = False

        # load image and segmentation labels
        image = Image.open(BytesIO(base64.urlsafe_b64decode(base64_str))).convert("RGB")
        label = Image.open(BytesIO(base64.urlsafe_b64decode(seg64_str)))
        label = np.asarray(label)
        label = cv2.resize(label, [self.patch_image_size, self.patch_image_size], interpolation=cv2.INTER_NEAREST)

        w, h = image.size
        patch_image = self.positioning_transform(image, target=None)
        resize_h = self.patch_image_size
        resize_w = self.patch_image_size
        patch_mask = torch.tensor([True])

        if train:
            prob = np.random.uniform()
            if prob < 0.5:
                polygons_interpolated = string_to_polygons(poly_interpolated)
                ds_rate = np.random.randint(25, 41)
                polygons_augmented = downsample_polygons(polygons_interpolated, ds_rate)
                poly = polygons_to_string(polygons_augmented)
            else:
                poly = poly_original

        polygons = string_to_polygons(poly)
        polygons_scaled = []
        for polygon in polygons:
            n_point = len(polygon) // 2
            scale = np.concatenate([np.array([w, h]) for _ in range(n_point)], 0)
            polygon = polygon / scale
            polygon = polygon.reshape(n_point, 2)
            polygons_scaled.append(polygon)

        x0, y0, x1, y1 = region_coord.strip().split(',')
        region_points = [float(x0), float(y0), float(x1), float(y1)]
        region = np.array(region_points)

        region_points = region_points / np.array([w, h, w, h])  # scaled to [0,1]
        region_points = torch.tensor(region_points.reshape(2, 2))

        quant_box = region_points * (self.num_bins - 1)
        quant_box11 = [[math.floor(p[0]), math.floor(p[1])] for p in quant_box]
        quant_box21 = [[math.ceil(p[0]), math.floor(p[1])] for p in quant_box]
        quant_box12 = [[math.floor(p[0]), math.ceil(p[1])] for p in quant_box]
        quant_box22 = [[math.ceil(p[0]), math.ceil(p[1])] for p in quant_box]

        quant_poly = [poly * (self.num_bins - 1) for poly in polygons_scaled]
        quant_poly11 = [[[math.floor(p[0]), math.floor(p[1])] for p in poly] for poly in quant_poly]
        quant_poly21 = [[[math.ceil(p[0]), math.floor(p[1])] for p in poly] for poly in quant_poly]
        quant_poly12 = [[[math.floor(p[0]), math.ceil(p[1])] for p in poly] for poly in quant_poly]
        quant_poly22 = [[[math.ceil(p[0]), math.ceil(p[1])] for p in poly] for poly in quant_poly]

        region_coord11, _ = points_to_token_string(quant_box11, quant_poly11)
        region_coord21, _ = points_to_token_string(quant_box21, quant_poly21)
        region_coord12, _ = points_to_token_string(quant_box12, quant_poly12)
        region_coord22, token_type = points_to_token_string(quant_box22, quant_poly22)

        # compute bilinear interpolation coefficient
        delta_x1 = [0] + [p[0] - math.floor(p[0]) for p in quant_box]  # [0] for bos token
        for polygon in quant_poly:
            delta = [poly_point[0] - math.floor(poly_point[0]) for poly_point in polygon]
            delta_x1.extend(delta)
            delta_x1.extend([0])  # for separator token
        delta_x1 = delta_x1[:-1]  # there is no separator token in the end
        delta_x1 = torch.tensor(delta_x1)
        delta_x2 = 1 - delta_x1

        delta_y1 = [0] + [p[1] - math.floor(p[1]) for p in quant_box]  # [0] for bos token
        for polygon in quant_poly:
            delta = [poly_point[1] - math.floor(poly_point[1]) for poly_point in polygon]
            delta_y1.extend(delta)
            delta_y1.extend([0])  # for separator token
        delta_y1 = delta_y1[:-1]  # there is no separator token in the end
        delta_y1 = torch.tensor(delta_y1)
        delta_y2 = 1 - delta_y1

        token_type.append(2)  # 2 for eos token

        src_caption = self.pre_caption(text, self.max_src_length)

        prompt = ' which region does the text " {} " describe?'.format(src_caption)

        # tgt for input
        tgt_item11 = self.encode_text(region_coord11, use_bpe=False)
        tgt_item12 = self.encode_text(region_coord12, use_bpe=False)
        tgt_item21 = self.encode_text(region_coord21, use_bpe=False)
        tgt_item22 = self.encode_text(region_coord22, use_bpe=False)

        # tgt for output
        target_item = region_points
        for poly in polygons_scaled:
            target_item = torch.cat([target_item, torch.tensor(poly), torch.tensor([[0, 0]])], dim=0)  # [0, 0] is padding token for separator and eos

        #target_item = torch.cat([tgt_item, self.eos_item])
        prev_output_item11 = torch.cat([self.bos_item, tgt_item11])
        prev_output_item12 = torch.cat([self.bos_item, tgt_item12])
        prev_output_item21 = torch.cat([self.bos_item, tgt_item21])
        prev_output_item22 = torch.cat([self.bos_item, tgt_item22])
        example = {
            "id": uniq_id,
            "source": prompt,
            "patch_image": patch_image,
            "patch_mask": patch_mask,
            "target": target_item,
            "prev_output_tokens_11": prev_output_item11,
            "prev_output_tokens_12": prev_output_item12,
            "prev_output_tokens_21": prev_output_item21,
            "prev_output_tokens_22": prev_output_item22,
            "delta_x1": delta_x1,
            "delta_y1": delta_y1,
            "delta_x2": delta_x2,
            "delta_y2": delta_y2,
            "w_resize_ratio": torch.tensor(resize_w / w),
            "h_resize_ratio": torch.tensor(resize_h / h),
            "region_coord": torch.tensor(region),
            "token_type": torch.tensor(token_type),
            "w": torch.tensor(w),
            "h": torch.tensor(h),
            "label": label,
            "n_poly": len(polygons),
            "text": src_caption
        }
        return example

    def collate(self, samples, pad_idx, eos_idx):
        if len(samples) == 0:
            return {}

        def merge(key, padding_item):
            return data_utils.collate_tokens(
                [s[key] for s in samples],
                padding_item,
                eos_idx=eos_idx,
            )

        id = np.array([s["id"] for s in samples])
        captions = [s["source"] for s in samples]
        tokenized = self.tokenizer.batch_encode_plus(captions, padding="longest", return_tensors="pt")
        src_tokens = tokenized["input_ids"]
        att_masks = tokenized["attention_mask"]
        src_lengths = torch.LongTensor(att_masks.ne(0).long().sum())

        patch_images = torch.stack([sample['patch_image'] for sample in samples], dim=0)
        patch_masks = torch.cat([sample['patch_mask'] for sample in samples])

        w_resize_ratios = torch.stack([s["w_resize_ratio"] for s in samples], dim=0)
        h_resize_ratios = torch.stack([s["h_resize_ratio"] for s in samples], dim=0)

        delta_x1 = merge("delta_x1", 0)
        delta_y1 = merge("delta_y1", 0)
        delta_x2 = merge("delta_x2", 1)
        delta_y2 = merge("delta_y2", 1)

        region_coords = torch.stack([s['region_coord'] for s in samples], dim=0)

        target = merge("target", pad_idx)
        tgt_lengths = torch.LongTensor([s["target"].shape[0] for s in samples])
        ntokens = tgt_lengths.sum().item()

        prev_output_tokens_11 = merge("prev_output_tokens_11", pad_idx)
        prev_output_tokens_12 = merge("prev_output_tokens_12", pad_idx)
        prev_output_tokens_21 = merge("prev_output_tokens_21", pad_idx)
        prev_output_tokens_22 = merge("prev_output_tokens_22", pad_idx)

        token_type = merge("token_type", -1)
        w = torch.stack([s["w"] for s in samples], dim=0)
        h = torch.stack([s["h"] for s in samples], dim=0)
        n_poly = [s['n_poly'] for s in samples]

        labels = np.stack([sample['label'] for sample in samples], 0)
        text = [s["text"] for s in samples]
        batch = {
            "id": id,
            "nsentences": len(samples),
            "ntokens": ntokens,
            "net_input": {
                "src_tokens": src_tokens,
                "src_lengths": src_lengths,
                "att_masks": att_masks,
                "patch_images": patch_images,
                "patch_masks": patch_masks,
                "prev_output_tokens_11": prev_output_tokens_11,
                "prev_output_tokens_12": prev_output_tokens_12,
                "prev_output_tokens_21": prev_output_tokens_21,
                "prev_output_tokens_22": prev_output_tokens_22,
                "delta_x1": delta_x1,
                "delta_y1": delta_y1,
                "delta_x2": delta_x2,
                "delta_y2": delta_y2
            },
            "target": target,
            "w_resize_ratios": w_resize_ratios,
            "h_resize_ratios": h_resize_ratios,
            "region_coords": region_coords,
            "label": labels,
            "token_type": token_type,
            "w": w,
            "h": h,
            "n_poly": n_poly,
            "text": text
        }

        return batch

    def collater(self, samples, pad_to_length=None):
        """Merge a list of samples to form a mini-batch.
        Args:
            samples (List[dict]): samples to collate
        Returns:
            dict: a mini-batch containing the data of the task
        """
        return self.collate(samples, pad_idx=self.pad, eos_idx=self.eos)