Spaces:
Sleeping
Sleeping
| # Copyright (c) OpenMMLab. All rights reserved. | |
| from typing import Dict, List, Union | |
| import numpy as np | |
| def remove_pipeline_elements(results: Dict, | |
| remove_inds: Union[List[int], | |
| np.ndarray]) -> Dict: | |
| """Remove elements in the pipeline given target indexes. | |
| Args: | |
| results (dict): Result dict from loading pipeline. | |
| remove_inds (list(int) or np.ndarray): The element indexes to be | |
| removed. | |
| Required Keys: | |
| - gt_polygons (optional) | |
| - gt_bboxes (optional) | |
| - gt_bboxes_labels (optional) | |
| - gt_ignored (optional) | |
| - gt_texts (optional) | |
| Modified Keys: | |
| - gt_polygons (optional) | |
| - gt_bboxes (optional) | |
| - gt_bboxes_labels (optional) | |
| - gt_ignored (optional) | |
| - gt_texts (optional) | |
| Returns: | |
| dict: The results with element removed. | |
| """ | |
| keys = [ | |
| 'gt_polygons', 'gt_bboxes', 'gt_bboxes_labels', 'gt_ignored', | |
| 'gt_texts' | |
| ] | |
| num_elements = -1 | |
| for key in keys: | |
| if key in results: | |
| num_elements = len(results[key]) | |
| break | |
| if num_elements == -1: | |
| return results | |
| kept_inds = np.array( | |
| [i for i in range(num_elements) if i not in remove_inds]) | |
| for key in keys: | |
| if key in results: | |
| if isinstance(results[key], np.ndarray): | |
| results[key] = results[key][kept_inds] | |
| elif isinstance(results[key], list): | |
| results[key] = [results[key][i] for i in kept_inds] | |
| else: | |
| raise TypeError( | |
| f'Unsupported type {type(results[key])} for key {key}') | |
| return results | |