mmdet/evaluation/metrics/youtube_vis_metric.py

# Copyright (c) OpenMMLab. All rights reserved.
import os.path as osp
import tempfile
import warnings
import zipfile
from collections import OrderedDict, defaultdict
from typing import Dict, List, Optional, Sequence, Tuple, Union

import mmengine
import numpy as np
from mmengine.dist import (all_gather_object, barrier, broadcast_object_list,
                           is_main_process)
from mmengine.logging import MMLogger

from mmdet.registry import METRICS
from mmdet.structures.mask import encode_mask_results
from ..functional import YTVIS, YTVISeval
from .base_video_metric import BaseVideoMetric, collect_tracking_results


@METRICS.register_module()
class YouTubeVISMetric(BaseVideoMetric):
    """mAP evaluation metrics for the VIS task.

    Args:
        metric (str | list[str]): Metrics to be evaluated.
            Default value is `youtube_vis_ap`.
        metric_items (List[str], optional): Metric result names to be
            recorded in the evaluation result. Defaults to None.
        outfile_prefix (str | None): The prefix of json files. It includes
            the file path and the prefix of filename, e.g., "a/b/prefix".
            If not specified, a temp file will be created. Defaults to None.
        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 homonyms metrics of different evaluators.
            If prefix is not provided in the argument, self.default_prefix
            will be used instead. Default: None
        format_only (bool): If True, only formatting the results to the
            official format and not performing evaluation. Defaults to False.
    """

    default_prefix: Optional[str] = 'youtube_vis'

    def __init__(self,
                 metric: Union[str, List[str]] = 'youtube_vis_ap',
                 metric_items: Optional[Sequence[str]] = None,
                 outfile_prefix: Optional[str] = None,
                 collect_device: str = 'cpu',
                 prefix: Optional[str] = None,
                 format_only: bool = False) -> None:
        super().__init__(collect_device=collect_device, prefix=prefix)
        # vis evaluation metrics
        self.metrics = metric if isinstance(metric, list) else [metric]
        self.format_only = format_only
        allowed_metrics = ['youtube_vis_ap']
        for metric in self.metrics:
            if metric not in allowed_metrics:
                raise KeyError(
                    f"metric should be 'youtube_vis_ap', but got {metric}.")

        self.metric_items = metric_items
        self.outfile_prefix = outfile_prefix
        self.per_video_res = []
        self.categories = []
        self._vis_meta_info = defaultdict(list)  # record video and image infos

    def process_video(self, data_samples):

        video_length = len(data_samples)
        for frame_id in range(video_length):
            result = dict()
            img_data_sample = data_samples[frame_id].to_dict()
            pred = img_data_sample['pred_track_instances']
            video_id = img_data_sample['video_id']

            result['img_id'] = img_data_sample['img_id']
            result['bboxes'] = pred['bboxes'].cpu().numpy()
            result['scores'] = pred['scores'].cpu().numpy()
            result['labels'] = pred['labels'].cpu().numpy()
            result['instances_id'] = pred['instances_id'].cpu().numpy()
            # encode mask to RLE
            assert 'masks' in pred, \
                'masks must exist in YouTube-VIS metric'
            result['masks'] = encode_mask_results(
                pred['masks'].detach().cpu().numpy())

            # parse gt
            gt = dict()
            gt['width'] = img_data_sample['ori_shape'][1]
            gt['height'] = img_data_sample['ori_shape'][0]
            gt['img_id'] = img_data_sample['img_id']
            gt['frame_id'] = frame_id
            gt['video_id'] = video_id
            gt['video_length'] = video_length

            if 'instances' in img_data_sample:
                gt['anns'] = img_data_sample['instances']
            else:
                gt['anns'] = dict()
            self.per_video_res.append((result, gt))

        preds, gts = zip(*self.per_video_res)
        # format the results
        # we must format gts first to update self._vis_meta_info
        gt_results = self._format_one_video_gts(gts)
        pred_results = self._format_one_video_preds(preds)
        self.per_video_res.clear()
        # add converted result to the results list
        self.results.append((pred_results, gt_results))

    def compute_metrics(self, results: List) -> Dict[str, float]:
        """Compute the metrics from processed results.

        Args:
            results (List): The processed results of each batch.

        Returns:
            Dict[str, float]: The computed metrics. The keys are the names of
            the metrics, and the values are corresponding results.
        """
        # split gt and prediction list
        tmp_pred_results, tmp_gt_results = zip(*results)
        gt_results = self.format_gts(tmp_gt_results)
        pred_results = self.format_preds(tmp_pred_results)

        if self.format_only:
            self.save_pred_results(pred_results)
            return dict()

        ytvis = YTVIS(gt_results)

        ytvis_dets = ytvis.loadRes(pred_results)
        vid_ids = ytvis.getVidIds()

        iou_type = metric = 'segm'
        eval_results = OrderedDict()
        ytvisEval = YTVISeval(ytvis, ytvis_dets, iou_type)
        ytvisEval.params.vidIds = vid_ids
        ytvisEval.evaluate()
        ytvisEval.accumulate()
        ytvisEval.summarize()

        coco_metric_names = {
            'mAP': 0,
            'mAP_50': 1,
            'mAP_75': 2,
            'mAP_s': 3,
            'mAP_m': 4,
            'mAP_l': 5,
            'AR@1': 6,
            'AR@10': 7,
            'AR@100': 8,
            'AR_s@100': 9,
            'AR_m@100': 10,
            'AR_l@100': 11
        }
        metric_items = self.metric_items
        if metric_items is not None:
            for metric_item in metric_items:
                if metric_item not in coco_metric_names:
                    raise KeyError(
                        f'metric item "{metric_item}" is not supported')

        if metric_items is None:
            metric_items = [
                'mAP', 'mAP_50', 'mAP_75', 'mAP_s', 'mAP_m', 'mAP_l'
            ]
        for metric_item in metric_items:
            key = f'{metric}_{metric_item}'
            val = float(
                f'{ytvisEval.stats[coco_metric_names[metric_item]]:.3f}')
            eval_results[key] = val

        return eval_results

    def format_gts(self, gts: Tuple[List]) -> dict:
        """Gather all ground-truth from self.results."""
        self.categories = [
            dict(id=id + 1, name=name)
            for id, name in enumerate(self.dataset_meta['classes'])
        ]
        gt_results = dict(
            categories=self.categories,
            videos=self._vis_meta_info['videos'],
            annotations=[])
        for gt_result in gts:
            gt_results['annotations'].extend(gt_result)
        return gt_results

    def format_preds(self, preds: Tuple[List]) -> List:
        """Gather all predictions from self.results."""
        pred_results = []
        for pred_result in preds:
            pred_results.extend(pred_result)
        return pred_results

    def _format_one_video_preds(self, pred_dicts: Tuple[dict]) -> List:
        """Convert the annotation to the format of YouTube-VIS.

        This operation is to make it easier to use the official eval API.

        Args:
            pred_dicts (Tuple[dict]): Prediction of the dataset.

        Returns:
            List: The formatted predictions.
        """
        # Collate preds scatters (tuple of dict to dict of list)
        preds = defaultdict(list)
        for pred in pred_dicts:
            for key in pred.keys():
                preds[key].append(pred[key])

        img_infos = self._vis_meta_info['images']
        vid_infos = self._vis_meta_info['videos']
        inds = [i for i, _ in enumerate(img_infos) if _['frame_id'] == 0]
        inds.append(len(img_infos))
        json_results = []
        video_id = vid_infos[-1]['id']
        # collect data for each instances in a video.
        collect_data = dict()
        for frame_id, (masks, scores, labels, ids) in enumerate(
                zip(preds['masks'], preds['scores'], preds['labels'],
                    preds['instances_id'])):

            assert len(masks) == len(labels)
            for j, id in enumerate(ids):
                if id not in collect_data:
                    collect_data[id] = dict(
                        category_ids=[], scores=[], segmentations=dict())
                collect_data[id]['category_ids'].append(labels[j])
                collect_data[id]['scores'].append(scores[j])
                if isinstance(masks[j]['counts'], bytes):
                    masks[j]['counts'] = masks[j]['counts'].decode()
                collect_data[id]['segmentations'][frame_id] = masks[j]

        # transform the collected data into official format
        for id, id_data in collect_data.items():
            output = dict()
            output['video_id'] = video_id
            output['score'] = np.array(id_data['scores']).mean().item()
            # majority voting for sequence category
            output['category_id'] = np.bincount(
                np.array(id_data['category_ids'])).argmax().item() + 1
            output['segmentations'] = []
            for frame_id in range(inds[-1] - inds[-2]):
                if frame_id in id_data['segmentations']:
                    output['segmentations'].append(
                        id_data['segmentations'][frame_id])
                else:
                    output['segmentations'].append(None)
            json_results.append(output)

        return json_results

    def _format_one_video_gts(self, gt_dicts: Tuple[dict]) -> List:
        """Convert the annotation to the format of YouTube-VIS.

        This operation is to make it easier to use the official eval API.

        Args:
            gt_dicts (Tuple[dict]): Ground truth of the dataset.

        Returns:
            list: The formatted gts.
        """
        video_infos = []
        image_infos = []
        instance_infos = defaultdict(list)
        len_videos = dict()  # mapping from instance_id to video_length
        vis_anns = []

        # get video infos
        for gt_dict in gt_dicts:
            frame_id = gt_dict['frame_id']
            video_id = gt_dict['video_id']
            img_id = gt_dict['img_id']
            image_info = dict(
                id=img_id,
                width=gt_dict['width'],
                height=gt_dict['height'],
                frame_id=frame_id,
                file_name='')
            image_infos.append(image_info)
            if frame_id == 0:
                video_info = dict(
                    id=video_id,
                    width=gt_dict['width'],
                    height=gt_dict['height'],
                    file_name='')
                video_infos.append(video_info)

            for ann in gt_dict['anns']:
                label = ann['bbox_label']
                bbox = ann['bbox']
                instance_id = ann['instance_id']
                # update video length
                len_videos[instance_id] = gt_dict['video_length']
                coco_bbox = [
                    bbox[0],
                    bbox[1],
                    bbox[2] - bbox[0],
                    bbox[3] - bbox[1],
                ]

                annotation = dict(
                    video_id=video_id,
                    frame_id=frame_id,
                    bbox=coco_bbox,
                    instance_id=instance_id,
                    iscrowd=ann.get('ignore_flag', 0),
                    category_id=int(label) + 1,
                    area=coco_bbox[2] * coco_bbox[3])
                if ann.get('mask', None):
                    mask = ann['mask']
                    # area = mask_util.area(mask)
                    if isinstance(mask, dict) and isinstance(
                            mask['counts'], bytes):
                        mask['counts'] = mask['counts'].decode()
                    annotation['segmentation'] = mask

                instance_infos[instance_id].append(annotation)

        # update vis meta info
        self._vis_meta_info['images'].extend(image_infos)
        self._vis_meta_info['videos'].extend(video_infos)

        for instance_id, ann_infos in instance_infos.items():
            cur_video_len = len_videos[instance_id]
            segm = [None] * cur_video_len
            bbox = [None] * cur_video_len
            area = [None] * cur_video_len
            # In the official format, no instances are represented by
            # 'None', however, only images with instances are recorded
            # in the current annotations, so we need to use 'None' to
            # initialize these lists.
            for ann_info in ann_infos:
                frame_id = ann_info['frame_id']
                segm[frame_id] = ann_info['segmentation']
                bbox[frame_id] = ann_info['bbox']
                area[frame_id] = ann_info['area']
            instance = dict(
                category_id=ann_infos[0]['category_id'],
                segmentations=segm,
                bboxes=bbox,
                video_id=ann_infos[0]['video_id'],
                areas=area,
                id=instance_id,
                iscrowd=ann_infos[0]['iscrowd'])
            vis_anns.append(instance)
        return vis_anns

    def save_pred_results(self, pred_results: List) -> None:
        """Save the results to a zip file (standard format for YouTube-VIS
        Challenge).

        Args:
            pred_results (list): Testing results of the
                dataset.
        """
        logger: MMLogger = MMLogger.get_current_instance()
        if self.outfile_prefix is None:
            tmp_dir = tempfile.TemporaryDirectory()
            outfile_prefix = osp.join(tmp_dir.name, 'results')
        else:
            outfile_prefix = self.outfile_prefix
        mmengine.dump(pred_results, f'{outfile_prefix}.json')
        # zip the json file in order to submit to the test server.
        zip_file_name = f'{outfile_prefix}.submission_file.zip'
        zf = zipfile.ZipFile(zip_file_name, 'w', zipfile.ZIP_DEFLATED)
        logger.info(f"zip the 'results.json' into '{zip_file_name}', "
                    'please submmit the zip file to the test server')
        zf.write(f'{outfile_prefix}.json', 'results.json')
        zf.close()

    def evaluate(self, size: int) -> dict:
        """Evaluate the model performance of the whole dataset after processing
        all batches.

        Args:
            size (int): Length of the entire validation dataset.

        Returns:
            dict: Evaluation metrics dict on the val dataset. The keys are the
            names of the metrics, and the values are corresponding results.
        """
        # wait for all processes to complete prediction.
        barrier()

        if len(self.results) == 0:
            warnings.warn(
                f'{self.__class__.__name__} got empty `self.results`. Please '
                'ensure that the processed results are properly added into '
                '`self.results` in `process` method.')

        results = collect_tracking_results(self.results, self.collect_device)

        # gather seq_info
        gathered_seq_info = all_gather_object(self._vis_meta_info['videos'])
        all_seq_info = []
        for _seq_info in gathered_seq_info:
            all_seq_info.extend(_seq_info)
        # update self._vis_meta_info
        self._vis_meta_info = dict(videos=all_seq_info)

        if is_main_process():
            _metrics = self.compute_metrics(results)  # type: ignore
            # Add prefix to metric names
            if self.prefix:
                _metrics = {
                    '/'.join((self.prefix, k)): v
                    for k, v in _metrics.items()
                }
            metrics = [_metrics]
        else:
            metrics = [None]  # type: ignore

        broadcast_object_list(metrics)

        # reset the results list
        self.results.clear()
        # reset the vis_meta_info
        self._vis_meta_info.clear()
        return metrics[0]