Spaces:
Build error
Build error
| # -------------------------------------------------------- | |
| # X-Decoder -- Generalized Decoding for Pixel, Image, and Language | |
| # Copyright (c) 2022 Microsoft | |
| # Licensed under The MIT License [see LICENSE for details] | |
| # Modified by Xueyan Zou ([email protected]) | |
| # -------------------------------------------------------- | |
| import logging | |
| import torch | |
| from torchvision.ops import box_iou | |
| from detectron2.structures import BoxMode | |
| from detectron2.data import MetadataCatalog | |
| from detectron2.utils.comm import all_gather, is_main_process, synchronize | |
| from detectron2.evaluation.evaluator import DatasetEvaluator | |
| import matplotlib.pyplot as plt | |
| import numpy as np | |
| import os | |
| import copy | |
| class GroundingEvaluator(DatasetEvaluator): | |
| """ | |
| Evaluate grounding segmentation metrics. | |
| """ | |
| def __init__( | |
| self, | |
| dataset_name, | |
| compute_box=False, | |
| distributed=True, | |
| ): | |
| self._logger = logging.getLogger(__name__) | |
| self._dataset_name = dataset_name | |
| self._distributed = distributed | |
| self._cpu_device = torch.device("cpu") | |
| self._compute_box = compute_box | |
| meta = MetadataCatalog.get(dataset_name) | |
| def reset(self): | |
| self.cum_I = 0 | |
| self.cum_U = 0 | |
| self.mIoU = 0 | |
| self.mDice = 0 | |
| self.cum_mean_area = 0 | |
| self.eval_seg_iou_list = [.5, .6, .7, .8, .9] | |
| self.seg_correct = torch.zeros(len(self.eval_seg_iou_list), device=self._cpu_device) | |
| self.seg_total = 0 | |
| self.instance_results = [] | |
| if self._compute_box: | |
| self.mIoU_box = 0 | |
| self.seg_correct_box = torch.zeros(len(self.eval_seg_iou_list), device=self._cpu_device) | |
| def computeIoU(pred_seg, gd_seg): | |
| I = (pred_seg & gd_seg) | |
| U = (pred_seg | gd_seg) | |
| return I, U | |
| def get_metadata(self, _input): | |
| """ | |
| Extracts and returns specific metadata from the input dictionary. | |
| Parameters: | |
| _input (dict): A dictionary containing keys like 'file_name', 'image_id', and 'grounding_info'. | |
| The 'grounding_info' is a list of dictionaries with keys like 'area', 'iscrowd', etc. | |
| Returns: | |
| dict: A dictionary containing filtered metadata. | |
| """ | |
| _input = copy.deepcopy(_input) | |
| selected_input_keys = ['file_name', 'image_id', 'grounding_info'] | |
| selected_grounding_info_keys = ['area', 'mask_file', 'iscrowd', 'image_id', 'category_id', 'id', 'file_name', 'split', 'ann_id', 'ref_id'] | |
| filtered_input = {key: _input[key] for key in selected_input_keys if key in _input} | |
| # Check if grounding_info is present and is a list | |
| if 'grounding_info' in filtered_input and isinstance(filtered_input['grounding_info'], list): | |
| # Filter each grounding_info dictionary | |
| filtered_input['grounding_info'] = [ | |
| {key: info[key] for key in selected_grounding_info_keys if key in info} | |
| for info in filtered_input['grounding_info'] | |
| ] | |
| return filtered_input | |
| def process(self, inputs, outputs): | |
| for input, output in zip(inputs, outputs): | |
| pred = output['grounding_mask'].sigmoid() > 0.5 | |
| # # save pixel probability | |
| # prob = output['grounding_mask'].sigmoid().cpu().numpy()[0] * 255 | |
| # pred_file = input['file_name'].split('.')[0].replace('test/', 'test_pred/') + '_' + input['groundings']['texts'][0].replace(' ', '+') + '.png' | |
| # if not os.path.exists('/'.join(pred_file.split('/')[:-1])): | |
| # os.makedirs('/'.join(pred_file.split('/')[:-1]), exist_ok=True) | |
| # plt.imsave(pred_file, | |
| # prob.astype(np.uint8), cmap='gray') | |
| gt = input['groundings']['masks'].bool() | |
| bsi = len(pred) | |
| I, U = self.computeIoU(pred, gt) | |
| self.cum_I += I.sum().cpu() | |
| self.cum_U += U.sum().cpu() | |
| IoU = I.reshape(bsi,-1).sum(-1)*1.0 / (U.reshape(bsi,-1).sum(-1) + 1e-6) | |
| self.mIoU += IoU.sum().cpu() | |
| # Add Dice score in eval | |
| Dice = I.reshape(bsi,-1).sum(-1)*2.0 / (gt.reshape(bsi,-1).sum(-1) + pred.reshape(bsi,-1).sum(-1) + 1e-6) | |
| self.mDice += Dice.sum().cpu() | |
| self.cum_mean_area += ((gt.reshape(bsi,-1).sum(-1) + pred.reshape(bsi,-1).sum(-1)) / 2.0).sum().cpu() | |
| if self._compute_box: | |
| pred_box = BoxMode.convert(output['grounding_box'], BoxMode.XYWH_ABS, BoxMode.XYXY_ABS) | |
| gt_box = BoxMode.convert(input['groundings']['boxes'], BoxMode.XYWH_ABS, BoxMode.XYXY_ABS).cpu() | |
| IoU_box = box_iou(pred_box, gt_box).diagonal() | |
| self.mIoU_box += IoU_box.sum() | |
| for idx in range(len(self.eval_seg_iou_list)): | |
| eval_seg_iou = self.eval_seg_iou_list[idx] | |
| self.seg_correct[idx] += (IoU >= eval_seg_iou).sum().cpu() | |
| if self._compute_box: | |
| self.seg_correct_box[idx] += (IoU_box >= eval_seg_iou).sum().cpu() | |
| self.seg_total += bsi | |
| instance_result = { | |
| 'metadata': self.get_metadata(input), | |
| 'IoU': IoU.cpu().numpy().tolist(), | |
| 'Dice': Dice.cpu().numpy().tolist(), | |
| 'I': I.sum(dim=(1, 2)).cpu().numpy().tolist(), | |
| 'U': U.sum(dim=(1, 2)).cpu().numpy().tolist(), | |
| 'IoU_box': IoU_box.cpu().numpy().tolist() if self._compute_box else '', | |
| 'pred_area': pred.reshape(bsi,-1).sum(-1).cpu().numpy().tolist(), | |
| } | |
| iou_len = IoU.shape[0] | |
| grounding_info_len = len(self.get_metadata(input)['grounding_info']) | |
| assert iou_len == grounding_info_len, f'Number of IoU scores ({iou_len}) and grounding info ({grounding_info_len}) do not match.' | |
| self.instance_results.append(instance_result) | |
| def evaluate(self): | |
| if self._distributed: | |
| synchronize() | |
| self.cum_I = torch.stack(all_gather(self.cum_I)).sum() | |
| self.cum_U = torch.stack(all_gather(self.cum_U)).sum() | |
| self.mIoU = torch.stack(all_gather(self.mIoU)).sum() | |
| self.mDice = torch.stack(all_gather(self.mDice)).sum() | |
| self.cum_mean_area = torch.stack(all_gather(self.cum_mean_area)).sum() | |
| self.seg_correct = torch.stack(all_gather(self.seg_correct)).sum(0) | |
| self.seg_total = sum(all_gather(self.seg_total)) | |
| self.instance_results = sum(all_gather(self.instance_results), []) | |
| if self._compute_box: | |
| self.mIoU_box = torch.stack(all_gather(self.mIoU_box)).sum() | |
| self.seg_correct_box = torch.stack(all_gather(self.seg_correct_box)).sum(0) | |
| if not is_main_process(): | |
| return | |
| results = {} | |
| for idx in range(len(self.eval_seg_iou_list)): | |
| result_str = 'precision@{}'.format(self.eval_seg_iou_list[idx]) | |
| results[result_str] = (self.seg_correct[idx]*100 / self.seg_total).item() | |
| results['cIoU'] = (self.cum_I*100./self.cum_U).item() | |
| results['mIoU'] = (self.mIoU*100./self.seg_total).item() | |
| results['cDice'] = (self.cum_I*100./self.cum_mean_area).item() | |
| results['mDice'] = (self.mDice*100./self.seg_total).item() | |
| if self._compute_box: | |
| for idx in range(len(self.eval_seg_iou_list)): | |
| result_str = 'precisionB@{}'.format(self.eval_seg_iou_list[idx]) | |
| results[result_str] = (self.seg_correct_box[idx]*100 / self.seg_total).item() | |
| results['mBIoU'] = (self.mIoU_box*100./self.seg_total).item() | |
| self._logger.info(results) | |
| return {'grounding': {'scores': results, 'instance_results': self.instance_results}} |