tests/test_dataset/test_loading.py

# Copyright (c) OpenMMLab. All rights reserved.
import copy

import numpy as np

from mmocr.datasets.pipelines import LoadImageFromNdarray, LoadTextAnnotations


def _create_dummy_ann():
    results = {}
    results['img_info'] = {}
    results['img_info']['height'] = 1000
    results['img_info']['width'] = 1000
    results['ann_info'] = {}
    results['ann_info']['masks'] = []
    results['mask_fields'] = []
    results['ann_info']['masks_ignore'] = [
        [[499, 94, 531, 94, 531, 124, 499, 124]],
        [[3, 156, 81, 155, 78, 181, 0, 182]],
        [[11, 223, 59, 221, 59, 234, 11, 236]],
        [[500, 156, 551, 156, 550, 165, 499, 165]]
    ]

    return results


def test_loadtextannotation():

    results = _create_dummy_ann()
    with_bbox = True
    with_label = True
    with_mask = True
    with_seg = False
    poly2mask = False

    # If no 'ori_shape' in result but use_img_shape=True,
    # result['img_info']['height'] and result['img_info']['width']
    # will be used to generate mask.
    loader = LoadTextAnnotations(
        with_bbox,
        with_label,
        with_mask,
        with_seg,
        poly2mask,
        use_img_shape=True)
    tmp_results = copy.deepcopy(results)
    output = loader._load_masks(tmp_results)
    assert len(output['gt_masks_ignore']) == 4
    assert np.allclose(output['gt_masks_ignore'].masks[0],
                       [[499, 94, 531, 94, 531, 124, 499, 124]])
    assert output['gt_masks_ignore'].height == results['img_info']['height']

    # If 'ori_shape' in result and use_img_shape=True,
    # result['ori_shape'] will be used to generate mask.
    loader = LoadTextAnnotations(
        with_bbox,
        with_label,
        with_mask,
        with_seg,
        poly2mask=True,
        use_img_shape=True)
    tmp_results = copy.deepcopy(results)
    tmp_results['ori_shape'] = (640, 640, 3)
    output = loader._load_masks(tmp_results)
    assert output['img_info']['height'] == 640
    assert output['gt_masks_ignore'].height == 640


def test_load_img_from_numpy():
    result = {'img': np.ones((32, 100, 3), dtype=np.uint8)}

    load = LoadImageFromNdarray(color_type='color')
    output = load(result)

    assert output['img'].shape[2] == 3
    assert len(output['img'].shape) == 3

    result = {'img': np.ones((32, 100, 1), dtype=np.uint8)}
    load = LoadImageFromNdarray(color_type='color')
    output = load(result)
    assert output['img'].shape[2] == 3

    result = {'img': np.ones((32, 100, 3), dtype=np.uint8)}
    load = LoadImageFromNdarray(color_type='grayscale', to_float32=True)
    output = load(result)
    assert output['img'].shape[2] == 1