# Copyright (c) Facebook, Inc. and its affiliates. import itertools import json import logging import numpy as np import os from collections import OrderedDict from typing import Optional, Union import custom_pycocotools.mask as mask_util import torch from PIL import Image from custom_detectron2.data import DatasetCatalog, MetadataCatalog from custom_detectron2.utils.comm import all_gather, is_main_process, synchronize from custom_detectron2.utils.file_io import PathManager from .evaluator import DatasetEvaluator _CV2_IMPORTED = True try: import cv2 # noqa except ImportError: # OpenCV is an optional dependency at the moment _CV2_IMPORTED = False def load_image_into_numpy_array( filename: str, copy: bool = False, dtype: Optional[Union[np.dtype, str]] = None, ) -> np.ndarray: with PathManager.open(filename, "rb") as f: array = np.array(Image.open(f), copy=copy, dtype=dtype) return array class SemSegEvaluator(DatasetEvaluator): """ Evaluate semantic segmentation metrics. """ def __init__( self, dataset_name, distributed=True, output_dir=None, *, sem_seg_loading_fn=load_image_into_numpy_array, 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. sem_seg_loading_fn: function to read sem seg file and load into numpy array. Default provided, but projects can customize. 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) # Dict that maps contiguous training ids to COCO category ids 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.sem_seg_loading_fn = sem_seg_loading_fn self._num_classes = len(meta.stuff_classes) 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 # This is because cv2.erode did not work for int datatype. Only works for uint8. self._compute_boundary_iou = True if not _CV2_IMPORTED: self._compute_boundary_iou = False self._logger.warn( """Boundary IoU calculation requires OpenCV. B-IoU metrics are not going to be computed because OpenCV is not available to import.""" ) if self._num_classes >= np.iinfo(np.uint8).max: self._compute_boundary_iou = False self._logger.warn( f"""SemSegEvaluator(num_classes) is more than supported value for Boundary IoU calculation! B-IoU metrics are not going to be computed. Max allowed value (exclusive) for num_classes for calculating Boundary IoU is {np.iinfo(np.uint8).max}. The number of classes of dataset {self._dataset_name} is {self._num_classes}""" ) def reset(self): self._conf_matrix = np.zeros((self._num_classes + 1, self._num_classes + 1), dtype=np.int64) self._b_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) gt_filename = self.input_file_to_gt_file[input["file_name"]] gt = self.sem_seg_loading_fn(gt_filename, dtype=np.int) gt[gt == self._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) if self._compute_boundary_iou: b_gt = self._mask_to_boundary(gt.astype(np.uint8)) b_pred = self._mask_to_boundary(pred.astype(np.uint8)) self._b_conf_matrix += np.bincount( (self._num_classes + 1) * b_pred.reshape(-1) + b_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) b_conf_matrix_list = all_gather(self._b_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 self._b_conf_matrix = np.zeros_like(self._b_conf_matrix) for b_conf_matrix in b_conf_matrix_list: self._b_conf_matrix += b_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] union = pos_gt + pos_pred - tp iou_valid = np.logical_and(acc_valid, union > 0) iou[iou_valid] = tp[iou_valid] / union[iou_valid] macc = np.sum(acc[acc_valid]) / np.sum(acc_valid) miou = np.sum(iou[iou_valid]) / np.sum(iou_valid) fiou = np.sum(iou[iou_valid] * class_weights[iou_valid]) pacc = np.sum(tp) / np.sum(pos_gt) if self._compute_boundary_iou: b_iou = np.full(self._num_classes, np.nan, dtype=np.float) b_tp = self._b_conf_matrix.diagonal()[:-1].astype(np.float) b_pos_gt = np.sum(self._b_conf_matrix[:-1, :-1], axis=0).astype(np.float) b_pos_pred = np.sum(self._b_conf_matrix[:-1, :-1], axis=1).astype(np.float) b_union = b_pos_gt + b_pos_pred - b_tp b_iou_valid = b_union > 0 b_iou[b_iou_valid] = b_tp[b_iou_valid] / b_union[b_iou_valid] res = {} res["mIoU"] = 100 * miou res["fwIoU"] = 100 * fiou for i, name in enumerate(self._class_names): res[f"IoU-{name}"] = 100 * iou[i] if self._compute_boundary_iou: res[f"BoundaryIoU-{name}"] = 100 * b_iou[i] res[f"min(IoU, B-Iou)-{name}"] = 100 * min(iou[i], b_iou[i]) res["mACC"] = 100 * macc res["pACC"] = 100 * pacc for i, name in enumerate(self._class_names): res[f"ACC-{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 def _mask_to_boundary(self, mask: np.ndarray, dilation_ratio=0.02): assert mask.ndim == 2, "mask_to_boundary expects a 2-dimensional image" h, w = mask.shape diag_len = np.sqrt(h**2 + w**2) dilation = max(1, int(round(dilation_ratio * diag_len))) kernel = np.ones((3, 3), dtype=np.uint8) padded_mask = cv2.copyMakeBorder(mask, 1, 1, 1, 1, cv2.BORDER_CONSTANT, value=0) eroded_mask_with_padding = cv2.erode(padded_mask, kernel, iterations=dilation) eroded_mask = eroded_mask_with_padding[1:-1, 1:-1] boundary = mask - eroded_mask return boundary