Spaces:
Sleeping
Sleeping
File size: 1,707 Bytes
9bf4bd7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 |
# 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
|