tasks/refcoco_pretrain.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 dataclasses import dataclass, field
import os
import logging
from typing import Optional
import math
import numpy as np
import torch
from fairseq import metrics
from fairseq.tasks import register_task

from tasks.base_task import BaseTask, BaseConfig
from data.refcoco_pretrain_dataset import RefcocoPretrainDataset
from data.file_dataset import FileDataset
from tasks.base_task import BaseTask, BaseConfig, load_bert_pretrained_weights

logger = logging.getLogger(__name__)


@dataclass
class RefcocoPretrainConfig(BaseConfig):
    eval_acc: bool = field(
        default=False, metadata={"help": "evaluation with accuracy"}
    )
    eval_args: Optional[str] = field(
        default='{}',
        metadata={
            "help": 'generation args, e.g., \'{"beam": 4, "lenpen": 0.6}\', as JSON string'
        },
    )
    uses_ema: Optional[bool] = field(
        default=False,
        metadata={"help": "whether to use ema"},
    )
    eval_print_samples: bool = field(
        default=False, metadata={"help": "print sample generations during validation"}
    )

    max_image_size: int = field(
        default=512, metadata={"help": "max image size for normalization"}
    )
    scst: bool = field(
        default=False, metadata={"help": "Self-critical sequence training"}
    )
    scst_args: str = field(
        default='{}',
        metadata={
            "help": 'generation args for Self-critical sequence training, as JSON string'
        },
    )


