mmdet/evaluation/metrics/coco_panoptic_metric.py

# Copyright (c) OpenMMLab. All rights reserved.
import datetime
import itertools
import os.path as osp
import tempfile
from typing import Dict, Optional, Sequence, Tuple, Union

import mmcv
import numpy as np
from mmengine.evaluator import BaseMetric
from mmengine.fileio import FileClient, dump, load
from mmengine.logging import MMLogger, print_log
from terminaltables import AsciiTable

from mmdet.datasets.api_wrappers import COCOPanoptic
from mmdet.registry import METRICS
from ..functional import (INSTANCE_OFFSET, pq_compute_multi_core,
                          pq_compute_single_core)

try:
    import panopticapi
    from panopticapi.evaluation import VOID, PQStat
    from panopticapi.utils import id2rgb, rgb2id
except ImportError:
    panopticapi = None
    id2rgb = None
    rgb2id = None
    VOID = None
    PQStat = None


@METRICS.register_module()
class CocoPanopticMetric(BaseMetric):
    """COCO panoptic segmentation evaluation metric.

    Evaluate PQ, SQ RQ for panoptic segmentation tasks. Please refer to
    https://cocodataset.org/#panoptic-eval for more details.

    Args:
        ann_file (str, optional): Path to the coco format annotation file.
            If not specified, ground truth annotations from the dataset will
            be converted to coco format. Defaults to None.
        seg_prefix (str, optional): Path to the directory which contains the
            coco panoptic segmentation mask. It should be specified when
            evaluate. Defaults to None.
        classwise (bool): Whether to evaluate the metric class-wise.
            Defaults to False.
        outfile_prefix (str, optional): 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.
            It should be specified when format_only is True. Defaults to None.
        format_only (bool): Format the output results without perform
            evaluation. It is useful when you want to format the result
            to a specific format and submit it to the test server.
            Defaults to False.
        nproc (int): Number of processes for panoptic quality computing.
            Defaults to 32. When ``nproc`` exceeds the number of cpu cores,
            the number of cpu cores is used.
        file_client_args (dict): Arguments to instantiate a FileClient.
            See :class:`mmengine.fileio.FileClient` for details.
            Defaults to ``dict(backend='disk')``.
        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. Defaults to None.
    """
    default_prefix: Optional[str] = 'coco_panoptic'

    def __init__(self,
                 ann_file: Optional[str] = None,
                 seg_prefix: Optional[str] = None,
                 classwise: bool = False,
                 format_only: bool = False,
                 outfile_prefix: Optional[str] = None,
                 nproc: int = 32,
                 file_client_args: dict = dict(backend='disk'),
                 collect_device: str = 'cpu',
                 prefix: Optional[str] = None) -> None:
        if panopticapi is None:
            raise RuntimeError(
                'panopticapi is not installed, please install it by: '
                'pip install git+https://github.com/cocodataset/'
                'panopticapi.git.')

        super().__init__(collect_device=collect_device, prefix=prefix)
        self.classwise = classwise
        self.format_only = format_only
        if self.format_only:
            assert outfile_prefix is not None, 'outfile_prefix must be not'
            'None when format_only is True, otherwise the result files will'
            'be saved to a temp directory which will be cleaned up at the end.'

        self.tmp_dir = None
        # outfile_prefix should be a prefix of a path which points to a shared
        # storage when train or test with multi nodes.
        self.outfile_prefix = outfile_prefix
        if outfile_prefix is None:
            self.tmp_dir = tempfile.TemporaryDirectory()
            self.outfile_prefix = osp.join(self.tmp_dir.name, 'results')
        # the directory to save predicted panoptic segmentation mask
        self.seg_out_dir = f'{self.outfile_prefix}.panoptic'
        self.nproc = nproc
        self.seg_prefix = seg_prefix

        self.cat_ids = None
        self.cat2label = None

        self.file_client_args = file_client_args
        self.file_client = FileClient(**file_client_args)

        if ann_file:
            with self.file_client.get_local_path(ann_file) as local_path:
                self._coco_api = COCOPanoptic(local_path)
            self.categories = self._coco_api.cats
        else:
            self._coco_api = None
            self.categories = None

        self.file_client = FileClient(**file_client_args)

    def __del__(self) -> None:
        """Clean up."""
        if self.tmp_dir is not None:
            self.tmp_dir.cleanup()

    def gt_to_coco_json(self, gt_dicts: Sequence[dict],
                        outfile_prefix: str) -> Tuple[str, str]:
        """Convert ground truth to coco panoptic segmentation format json file.

        Args:
            gt_dicts (Sequence[dict]): Ground truth of the dataset.
            outfile_prefix (str): The filename prefix of the json file. If the
                prefix is "somepath/xxx", the json file will be named
                "somepath/xxx.gt.json".

        Returns:
            Tuple[str, str]: The filename of the json file and the name of the\
                directory which contains panoptic segmentation masks.
        """
        assert len(gt_dicts) > 0, 'gt_dicts is empty.'
        gt_folder = osp.dirname(gt_dicts[0]['seg_map_path'])
        converted_json_path = f'{outfile_prefix}.gt.json'

        categories = []
        for id, name in enumerate(self.dataset_meta['classes']):
            isthing = 1 if name in self.dataset_meta['thing_classes'] else 0
            categories.append({'id': id, 'name': name, 'isthing': isthing})

        image_infos = []
        annotations = []
        for gt_dict in gt_dicts:
            img_id = gt_dict['image_id']
            image_info = {
                'id': img_id,
                'width': gt_dict['width'],
                'height': gt_dict['height'],
                'file_name': osp.split(gt_dict['seg_map_path'])[-1]
            }
            image_infos.append(image_info)

            pan_png = mmcv.imread(gt_dict['seg_map_path']).squeeze()
            pan_png = pan_png[:, :, ::-1]
            pan_png = rgb2id(pan_png)
            segments_info = []
            for segment_info in gt_dict['segments_info']:
                id = segment_info['id']
                label = segment_info['category']
                mask = pan_png == id
                isthing = categories[label]['isthing']
                if isthing:
                    iscrowd = 1 if not segment_info['is_thing'] else 0
                else:
                    iscrowd = 0

                new_segment_info = {
                    'id': id,
                    'category_id': label,
                    'isthing': isthing,
                    'iscrowd': iscrowd,
                    'area': mask.sum()
                }
                segments_info.append(new_segment_info)

            segm_file = image_info['file_name'].replace('jpg', 'png')
            annotation = dict(
                image_id=img_id,
                segments_info=segments_info,
                file_name=segm_file)
            annotations.append(annotation)
            pan_png = id2rgb(pan_png)

        info = dict(
            date_created=str(datetime.datetime.now()),
            description='Coco json file converted by mmdet CocoPanopticMetric.'
        )
        coco_json = dict(
            info=info,
            images=image_infos,
            categories=categories,
            licenses=None,
        )
        if len(annotations) > 0:
            coco_json['annotations'] = annotations
        dump(coco_json, converted_json_path)
        return converted_json_path, gt_folder

    def result2json(self, results: Sequence[dict],
                    outfile_prefix: str) -> Tuple[str, str]:
        """Dump the panoptic results to a COCO style json file and a directory.

        Args:
            results (Sequence[dict]): Testing results of the dataset.
            outfile_prefix (str): The filename prefix of the json files and the
                directory.

        Returns:
            Tuple[str, str]: The json file and the directory which contains \
                panoptic segmentation masks. The filename of the json is
                "somepath/xxx.panoptic.json" and name of the directory is
                "somepath/xxx.panoptic".
        """
        label2cat = dict((v, k) for (k, v) in self.cat2label.items())
        pred_annotations = []
        for idx in range(len(results)):
            result = results[idx]
            for segment_info in result['segments_info']:
                sem_label = segment_info['category_id']
                # convert sem_label to json label
                cat_id = label2cat[sem_label]
                segment_info['category_id'] = label2cat[sem_label]
                is_thing = self.categories[cat_id]['isthing']
                segment_info['isthing'] = is_thing
            pred_annotations.append(result)
        pan_json_results = dict(annotations=pred_annotations)
        json_filename = f'{outfile_prefix}.panoptic.json'
        dump(pan_json_results, json_filename)
        return json_filename, (
            self.seg_out_dir
            if self.tmp_dir is None else tempfile.gettempdir())

    def _parse_predictions(self,
                           pred: dict,
                           img_id: int,
                           segm_file: str,
                           label2cat=None) -> dict:
        """Parse panoptic segmentation predictions.

        Args:
            pred (dict): Panoptic segmentation predictions.
            img_id (int): Image id.
            segm_file (str): Segmentation file name.
            label2cat (dict): Mapping from label to category id.
                Defaults to None.

        Returns:
            dict: Parsed predictions.
        """
        result = dict()
        result['img_id'] = img_id
        # shape (1, H, W) -> (H, W)
        pan = pred['pred_panoptic_seg']['sem_seg'].cpu().numpy()[0]
        pan_labels = np.unique(pan)
        segments_info = []
        for pan_label in pan_labels:
            sem_label = pan_label % INSTANCE_OFFSET
            # We reserve the length of dataset_meta['classes'] for VOID label
            if sem_label == len(self.dataset_meta['classes']):
                continue
            mask = pan == pan_label
            area = mask.sum()
            segments_info.append({
                'id':
                int(pan_label),
                # when ann_file provided, sem_label should be cat_id, otherwise
                # sem_label should be a continuous id, not the cat_id
                # defined in dataset
                'category_id':
                label2cat[sem_label] if label2cat else sem_label,
                'area':
                int(area)
            })
        # evaluation script uses 0 for VOID label.
        pan[pan % INSTANCE_OFFSET == len(self.dataset_meta['classes'])] = VOID
        pan = id2rgb(pan).astype(np.uint8)
        mmcv.imwrite(pan[:, :, ::-1], osp.join(self.seg_out_dir, segm_file))
        result = {
            'image_id': img_id,
            'segments_info': segments_info,
            'file_name': segm_file
        }

        return result

    def _compute_batch_pq_stats(self, data_samples: Sequence[dict]):
        """Process gts and predictions when ``outfile_prefix`` is not set, gts
        are from dataset or a json file which is defined by ``ann_file``.

        Intermediate results, ``pq_stats``, are computed here and put into
        ``self.results``.
        """
        if self._coco_api is None:
            categories = dict()
            for id, name in enumerate(self.dataset_meta['classes']):
                isthing = 1 if name in self.dataset_meta['thing_classes']\
                    else 0
                categories[id] = {'id': id, 'name': name, 'isthing': isthing}
            label2cat = None
        else:
            categories = self.categories
            cat_ids = self._coco_api.get_cat_ids(
                cat_names=self.dataset_meta['classes'])
            label2cat = {i: cat_id for i, cat_id in enumerate(cat_ids)}

        for data_sample in data_samples:
            # parse pred
            img_id = data_sample['img_id']
            segm_file = osp.basename(data_sample['img_path']).replace(
                'jpg', 'png')
            result = self._parse_predictions(
                pred=data_sample,
                img_id=img_id,
                segm_file=segm_file,
                label2cat=label2cat)

            # parse gt
            gt = dict()
            gt['image_id'] = img_id
            gt['width'] = data_sample['ori_shape'][1]
            gt['height'] = data_sample['ori_shape'][0]
            gt['file_name'] = segm_file

            if self._coco_api is None:
                # get segments_info from data_sample
                seg_map_path = osp.join(self.seg_prefix, segm_file)
                pan_png = mmcv.imread(seg_map_path).squeeze()
                pan_png = pan_png[:, :, ::-1]
                pan_png = rgb2id(pan_png)
                segments_info = []

                for segment_info in data_sample['segments_info']:
                    id = segment_info['id']
                    label = segment_info['category']
                    mask = pan_png == id
                    isthing = categories[label]['isthing']
                    if isthing:
                        iscrowd = 1 if not segment_info['is_thing'] else 0
                    else:
                        iscrowd = 0

                    new_segment_info = {
                        'id': id,
                        'category_id': label,
                        'isthing': isthing,
                        'iscrowd': iscrowd,
                        'area': mask.sum()
                    }
                    segments_info.append(new_segment_info)
            else:
                # get segments_info from annotation file
                segments_info = self._coco_api.imgToAnns[img_id]

            gt['segments_info'] = segments_info

            pq_stats = pq_compute_single_core(
                proc_id=0,
                annotation_set=[(gt, result)],
                gt_folder=self.seg_prefix,
                pred_folder=self.seg_out_dir,
                categories=categories,
                file_client=self.file_client)

            self.results.append(pq_stats)

    def _process_gt_and_predictions(self, data_samples: Sequence[dict]):
        """Process gts and predictions when ``outfile_prefix`` is set.

        The predictions will be saved to directory specified by
        ``outfile_predfix``. The matched pair (gt, result) will be put into
        ``self.results``.
        """
        for data_sample in data_samples:
            # parse pred
            img_id = data_sample['img_id']
            segm_file = osp.basename(data_sample['img_path']).replace(
                'jpg', 'png')
            result = self._parse_predictions(
                pred=data_sample, img_id=img_id, segm_file=segm_file)

            # parse gt
            gt = dict()
            gt['image_id'] = img_id
            gt['width'] = data_sample['ori_shape'][1]
            gt['height'] = data_sample['ori_shape'][0]

            if self._coco_api is None:
                # get segments_info from dataset
                gt['segments_info'] = data_sample['segments_info']
                gt['seg_map_path'] = data_sample['seg_map_path']

            self.results.append((gt, result))

    # TODO: data_batch is no longer needed, consider adjusting the
    #  parameter position
    def process(self, data_batch: dict, data_samples: Sequence[dict]) -> None:
        """Process one batch of data samples and predictions. The processed
        results should be stored in ``self.results``, which will be used to
        compute the metrics when all batches have been processed.

        Args:
            data_batch (dict): A batch of data from the dataloader.
            data_samples (Sequence[dict]): A batch of data samples that
                contain annotations and predictions.
        """
        # If ``self.tmp_dir`` is none, it will save gt and predictions to
        # self.results, otherwise, it will compute pq_stats here.
        if self.tmp_dir is None:
            self._process_gt_and_predictions(data_samples)
        else:
            self._compute_batch_pq_stats(data_samples)

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

        Args:
            results (list): The processed results of each batch. There
                are two cases:

                - When ``outfile_prefix`` is not provided, the elements in
                  results are pq_stats which can be summed directly to get PQ.
                - When ``outfile_prefix`` is provided, the elements in
                  results are tuples like (gt, pred).

        Returns:
            Dict[str, float]: The computed metrics. The keys are the names of
                the metrics, and the values are corresponding results.
        """
        logger: MMLogger = MMLogger.get_current_instance()

        if self.tmp_dir is None:
            # do evaluation after collect all the results

            # split gt and prediction list
            gts, preds = zip(*results)

            if self._coco_api is None:
                # use converted gt json file to initialize coco api
                logger.info('Converting ground truth to coco format...')
                coco_json_path, gt_folder = self.gt_to_coco_json(
                    gt_dicts=gts, outfile_prefix=self.outfile_prefix)
                self._coco_api = COCOPanoptic(coco_json_path)
            else:
                gt_folder = self.seg_prefix

            self.cat_ids = self._coco_api.get_cat_ids(
                cat_names=self.dataset_meta['classes'])
            self.cat2label = {
                cat_id: i
                for i, cat_id in enumerate(self.cat_ids)
            }
            self.img_ids = self._coco_api.get_img_ids()
            self.categories = self._coco_api.cats

            # convert predictions to coco format and dump to json file
            json_filename, pred_folder = self.result2json(
                results=preds, outfile_prefix=self.outfile_prefix)

            if self.format_only:
                logger.info('results are saved in '
                            f'{osp.dirname(self.outfile_prefix)}')
                return dict()

            imgs = self._coco_api.imgs
            gt_json = self._coco_api.img_ann_map
            gt_json = [{
                'image_id': k,
                'segments_info': v,
                'file_name': imgs[k]['segm_file']
            } for k, v in gt_json.items()]
            pred_json = load(json_filename)
            pred_json = dict(
                (el['image_id'], el) for el in pred_json['annotations'])

            # match the gt_anns and pred_anns in the same image
            matched_annotations_list = []
            for gt_ann in gt_json:
                img_id = gt_ann['image_id']
                if img_id not in pred_json.keys():
                    raise Exception('no prediction for the image'
                                    ' with id: {}'.format(img_id))
                matched_annotations_list.append((gt_ann, pred_json[img_id]))

            pq_stat = pq_compute_multi_core(
                matched_annotations_list,
                gt_folder,
                pred_folder,
                self.categories,
                file_client=self.file_client,
                nproc=self.nproc)

        else:
            # aggregate the results generated in process
            if self._coco_api is None:
                categories = dict()
                for id, name in enumerate(self.dataset_meta['classes']):
                    isthing = 1 if name in self.dataset_meta[
                        'thing_classes'] else 0
                    categories[id] = {
                        'id': id,
                        'name': name,
                        'isthing': isthing
                    }
                self.categories = categories

            pq_stat = PQStat()
            for result in results:
                pq_stat += result

        metrics = [('All', None), ('Things', True), ('Stuff', False)]
        pq_results = {}

        for name, isthing in metrics:
            pq_results[name], classwise_results = pq_stat.pq_average(
                self.categories, isthing=isthing)
            if name == 'All':
                pq_results['classwise'] = classwise_results

        classwise_results = None
        if self.classwise:
            classwise_results = {
                k: v
                for k, v in zip(self.dataset_meta['classes'],
                                pq_results['classwise'].values())
            }

        print_panoptic_table(pq_results, classwise_results, logger=logger)
        results = parse_pq_results(pq_results)

        return results


def parse_pq_results(pq_results: dict) -> dict:
    """Parse the Panoptic Quality results.

    Args:
        pq_results (dict): Panoptic Quality results.

    Returns:
        dict: Panoptic Quality results parsed.
    """
    result = dict()
    result['PQ'] = 100 * pq_results['All']['pq']
    result['SQ'] = 100 * pq_results['All']['sq']
    result['RQ'] = 100 * pq_results['All']['rq']
    result['PQ_th'] = 100 * pq_results['Things']['pq']
    result['SQ_th'] = 100 * pq_results['Things']['sq']
    result['RQ_th'] = 100 * pq_results['Things']['rq']
    result['PQ_st'] = 100 * pq_results['Stuff']['pq']
    result['SQ_st'] = 100 * pq_results['Stuff']['sq']
    result['RQ_st'] = 100 * pq_results['Stuff']['rq']
    return result


def print_panoptic_table(
        pq_results: dict,
        classwise_results: Optional[dict] = None,
        logger: Optional[Union['MMLogger', str]] = None) -> None:
    """Print the panoptic evaluation results table.

    Args:
        pq_results(dict): The Panoptic Quality results.
        classwise_results(dict, optional): The classwise Panoptic Quality.
            results. The keys are class names and the values are metrics.
            Defaults to None.
        logger (:obj:`MMLogger` | str, optional): Logger used for printing
            related information during evaluation. Default: None.
    """

    headers = ['', 'PQ', 'SQ', 'RQ', 'categories']
    data = [headers]
    for name in ['All', 'Things', 'Stuff']:
        numbers = [
            f'{(pq_results[name][k] * 100):0.3f}' for k in ['pq', 'sq', 'rq']
        ]
        row = [name] + numbers + [pq_results[name]['n']]
        data.append(row)
    table = AsciiTable(data)
    print_log('Panoptic Evaluation Results:\n' + table.table, logger=logger)

    if classwise_results is not None:
        class_metrics = [(name, ) + tuple(f'{(metrics[k] * 100):0.3f}'
                                          for k in ['pq', 'sq', 'rq'])
                         for name, metrics in classwise_results.items()]
        num_columns = min(8, len(class_metrics) * 4)
        results_flatten = list(itertools.chain(*class_metrics))
        headers = ['category', 'PQ', 'SQ', 'RQ'] * (num_columns // 4)
        results_2d = itertools.zip_longest(
            *[results_flatten[i::num_columns] for i in range(num_columns)])
        data = [headers]
        data += [result for result in results_2d]
        table = AsciiTable(data)
        print_log(
            'Classwise Panoptic Evaluation Results:\n' + table.table,
            logger=logger)