| | |
| | import itertools |
| | import json |
| | import logging |
| | import numpy as np |
| | import os |
| | from collections import OrderedDict |
| | import PIL.Image as Image |
| | import pycocotools.mask as mask_util |
| | import torch |
| |
|
| | from detectron2.data import DatasetCatalog, MetadataCatalog |
| | from detectron2.utils.comm import all_gather, is_main_process |
| | from detectron2.utils.file_io import PathManager |
| | from detectron2.evaluation.evaluator import DatasetEvaluator |
| | from detectron2.utils.comm import synchronize |
| |
|
| | from ..semseg_loader import load_semseg |
| |
|
| |
|
| | class SemSegEvaluator(DatasetEvaluator): |
| | """ |
| | Evaluate semantic segmentation metrics. |
| | """ |
| |
|
| | def __init__( |
| | self, |
| | dataset_name, |
| | distributed=True, |
| | output_dir=None, |
| | *, |
| | num_classes=None, |
| | ignore_label=None, |
| | ): |
| | """ |
| | Args: |
| | dataset_name (str): name of the dataset to be evaluated. |
| | distributed (bool): if True, will collect results from all ranks for evaluation. |
| | Otherwise, will evaluate the results in the current process. |
| | output_dir (str): an output directory to dump results. |
| | num_classes, ignore_label: deprecated argument |
| | """ |
| | self._logger = logging.getLogger(__name__) |
| | if num_classes is not None: |
| | self._logger.warn( |
| | "SemSegEvaluator(num_classes) is deprecated! It should be obtained from metadata." |
| | ) |
| | if ignore_label is not None: |
| | self._logger.warn( |
| | "SemSegEvaluator(ignore_label) is deprecated! It should be obtained from metadata." |
| | ) |
| | self._dataset_name = dataset_name |
| | self._distributed = distributed |
| | self._output_dir = output_dir |
| |
|
| | self._cpu_device = torch.device("cpu") |
| |
|
| | self.input_file_to_gt_file = { |
| | dataset_record["file_name"]: dataset_record["sem_seg_file_name"] |
| | for dataset_record in DatasetCatalog.get(dataset_name) |
| | } |
| |
|
| | meta = MetadataCatalog.get(dataset_name) |
| | |
| | try: |
| | c2d = meta.stuff_dataset_id_to_contiguous_id |
| | self._contiguous_id_to_dataset_id = {v: k for k, v in c2d.items()} |
| | except AttributeError: |
| | self._contiguous_id_to_dataset_id = None |
| | self._class_names = meta.stuff_classes |
| | self._class_offset = meta.class_offset if hasattr(meta, 'class_offset') else 0 |
| | self._num_classes = len(meta.stuff_classes) |
| | self._semseg_loader = meta.semseg_loader if hasattr(meta, 'semseg_loader') else 'PIL' |
| |
|
| | if num_classes is not None: |
| | assert self._num_classes == num_classes, f"{self._num_classes} != {num_classes}" |
| | self._ignore_label = ignore_label if ignore_label is not None else meta.ignore_label |
| |
|
| | def reset(self): |
| | self._conf_matrix = np.zeros((self._num_classes + 1, self._num_classes + 1), dtype=np.int64) |
| | self._predictions = [] |
| |
|
| | def process(self, inputs, outputs): |
| | """ |
| | Args: |
| | inputs: the inputs to a model. |
| | It is a list of dicts. Each dict corresponds to an image and |
| | contains keys like "height", "width", "file_name". |
| | outputs: the outputs of a model. It is either list of semantic segmentation predictions |
| | (Tensor [H, W]) or list of dicts with key "sem_seg" that contains semantic |
| | segmentation prediction in the same format. |
| | """ |
| | for input, output in zip(inputs, outputs): |
| | output = output["sem_seg"].argmax(dim=0).to(self._cpu_device) |
| | pred = np.array(output, dtype=np.int) |
| | |
| | with PathManager.open(self.input_file_to_gt_file[input["file_name"]], "rb") as f: |
| | gt = load_semseg(f, self._semseg_loader) - self._class_offset |
| | |
| | if isinstance(self._ignore_label, int): |
| | ignore_label = self._ignore_label - self._class_offset |
| | gt[gt == self._ignore_label] = self._num_classes |
| | elif isinstance(self._ignore_label, list): |
| | for ignore_label in self._ignore_label: |
| | ignore_label = ignore_label - self._class_offset |
| | gt[gt == ignore_label] = self._num_classes |
| | |
| | self._conf_matrix += np.bincount( |
| | (self._num_classes + 1) * pred.reshape(-1) + gt.reshape(-1), |
| | minlength=self._conf_matrix.size, |
| | ).reshape(self._conf_matrix.shape) |
| | |
| | self._predictions.extend(self.encode_json_sem_seg(pred, input["file_name"])) |
| |
|
| | def evaluate(self): |
| | """ |
| | Evaluates standard semantic segmentation metrics (http://cocodataset.org/#stuff-eval): |
| | |
| | * Mean intersection-over-union averaged across classes (mIoU) |
| | * Frequency Weighted IoU (fwIoU) |
| | * Mean pixel accuracy averaged across classes (mACC) |
| | * Pixel Accuracy (pACC) |
| | """ |
| | if self._distributed: |
| | synchronize() |
| | conf_matrix_list = all_gather(self._conf_matrix) |
| | self._predictions = all_gather(self._predictions) |
| | self._predictions = list(itertools.chain(*self._predictions)) |
| | if not is_main_process(): |
| | return |
| | self._conf_matrix = np.zeros_like(self._conf_matrix) |
| | for conf_matrix in conf_matrix_list: |
| | self._conf_matrix += conf_matrix |
| |
|
| | if self._output_dir: |
| | PathManager.mkdirs(self._output_dir) |
| | file_path = os.path.join(self._output_dir, "sem_seg_predictions.json") |
| | with PathManager.open(file_path, "w") as f: |
| | f.write(json.dumps(self._predictions)) |
| |
|
| | acc = np.full(self._num_classes, np.nan, dtype=np.float) |
| | iou = np.full(self._num_classes, np.nan, dtype=np.float) |
| | tp = self._conf_matrix.diagonal()[:-1].astype(np.float) |
| | pos_gt = np.sum(self._conf_matrix[:-1, :-1], axis=0).astype(np.float) |
| | class_weights = pos_gt / np.sum(pos_gt) |
| | pos_pred = np.sum(self._conf_matrix[:-1, :-1], axis=1).astype(np.float) |
| | acc_valid = pos_gt > 0 |
| | acc[acc_valid] = tp[acc_valid] / pos_gt[acc_valid] |
| | iou_valid = (pos_gt + pos_pred) > 0 |
| | union = pos_gt + pos_pred - tp |
| | iou[acc_valid] = tp[acc_valid] / union[acc_valid] |
| | macc = np.sum(acc[acc_valid]) / np.sum(acc_valid) |
| | miou = np.sum(iou[acc_valid]) / np.sum(iou_valid) |
| | fiou = np.sum(iou[acc_valid] * class_weights[acc_valid]) |
| | pacc = np.sum(tp) / np.sum(pos_gt) |
| |
|
| | res = {} |
| | res["mIoU"] = 100 * miou |
| | res["fwIoU"] = 100 * fiou |
| | for i, name in enumerate(self._class_names): |
| | res["IoU-{}".format(name)] = 100 * iou[i] |
| | res["mACC"] = 100 * macc |
| | res["pACC"] = 100 * pacc |
| | for i, name in enumerate(self._class_names): |
| | res["ACC-{}".format(name)] = 100 * acc[i] |
| |
|
| | if self._output_dir: |
| | file_path = os.path.join(self._output_dir, "sem_seg_evaluation.pth") |
| | with PathManager.open(file_path, "wb") as f: |
| | torch.save(res, f) |
| | results = OrderedDict({"sem_seg": res}) |
| | self._logger.info(results) |
| | return results |
| |
|
| | def encode_json_sem_seg(self, sem_seg, input_file_name): |
| | """ |
| | Convert semantic segmentation to COCO stuff format with segments encoded as RLEs. |
| | See http://cocodataset.org/#format-results |
| | """ |
| | json_list = [] |
| | for label in np.unique(sem_seg): |
| | if self._contiguous_id_to_dataset_id is not None: |
| | assert ( |
| | label in self._contiguous_id_to_dataset_id |
| | ), "Label {} is not in the metadata info for {}".format(label, self._dataset_name) |
| | dataset_id = self._contiguous_id_to_dataset_id[label] |
| | else: |
| | dataset_id = int(label) |
| | mask = (sem_seg == label).astype(np.uint8) |
| | mask_rle = mask_util.encode(np.array(mask[:, :, None], order="F"))[0] |
| | mask_rle["counts"] = mask_rle["counts"].decode("utf-8") |
| | json_list.append( |
| | {"file_name": input_file_name, "category_id": dataset_id, "segmentation": mask_rle} |
| | ) |
| | return json_list |
| |
|