@register_task("refcoco_pretrain", dataclass=RefcocoPretrainConfig)
class RefcocoPretrainTask(BaseTask):
    def __init__(self, cfg: RefcocoPretrainConfig, src_dict, tgt_dict):
        super().__init__(cfg, src_dict, tgt_dict)

    def load_dataset(self, split, epoch=1, combine=False, **kwargs):
        paths = self.cfg.data.split(',')
        assert len(paths) > 0

        if split == 'train':
            file_path = paths[(epoch - 1) % (len(paths) - 1)]
        else:
            file_path = paths[-1]
        dataset = FileDataset(file_path, self.cfg.selected_cols)

        self.datasets[split] = RefcocoPretrainDataset(
            split,
            dataset,
            self.bpe,
            self.src_dict,
            self.tgt_dict,
            max_src_length=self.cfg.max_src_length,
            max_tgt_length=self.cfg.max_tgt_length,
            patch_image_size=self.cfg.patch_image_size,
            imagenet_default_mean_and_std=self.cfg.imagenet_default_mean_and_std,
            num_bins=self.cfg.num_bins,
            max_image_size=self.cfg.max_image_size
        )

    def build_model(self, cfg):
        model = super().build_model(cfg)
        bert_path = "../../pretrained_weights/bert-base-uncased-pytorch_model.bin"
        if os.path.exists(bert_path):
            load_bert_pretrained_weights(model.encoder.bert, bert_path)
        if cfg._name == 'polyformer_b':
            swin_path = "../../pretrained_weights/swin_base_patch4_window12_384_22k.pth"
        else:
            swin_path = "../../pretrained_weights/swin_large_patch4_window12_384_22k.pth"
        if os.path.exists(swin_path):
            model.encoder.embed_images.init_weights(pretrained=swin_path)
        return model

    def _calculate_ap_score(self, hyps, refs, thresh=0.5):
        interacts = torch.cat(
            [torch.where(hyps[:, :2] < refs[:, :2], refs[:, :2], hyps[:, :2]),
             torch.where(hyps[:, 2:] < refs[:, 2:], hyps[:, 2:], refs[:, 2:])],
            dim=1
        )
        area_predictions = (hyps[:, 2] - hyps[:, 0]) * (hyps[:, 3] - hyps[:, 1])
        area_targets = (refs[:, 2] - refs[:, 0]) * (refs[:, 3] - refs[:, 1])
        interacts_w = interacts[:, 2] - interacts[:, 0]
        interacts_h = interacts[:, 3] - interacts[:, 1]
        area_interacts = interacts_w * interacts_h
        ious = area_interacts / (area_predictions + area_targets - area_interacts + 1e-6)
        return ((ious >= thresh) & (interacts_w > 0) & (interacts_h > 0)).float()

    def valid_step(self, sample, model, criterion):
        loss, sample_size, logging_output = criterion(model, sample)
        model.eval()
        if self.cfg.eval_acc:
            hyps, refs = self._inference(sample, model)
            scores = self._calculate_ap_score(hyps.float(), refs.float())
            logging_output["_score_sum"] = scores.sum().item()
            logging_output["_score_cnt"] = scores.size(0)

        return loss, sample_size, logging_output

    def reduce_metrics(self, logging_outputs, criterion):
        super().reduce_metrics(logging_outputs, criterion)

        def sum_logs(key):
            import torch
            result = sum(log.get(key, 0) for log in logging_outputs)
            if torch.is_tensor(result):
                result = result.cpu()
            return result

        def compute_score(meters):
            score = meters["_score_sum"].sum / meters["_score_cnt"].sum
            score = score if isinstance(score, float) else score.item()
            return round(score, 4)

        if sum_logs("_score_cnt") > 0:
            metrics.log_scalar("_score_sum", sum_logs("_score_sum"))
            metrics.log_scalar("_score_cnt", sum_logs("_score_cnt"))
            metrics.log_derived("score", compute_score)

    def _inference(self, sample, model):
        hyps = self.inference_step(model, sample)
        refs = sample['region_coords'].float()
        hyps = hyps * self.cfg.max_image_size
        hyps[:, ::2] /= sample['w_resize_ratios'].unsqueeze(1)
        hyps[:, 1::2] /= sample['h_resize_ratios'].unsqueeze(1)
        return hyps, refs

    def inference_step(self, model, sample):
        with torch.no_grad():
            if isinstance(model, list):
                model = model[0]
            total_len = 2
            model.eval()
            img = sample["net_input"]["patch_images"]
            b = img.shape[0]
            prev_output_token_11 = [[0] for _ in range(b)]
            prev_output_token_12 = [[0] for _ in range(b)]
            prev_output_token_21 = [[0] for _ in range(b)]
            prev_output_token_22 = [[0] for _ in range(b)]
            delta_x1 = [[0] for _ in range(b)]
            delta_y1 = [[0] for _ in range(b)]
            delta_x2 = [[1] for _ in range(b)]
            delta_y2 = [[1] for _ in range(b)]

            gen_out = [[] for _ in range(b)]

            n_bins = self.cfg.num_bins

            encoder_out = model.encoder(
                sample['net_input']['src_tokens'],
                src_lengths=sample['net_input']['src_lengths'],
                att_masks=sample['net_input']['att_masks'],
                patch_images=sample['net_input']['patch_images'],
                patch_masks=sample['net_input']['patch_masks'],
                token_embeddings=None,
                return_all_hiddens=False,
                sample_patch_num=None
            )

            for i in range(total_len):
                prev_output_tokens_11_tensor = torch.tensor(np.array(prev_output_token_11)).to(img.device).long()
                prev_output_tokens_12_tensor = torch.tensor(np.array(prev_output_token_12)).to(img.device).long()
                prev_output_tokens_21_tensor = torch.tensor(np.array(prev_output_token_21)).to(img.device).long()
                prev_output_tokens_22_tensor = torch.tensor(np.array(prev_output_token_22)).to(img.device).long()
                delta_x1_tensor = torch.tensor(np.array(delta_x1)).to(img.device)
                delta_x2_tensor = torch.tensor(np.array(delta_x2)).to(img.device)
                delta_y1_tensor = torch.tensor(np.array(delta_y1)).to(img.device)
                delta_y2_tensor = torch.tensor(np.array(delta_y2)).to(img.device)

                net_output = model.decoder(
                    prev_output_tokens_11_tensor,
                    prev_output_tokens_12_tensor,
                    prev_output_tokens_21_tensor,
                    prev_output_tokens_22_tensor,
                    delta_x1_tensor,
                    delta_y1_tensor,
                    delta_x2_tensor,
                    delta_y2_tensor,
                    code_masks=None,
                    encoder_out=encoder_out,
                    features_only=False,
                    alignment_layer=None,
                    alignment_heads=None,
                    src_lengths=sample['net_input']['src_lengths'],
                    return_all_hiddens=False
                )
                net_output = net_output[1]
                for j in range(b):
                    output_j_x, output_j_y = net_output[j, i].cpu().numpy()
                    gen_out[j].extend([output_j_x, output_j_y])

                    output_j_x = output_j_x * (n_bins - 1)
                    output_j_y = output_j_y * (n_bins - 1)

                    output_j_x_floor = math.floor(output_j_x)
                    output_j_y_floor = math.floor(output_j_y)
                    output_j_x_ceil = math.ceil(output_j_x)
                    output_j_y_ceil = math.ceil(output_j_y)

                    # convert to token
                    prev_output_token_11[j].append(output_j_x_floor * n_bins + output_j_y_floor + 4)
                    prev_output_token_12[j].append(output_j_x_floor * n_bins + output_j_y_ceil + 4)
                    prev_output_token_21[j].append(output_j_x_ceil * n_bins + output_j_y_floor + 4)
                    prev_output_token_22[j].append(output_j_x_ceil * n_bins + output_j_y_ceil + 4)

                    delta_x = output_j_x - output_j_x_floor
                    delta_y = output_j_y - output_j_y_floor
                    delta_x1[j].append(delta_x)
                    delta_y1[j].append(delta_y)
                    delta_x2[j].append(1-delta_x)
                    delta_y2[j].append(1-delta_y)
        return torch.tensor(gen_out).to(img.device)