|
|
|
|
|
from mmdet.datasets.builder import DATASETS |
|
|
|
|
|
import mmocr.utils as utils |
|
|
from mmocr.datasets.ocr_dataset import OCRDataset |
|
|
|
|
|
|
|
|
@DATASETS.register_module() |
|
|
class OCRSegDataset(OCRDataset): |
|
|
|
|
|
def pre_pipeline(self, results): |
|
|
results['img_prefix'] = self.img_prefix |
|
|
|
|
|
def _parse_anno_info(self, annotations): |
|
|
"""Parse char boxes annotations. |
|
|
Args: |
|
|
annotations (list[dict]): Annotations of one image, where |
|
|
each dict is for one character. |
|
|
|
|
|
Returns: |
|
|
dict: A dict containing the following keys: |
|
|
|
|
|
- chars (list[str]): List of character strings. |
|
|
- char_rects (list[list[float]]): List of char box, with each |
|
|
in style of rectangle: [x_min, y_min, x_max, y_max]. |
|
|
- char_quads (list[list[float]]): List of char box, with each |
|
|
in style of quadrangle: [x1, y1, x2, y2, x3, y3, x4, y4]. |
|
|
""" |
|
|
|
|
|
assert utils.is_type_list(annotations, dict) |
|
|
assert 'char_box' in annotations[0] |
|
|
assert 'char_text' in annotations[0] |
|
|
assert len(annotations[0]['char_box']) in [4, 8] |
|
|
|
|
|
chars, char_rects, char_quads = [], [], [] |
|
|
for ann in annotations: |
|
|
char_box = ann['char_box'] |
|
|
if len(char_box) == 4: |
|
|
char_box_type = ann.get('char_box_type', 'xyxy') |
|
|
if char_box_type == 'xyxy': |
|
|
char_rects.append(char_box) |
|
|
char_quads.append([ |
|
|
char_box[0], char_box[1], char_box[2], char_box[1], |
|
|
char_box[2], char_box[3], char_box[0], char_box[3] |
|
|
]) |
|
|
elif char_box_type == 'xywh': |
|
|
x1, y1, w, h = char_box |
|
|
x2 = x1 + w |
|
|
y2 = y1 + h |
|
|
char_rects.append([x1, y1, x2, y2]) |
|
|
char_quads.append([x1, y1, x2, y1, x2, y2, x1, y2]) |
|
|
else: |
|
|
raise ValueError(f'invalid char_box_type {char_box_type}') |
|
|
elif len(char_box) == 8: |
|
|
x_list, y_list = [], [] |
|
|
for i in range(4): |
|
|
x_list.append(char_box[2 * i]) |
|
|
y_list.append(char_box[2 * i + 1]) |
|
|
x_max, x_min = max(x_list), min(x_list) |
|
|
y_max, y_min = max(y_list), min(y_list) |
|
|
char_rects.append([x_min, y_min, x_max, y_max]) |
|
|
char_quads.append(char_box) |
|
|
else: |
|
|
raise Exception( |
|
|
f'invalid num in char box: {len(char_box)} not in (4, 8)') |
|
|
chars.append(ann['char_text']) |
|
|
|
|
|
ann = dict(chars=chars, char_rects=char_rects, char_quads=char_quads) |
|
|
|
|
|
return ann |
|
|
|
|
|
def prepare_train_img(self, index): |
|
|
"""Get training data and annotations from pipeline. |
|
|
|
|
|
Args: |
|
|
index (int): Index of data. |
|
|
|
|
|
Returns: |
|
|
dict: Training data and annotation after pipeline with new keys |
|
|
introduced by pipeline. |
|
|
""" |
|
|
img_ann_info = self.data_infos[index] |
|
|
img_info = { |
|
|
'filename': img_ann_info['file_name'], |
|
|
} |
|
|
ann_info = self._parse_anno_info(img_ann_info['annotations']) |
|
|
results = dict(img_info=img_info, ann_info=ann_info) |
|
|
|
|
|
self.pre_pipeline(results) |
|
|
|
|
|
return self.pipeline(results) |
|
|
|
