60 lines
1.9 KiB
Python
60 lines
1.9 KiB
Python
# Copyright (c) OpenMMLab. All rights reserved.
|
|
from typing import List, Optional
|
|
|
|
import mmengine
|
|
|
|
from mmpretrain.registry import METRICS
|
|
from mmpretrain.utils import require
|
|
from .caption import COCOCaption, save_result
|
|
|
|
try:
|
|
from pycocoevalcap.eval import COCOEvalCap
|
|
from pycocotools.coco import COCO
|
|
except ImportError:
|
|
COCOEvalCap = None
|
|
COCO = None
|
|
|
|
|
|
@METRICS.register_module()
|
|
class NocapsSave(COCOCaption):
|
|
"""Nocaps evaluation wrapper.
|
|
|
|
Save the generated captions and transform into coco format.
|
|
The dumped file can be submitted to the official evluation system.
|
|
|
|
Args:
|
|
collect_device (str): Device name used for collecting results from
|
|
different ranks during distributed training. Must be 'cpu' or
|
|
'gpu'. Defaults to 'cpu'.
|
|
prefix (str, optional): The prefix that will be added in the metric
|
|
names to disambiguate homonymous metrics of different evaluators.
|
|
If prefix is not provided in the argument, self.default_prefix
|
|
will be used instead. Should be modified according to the
|
|
`retrieval_type` for unambiguous results. Defaults to TR.
|
|
"""
|
|
|
|
@require('pycocoevalcap')
|
|
def __init__(self,
|
|
save_dir: str = './',
|
|
collect_device: str = 'cpu',
|
|
prefix: Optional[str] = None):
|
|
super(COCOCaption, self).__init__(
|
|
collect_device=collect_device, prefix=prefix)
|
|
self.save_dir = save_dir
|
|
|
|
def compute_metrics(self, results: List):
|
|
"""Compute the metrics from processed results.
|
|
|
|
Args:
|
|
results (dict): The processed results of each batch.
|
|
"""
|
|
mmengine.mkdir_or_exist(self.save_dir)
|
|
save_result(
|
|
result=results,
|
|
result_dir=self.save_dir,
|
|
filename='nocap_pred',
|
|
remove_duplicate='image_id',
|
|
)
|
|
|
|
return dict()
|