MAERec-Gradio / mmocr /utils /transform_utils.py
Mountchicken's picture
Upload 704 files
9bf4bd7
raw
history blame
1.71 kB
# 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