# 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. from dataclasses import dataclass, field import json import logging from typing import Optional from argparse import Namespace import torch from fairseq import metrics from fairseq.tasks import register_task from tasks.ofa_task import OFATask, OFAConfig from data.mm_data.refcoco_dataset import RefcocoDataset from data.file_dataset import FileDataset logger = logging.getLogger(__name__) from mapcalc import calculate_map, calculate_map_range from functools import partial @dataclass class RefcocoConfig(OFAConfig): 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' }, ) 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' }, ) acc_thresh: Optional[str] = field( default=None, metadata={"help": "acc thresh for refcoco"} ) metric: Optional[str] = field( default='acc', metadata={"help": "metric"} ) max_area_size: Optional[float] = field( default=None, metadata={"help": "max_area_size"} ) min_area_size: Optional[float] = field( default=None, metadata={"help": "min_area_size"} ) @register_task("refcoco", dataclass=RefcocoConfig) class RefcocoTask(OFATask): def __init__(self, cfg: RefcocoConfig, src_dict, tgt_dict): super().__init__(cfg, src_dict, tgt_dict) self.metric = cfg.metric self.min_area_size = cfg.min_area_size self.max_area_size = cfg.max_area_size try: self.acc_thresh = float(cfg.acc_thresh) except: self.acc_thresh = cfg.acc_thresh print(self.acc_thresh, self.metric, self.min_area_size, self.max_area_size) 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] = RefcocoDataset( 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) if self.cfg.eval_acc: gen_args = json.loads(self.cfg.eval_args) self.sequence_generator = self.build_generator( [model], Namespace(**gen_args) ) if self.cfg.scst: scst_args = json.loads(self.cfg.scst_args) self.scst_generator = self.build_generator( [model], Namespace(**scst_args) ) return model def _calculate_ap_score(self, hyps, refs, thresh=0.5, min_area_size=None, max_area_size=None): 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) if max_area_size is not None and min_area_size is not None: ious = ious * ((area_targets < max_area_size).float() + (area_targets > min_area_size).float())/2 elif min_area_size is not None: ious = ious * (area_targets > min_area_size).float() elif max_area_size is not None: ious = ious * (area_targets < max_area_size).float() if thresh is None: return ious else: return ((ious >= thresh) & (interacts_w > 0) & (interacts_h > 0)).float() def _calculate_map_score(self, hyps, refs, thresh=0.5): ground_truth = { 'boxes': refs.cpu().numpy().tolist(), 'labels': [1 for i in range(refs.shape[0])] } result_dict = { 'boxes': hyps.cpu().numpy().tolist(), 'labels': [1 for i in range(hyps.shape[0])], } score = calculate_map(ground_truth, result_dict, thresh) score = torch.tensor(score).unsqueeze(0).repeat(refs.shape[0]).to(hyps.device) return score 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(self.sequence_generator, sample, model) hyps = hyps / (self.cfg.num_bins - 1) * self.cfg.max_image_size refs = refs / (self.cfg.num_bins - 1) * self.cfg.max_image_size hyps[:, ::2] /= sample['w_resize_ratios'].unsqueeze(1) hyps[:, 1::2] /= sample['h_resize_ratios'].unsqueeze(1) refs[:, ::2] /= sample['w_resize_ratios'].unsqueeze(1) refs[:, 1::2] /= sample['h_resize_ratios'].unsqueeze(1) # scores = self._calculate_ap_score(hyps, refs) # scores = self._calculate_ap_score(hyps, sample['region_coords'].float()) # scores = self._calculate_ap_score(hyps, refs) scores = self._calculate_ap_score(hyps, sample['region_coords'].float(), thresh=self.acc_thresh) if self.min_area_size is not None: large_scores = self._calculate_ap_score(hyps, sample['region_coords'].float(), thresh=self.acc_thresh, min_area_size=self.min_area_size) logging_output["_large_score_sum"] = large_scores.sum().item() logging_output["_large_score_cnt"] = large_scores.size(0) if self.max_area_size is not None: small_scores = self._calculate_ap_score(hyps, sample['region_coords'].float(), thresh=self.acc_thresh, max_area_size=self.max_area_size) logging_output["_small_score_sum"] = small_scores.sum().item() logging_output["_small_score_cnt"] = small_scores.size(0) if self.metric == 'map': map_scores = self._calculate_map_score(hyps, sample['region_coords'].float(), thresh=self.acc_thresh) logging_output["_map_score_sum"] = map_scores.sum().item() logging_output["_map_score_cnt"] = map_scores.size(0) 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, prefix='_score'): score = meters[prefix+"_sum"].sum / meters[prefix+"_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) if self.metric == 'map': metrics.log_scalar("_map_score_sum", sum_logs("_map_score_sum")) metrics.log_scalar("_map_score_cnt", sum_logs("_map_score_cnt")) metrics.log_derived("score", partial(compute_score, prefix='_map_score')) if self.min_area_size is not None: metrics.log_scalar("_large_score_sum", sum_logs("_large_score_sum")) metrics.log_scalar("_large_score_cnt", sum_logs("_large_score_cnt")) metrics.log_derived("score", partial(compute_score, prefix='_large_score')) if self.max_area_size is not None: metrics.log_scalar("_small_score_sum", sum_logs("_small_score_sum")) metrics.log_scalar("_small_score_cnt", sum_logs("_small_score_cnt")) metrics.log_derived("score", partial(compute_score, prefix='_small_score')) def _inference(self, generator, sample, model): gen_out = self.inference_step(generator, [model], sample) hyps, refs = [], [] for i in range(len(gen_out)): hyps.append(gen_out[i][0]["tokens"][:-1] - len(self.src_dict) + self.cfg.num_bins) refs.append(sample["target"][i][:-1] - len(self.src_dict) + self.cfg.num_bins) if self.cfg.eval_print_samples: logger.info("example hypothesis: ", hyps[0]) logger.info("example reference: ", refs[0]) return torch.stack(hyps, dim=0), torch.stack(refs, dim=0)