from mmseg.datasets.builder import DATASETS from mmseg.datasets.custom import CustomDataset from .geospatial_pipelines import LoadGeospatialAnnotations import mmcv import numpy as np from mmcv.utils import print_log from prettytable import PrettyTable from collections import OrderedDict from mmseg.core import eval_metrics, intersect_and_union, pre_eval_to_metrics from collections import Counter import torch @DATASETS.register_module() class GeospatialDataset(CustomDataset): """GeospatialDataset dataset. """ def __init__(self, CLASSES=(0, 1), PALETTE=None, **kwargs): self.CLASSES = CLASSES self.PALETTE = PALETTE gt_seg_map_loader_cfg = kwargs.pop('gt_seg_map_loader_cfg') if 'gt_seg_map_loader_cfg' in kwargs else dict() reduce_zero_label = kwargs.pop('reduce_zero_label') if 'reduce_zero_label' in kwargs else False super(GeospatialDataset, self).__init__( reduce_zero_label=reduce_zero_label, **kwargs) self.gt_seg_map_loader = LoadGeospatialAnnotations(reduce_zero_label=reduce_zero_label, **gt_seg_map_loader_cfg) def pre_eval(self, preds, indices): """Collect eval result from each iteration. Args: preds (list[torch.Tensor] | torch.Tensor): the segmentation logit after argmax, shape (N, H, W). indices (list[int] | int): the prediction related ground truth indices. Returns: list[torch.Tensor]: (area_intersect, area_union, area_prediction, area_ground_truth). """ # In order to compat with batch inference if not isinstance(indices, list): indices = [indices] if not isinstance(preds, list): preds = [preds] crops=len(self.CLASSES) pre_eval_results = [] for pred, index in zip(preds, indices): seg_map = self.get_gt_seg_map_by_idx(index) crop_pre_eval_results=[] for i in range(crops): crop_pred=pred[i*2:(i+1)*2].argmax(axis=0) if len(seg_map.shape)==3: crop_seg_map=seg_map[i] else: crop_seg_map=seg_map crop_pre_eval_results.append( intersect_and_union( crop_pred, crop_seg_map, len(self.CLASSES[i]), self.ignore_index, # as the label map has already been applied and zero label # has already been reduced by get_gt_seg_map_by_idx() i.e. # LoadAnnotations.__call__(), these operations should not # be duplicated. See the following issues/PRs: # https://github.com/open-mmlab/mmsegmentation/issues/1415 # https://github.com/open-mmlab/mmsegmentation/pull/1417 # https://github.com/open-mmlab/mmsegmentation/pull/2504 # for more details label_map=dict(), reduce_zero_label=False)) pre_eval_results.append(crop_pre_eval_results) return pre_eval_results def evaluate(self, results, metric='mIoU', logger=None, gt_seg_maps=None, **kwargs): """Evaluate the dataset. Args: results (list[tuple[torch.Tensor]] | list[str]): per image pre_eval results or predict segmentation map for computing evaluation metric. metric (str | list[str]): Metrics to be evaluated. 'mIoU', 'mDice' and 'mFscore' are supported. logger (logging.Logger | None | str): Logger used for printing related information during evaluation. Default: None. gt_seg_maps (generator[ndarray]): Custom gt seg maps as input, used in ConcatDataset Returns: dict[str, float]: Default metrics. """ if isinstance(metric, str): metric = [metric] allowed_metrics = ['mIoU', 'mDice', 'mFscore'] if not set(metric).issubset(set(allowed_metrics)): raise KeyError('metric {} is not supported'.format(metric)) crop_classes=self.CLASSES crop_types=len(crop_classes) crop_eval_results=[] for i in range(crop_types): eval_results = {} crop_class=crop_classes[i] crop_results=results[i::crop_types] crop_results_flatten=[] [crop_results_flatten.extend(ele) for ele in crop_results] print("single crop result flatten:",len(crop_results_flatten)) # test a list of files # if mmcv.is_list_of(results, np.ndarray) or mmcv.is_list_of(results, str): if mmcv.is_list_of(crop_results_flatten, np.ndarray) or mmcv.is_list_of(crop_results_flatten, str): # print("results:",results[i]) if gt_seg_maps is None: gt_seg_maps = self.get_gt_seg_maps() print("gt_seg_maps:",gt_seg_maps) # num_classes = len(self.CLASSES) num_classes=len(crop_class) single_crop_results=results[i] crop_seg_maps=gt_seg_maps ret_metrics = eval_metrics( single_crop_results, crop_seg_maps, num_classes, self.ignore_index, metric, label_map=dict(), reduce_zero_label=False) else: single_crop_results=crop_results_flatten ret_metrics=pre_eval_to_metrics(single_crop_results,metric) if crop_class is None: class_names = tuple(range(num_classes)) else: # class_names = self.CLASSES class_names=crop_class # summary table ret_metrics_summary = OrderedDict({ ret_metric: np.round(np.nanmean(ret_metric_value) * 100, 2) for ret_metric, ret_metric_value in ret_metrics.items() }) # each class table ret_metrics.pop('aAcc', None) ret_metrics_class = OrderedDict({ ret_metric: np.round(ret_metric_value * 100, 2) for ret_metric, ret_metric_value in ret_metrics.items() }) ret_metrics_class.update({'Class': class_names}) ret_metrics_class.move_to_end('Class', last=False) # for logger class_table_data = PrettyTable() for key, val in ret_metrics_class.items(): class_table_data.add_column(key, val) summary_table_data = PrettyTable() for key, val in ret_metrics_summary.items(): if key == 'aAcc': summary_table_data.add_column(key, [val]) else: summary_table_data.add_column('m' + key, [val]) print_log('per class results:', logger) print_log('\n' + class_table_data.get_string(), logger=logger) print_log('Summary:', logger) print_log('\n' + summary_table_data.get_string(), logger=logger) # each metric dict for key, value in ret_metrics_summary.items(): if key == 'aAcc': eval_results[key] = value / 100.0 else: eval_results['m' + key] = value / 100.0 ret_metrics_class.pop('Class', None) for key, value in ret_metrics_class.items(): eval_results.update({ key + '.' + str(name): value[idx] / 100.0 for idx, name in enumerate(class_names) }) # return eval_results crop_eval_results.append(eval_results) merged_dict={} for i in range(crop_types): merged_dict = Counter(merged_dict) + Counter(crop_eval_results[i]) average_values = {k: v / crop_types for k, v in merged_dict.items()} print("average crop metric values:",average_values) return average_values