Spaces:
Runtime error
Runtime error
File size: 1,888 Bytes
3b96cb1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 |
# Copyright (c) OpenMMLab. All rights reserved.
import warnings
from typing import Sequence
from mmengine.evaluator import DumpResults
from mmengine.evaluator.metric import _to_cpu
from mmdet.registry import METRICS
from mmdet.structures.mask import encode_mask_results
@METRICS.register_module()
class DumpDetResults(DumpResults):
"""Dump model predictions to a pickle file for offline evaluation.
Different from `DumpResults` in MMEngine, it compresses instance
segmentation masks into RLE format.
Args:
out_file_path (str): Path of the dumped file. Must end with '.pkl'
or '.pickle'.
collect_device (str): Device name used for collecting results from
different ranks during distributed training. Must be 'cpu' or
'gpu'. Defaults to 'cpu'.
"""
def process(self, data_batch: dict, data_samples: Sequence[dict]) -> None:
"""transfer tensors in predictions to CPU."""
data_samples = _to_cpu(data_samples)
for data_sample in data_samples:
# remove gt
data_sample.pop('gt_instances', None)
data_sample.pop('ignored_instances', None)
data_sample.pop('gt_panoptic_seg', None)
if 'pred_instances' in data_sample:
pred = data_sample['pred_instances']
# encode mask to RLE
if 'masks' in pred:
pred['masks'] = encode_mask_results(pred['masks'].numpy())
if 'pred_panoptic_seg' in data_sample:
warnings.warn(
'Panoptic segmentation map will not be compressed. '
'The dumped file will be extremely large! '
'Suggest using `CocoPanopticMetric` to save the coco '
'format json and segmentation png files directly.')
self.results.extend(data_samples)
|