| import string | |
| import math | |
| import torch | |
| from data import data_utils | |
| def get_symbols_to_strip_from_output(generator): | |
| if hasattr(generator, "symbols_to_strip_from_output"): | |
| return generator.symbols_to_strip_from_output | |
| else: | |
| return {generator.bos, generator.eos} | |
| def decode_fn(x, tgt_dict, bpe, generator, tokenizer=None): | |
| x = tgt_dict.string(x.int().cpu(), extra_symbols_to_ignore=get_symbols_to_strip_from_output(generator)) | |
| if bpe is not None: | |
| x = bpe.decode(x) | |
| if tokenizer is not None: | |
| x = tokenizer.decode(x) | |
| return x | |
| def eval_caption(task, generator, models, sample): | |
| transtab = str.maketrans({key: None for key in string.punctuation}) | |
| hypos = task.inference_step(generator, models, sample) | |
| results = [] | |
| for i, sample_id in enumerate(sample["id"].tolist()): | |
| detok_hypo_str = decode_fn(hypos[i][0]["tokens"], task.tgt_dict, task.bpe, generator) | |
| results.append({"image_id": str(sample_id), "caption": detok_hypo_str.translate(transtab).strip()}) | |
| return results, None | |
| def eval_refcoco(task, generator, models, sample): | |
| def _calculate_ap_score(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() | |
| gen_out = task.inference_step(generator, models, sample) | |
| hyps = [] | |
| for i in range(len(gen_out)): | |
| hyps.append(gen_out[i][0]["tokens"][:-1] - len(task.src_dict) + task.cfg.num_bins) | |
| hyps = torch.stack(hyps, dim=0) | |
| hyps = hyps / (task.cfg.num_bins - 1) * task.cfg.max_image_size | |
| hyps[:, ::2] /= sample['w_resize_ratios'].unsqueeze(1) | |
| hyps[:, 1::2] /= sample['h_resize_ratios'].unsqueeze(1) | |
| results = [ | |
| {"uniq_id": sample_id, | |
| "box": [hyps[i][0].item(), hyps[i][1].item(), hyps[i][2].item(), hyps[i][3].item()]} | |
| for i, sample_id in enumerate(sample["id"].tolist()) | |
| ] | |
| scores = _calculate_ap_score(hyps, sample['region_coords'].float()) | |
| return results, scores | |
| def eval_step(task, generator, models, sample): | |
| if task.cfg._name == 'caption': | |
| return eval_caption(task, generator, models, sample) | |
| elif task.cfg._name == 'refcoco': | |
| return eval_refcoco(task, generator, models, sample) | |
| else: | |
| raise NotImplementedError | |