tests/test_utils/test_ocr.py

# Copyright (c) OpenMMLab. All rights reserved.
import io
import json
import os
import platform
import random
import sys
import tempfile
from pathlib import Path
from unittest import mock

import mmcv
import numpy as np
import pytest
import torch

from mmocr.apis import init_detector
from mmocr.datasets.kie_dataset import KIEDataset
from mmocr.utils.ocr import MMOCR


def test_ocr_init_errors():
    # Test assertions
    with pytest.raises(ValueError):
        _ = MMOCR(det='test')
    with pytest.raises(ValueError):
        _ = MMOCR(recog='test')
    with pytest.raises(ValueError):
        _ = MMOCR(kie='test')
    with pytest.raises(NotImplementedError):
        _ = MMOCR(det=None, recog=None, kie='SDMGR')
    with pytest.raises(NotImplementedError):
        _ = MMOCR(det='DB_r18', recog=None, kie='SDMGR')


cfg_default_prefix = os.path.join(str(Path.cwd()), 'configs/')


@pytest.mark.parametrize(
    'det, recog, kie, config_dir, gt_cfg, gt_ckpt',
    [('DB_r18', None, '', '',
      cfg_default_prefix + 'textdet/dbnet/dbnet_r18_fpnc_1200e_icdar2015.py',
      'https://download.openmmlab.com/mmocr/textdet/'
      'dbnet/dbnet_r18_fpnc_sbn_1200e_icdar2015_20210329-ba3ab597.pth'),
     (None, 'CRNN', '', '',
      cfg_default_prefix + 'textrecog/crnn/crnn_academic_dataset.py',
      'https://download.openmmlab.com/mmocr/textrecog/'
      'crnn/crnn_academic-a723a1c5.pth'),
     ('DB_r18', 'CRNN', 'SDMGR', '', [
         cfg_default_prefix +
         'textdet/dbnet/dbnet_r18_fpnc_1200e_icdar2015.py',
         cfg_default_prefix + 'textrecog/crnn/crnn_academic_dataset.py',
         cfg_default_prefix + 'kie/sdmgr/sdmgr_unet16_60e_wildreceipt.py'
     ], [
         'https://download.openmmlab.com/mmocr/textdet/'
         'dbnet/dbnet_r18_fpnc_sbn_1200e_icdar2015_20210329-ba3ab597.pth',
         'https://download.openmmlab.com/mmocr/textrecog/'
         'crnn/crnn_academic-a723a1c5.pth',
         'https://download.openmmlab.com/mmocr/kie/'
         'sdmgr/sdmgr_unet16_60e_wildreceipt_20210520-7489e6de.pth'
     ]),
     ('DB_r18', 'CRNN', 'SDMGR', 'test/', [
         'test/textdet/dbnet/dbnet_r18_fpnc_1200e_icdar2015.py',
         'test/textrecog/crnn/crnn_academic_dataset.py',
         'test/kie/sdmgr/sdmgr_unet16_60e_wildreceipt.py'
     ], [
         'https://download.openmmlab.com/mmocr/textdet/'
         'dbnet/dbnet_r18_fpnc_sbn_1200e_icdar2015_20210329-ba3ab597.pth',
         'https://download.openmmlab.com/mmocr/textrecog/'
         'crnn/crnn_academic-a723a1c5.pth',
         'https://download.openmmlab.com/mmocr/kie/'
         'sdmgr/sdmgr_unet16_60e_wildreceipt_20210520-7489e6de.pth'
     ])],
)
@mock.patch('mmocr.utils.ocr.init_detector')
@mock.patch('mmocr.utils.ocr.build_detector')
@mock.patch('mmocr.utils.ocr.Config.fromfile')
@mock.patch('mmocr.utils.ocr.load_checkpoint')
def test_ocr_init(mock_loading, mock_config, mock_build_detector,
                  mock_init_detector, det, recog, kie, config_dir, gt_cfg,
                  gt_ckpt):

    def loadcheckpoint_assert(*args, **kwargs):
        assert args[1] == gt_ckpt[-1]
        assert kwargs['map_location'] == torch.device(
            'cuda' if torch.cuda.is_available() else 'cpu')

    mock_loading.side_effect = loadcheckpoint_assert
    with mock.patch('mmocr.utils.ocr.revert_sync_batchnorm'):
        if kie == '':
            if config_dir == '':
                _ = MMOCR(det=det, recog=recog)
            else:
                _ = MMOCR(det=det, recog=recog, config_dir=config_dir)
        else:
            if config_dir == '':
                _ = MMOCR(det=det, recog=recog, kie=kie)
            else:
                _ = MMOCR(det=det, recog=recog, kie=kie, config_dir=config_dir)
        if isinstance(gt_cfg, str):
            gt_cfg = [gt_cfg]
        if isinstance(gt_ckpt, str):
            gt_ckpt = [gt_ckpt]

        i_range = range(len(gt_cfg))
        if kie:
            i_range = i_range[:-1]
            mock_config.assert_called_with(gt_cfg[-1])
            mock_build_detector.assert_called_once()
            mock_loading.assert_called_once()
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        calls = [
            mock.call(gt_cfg[i], gt_ckpt[i], device=device) for i in i_range
        ]
        mock_init_detector.assert_has_calls(calls)


@pytest.mark.parametrize(
    'det, det_config, det_ckpt, recog, recog_config, recog_ckpt,'
    'kie, kie_config, kie_ckpt, config_dir, gt_cfg, gt_ckpt',
    [('DB_r18', 'test.py', '', 'CRNN', 'test.py', '', 'SDMGR', 'test.py', '',
      'configs/', ['test.py', 'test.py', 'test.py'], [
          'https://download.openmmlab.com/mmocr/textdet/'
          'dbnet/dbnet_r18_fpnc_sbn_1200e_icdar2015_20210329-ba3ab597.pth',
          'https://download.openmmlab.com/mmocr/textrecog/'
          'crnn/crnn_academic-a723a1c5.pth',
          'https://download.openmmlab.com/mmocr/kie/'
          'sdmgr/sdmgr_unet16_60e_wildreceipt_20210520-7489e6de.pth'
      ]),
     ('DB_r18', '', 'test.ckpt', 'CRNN', '', 'test.ckpt', 'SDMGR', '',
      'test.ckpt', '', [
          'textdet/dbnet/dbnet_r18_fpnc_1200e_icdar2015.py',
          'textrecog/crnn/crnn_academic_dataset.py',
          'kie/sdmgr/sdmgr_unet16_60e_wildreceipt.py'
      ], ['test.ckpt', 'test.ckpt', 'test.ckpt']),
     ('DB_r18', 'test.py', 'test.ckpt', 'CRNN', 'test.py', 'test.ckpt',
      'SDMGR', 'test.py', 'test.ckpt', '', ['test.py', 'test.py', 'test.py'],
      ['test.ckpt', 'test.ckpt', 'test.ckpt'])])
@mock.patch('mmocr.utils.ocr.init_detector')
@mock.patch('mmocr.utils.ocr.build_detector')
@mock.patch('mmocr.utils.ocr.Config.fromfile')
@mock.patch('mmocr.utils.ocr.load_checkpoint')
def test_ocr_init_customize_config(mock_loading, mock_config,
                                   mock_build_detector, mock_init_detector,
                                   det, det_config, det_ckpt, recog,
                                   recog_config, recog_ckpt, kie, kie_config,
                                   kie_ckpt, config_dir, gt_cfg, gt_ckpt):

    def loadcheckpoint_assert(*args, **kwargs):
        assert args[1] == gt_ckpt[-1]

    mock_loading.side_effect = loadcheckpoint_assert
    with mock.patch('mmocr.utils.ocr.revert_sync_batchnorm'):
        _ = MMOCR(
            det=det,
            det_config=det_config,
            det_ckpt=det_ckpt,
            recog=recog,
            recog_config=recog_config,
            recog_ckpt=recog_ckpt,
            kie=kie,
            kie_config=kie_config,
            kie_ckpt=kie_ckpt,
            config_dir=config_dir)

        i_range = range(len(gt_cfg))
        if kie:
            i_range = i_range[:-1]
            mock_config.assert_called_with(gt_cfg[-1])
            mock_build_detector.assert_called_once()
            mock_loading.assert_called_once()
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        calls = [
            mock.call(gt_cfg[i], gt_ckpt[i], device=device) for i in i_range
        ]
        mock_init_detector.assert_has_calls(calls)


@mock.patch('mmocr.utils.ocr.init_detector')
@mock.patch('mmocr.utils.ocr.build_detector')
@mock.patch('mmocr.utils.ocr.Config.fromfile')
@mock.patch('mmocr.utils.ocr.load_checkpoint')
@mock.patch('mmocr.utils.ocr.model_inference')
def test_single_inference(mock_model_inference, mock_loading, mock_config,
                          mock_build_detector, mock_init_detector):

    def dummy_inference(model, arr, batch_mode):
        return arr

    mock_model_inference.side_effect = dummy_inference
    mmocr = MMOCR()

    data = list(range(20))
    model = 'dummy'
    res = mmocr.single_inference(model, data, batch_mode=False)
    assert (data == res)
    mock_model_inference.reset_mock()

    res = mmocr.single_inference(model, data, batch_mode=True)
    assert (data == res)
    mock_model_inference.assert_called_once()
    mock_model_inference.reset_mock()

    res = mmocr.single_inference(model, data, batch_mode=True, batch_size=100)
    assert (data == res)
    mock_model_inference.assert_called_once()
    mock_model_inference.reset_mock()

    res = mmocr.single_inference(model, data, batch_mode=True, batch_size=3)
    assert (data == res)


@mock.patch('mmocr.utils.ocr.init_detector')
@mock.patch('mmocr.utils.ocr.load_checkpoint')
def MMOCR_testobj(mock_loading, mock_init_detector, **kwargs):
    # returns an MMOCR object bypassing the
    # checkpoint initialization step
    def init_detector_skip_ckpt(config, ckpt, device):
        return init_detector(config, device=device)

    def modify_kie_class(model, ckpt, map_location):
        model.class_list = 'tests/data/kie_toy_dataset/class_list.txt'

    mock_init_detector.side_effect = init_detector_skip_ckpt
    mock_loading.side_effect = modify_kie_class
    kwargs['det'] = kwargs.get('det', 'DB_r18')
    kwargs['recog'] = kwargs.get('recog', 'CRNN')
    kwargs['kie'] = kwargs.get('kie', 'SDMGR')
    device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
    return MMOCR(**kwargs, device=device)


@pytest.mark.skipif(
    platform.system() == 'Windows',
    reason='Win container on Github Action does not have enough RAM to run')
@mock.patch('mmocr.utils.ocr.KIEDataset')
def test_readtext(mock_kiedataset):
    # Fixing the weights of models to prevent them from
    # generating invalid results and triggering other assertion errors
    torch.manual_seed(4)
    random.seed(4)
    mmocr = MMOCR_testobj()
    mmocr_det = MMOCR_testobj(kie='', recog='')
    mmocr_recog = MMOCR_testobj(kie='', det='', recog='CRNN_TPS')
    mmocr_det_recog = MMOCR_testobj(kie='')

    def readtext(imgs, ocr_obj=mmocr, **kwargs):
        # filename can be different depends on how
        # the the image was loaded
        e2e_res = ocr_obj.readtext(imgs, **kwargs)
        for res in e2e_res:
            res.pop('filename')
        return e2e_res

    def kiedataset_with_test_dict(**kwargs):
        kwargs['dict_file'] = 'tests/data/kie_toy_dataset/dict.txt'
        return KIEDataset(**kwargs)

    mock_kiedataset.side_effect = kiedataset_with_test_dict

    # Single image
    toy_dir = 'tests/data/toy_dataset/imgs/test/'
    toy_img1_path = toy_dir + 'img_1.jpg'
    str_e2e_res = readtext(toy_img1_path)
    toy_img1 = mmcv.imread(toy_img1_path)
    np_e2e_res = readtext(toy_img1)
    assert str_e2e_res == np_e2e_res

    # Multiple images
    toy_img2_path = toy_dir + 'img_2.jpg'
    toy_img2 = mmcv.imread(toy_img2_path)
    toy_imgs = [toy_img1, toy_img2]
    toy_img_paths = [toy_img1_path, toy_img2_path]
    np_e2e_results = readtext(toy_imgs)
    str_e2e_results = readtext(toy_img_paths)
    str_tuple_e2e_results = readtext(tuple(toy_img_paths))
    assert np_e2e_results == str_e2e_results
    assert str_e2e_results == str_tuple_e2e_results

    # Batch mode test
    toy_imgs.append(toy_dir + 'img_3.jpg')
    e2e_res = readtext(toy_imgs)
    full_batch_e2e_res = readtext(toy_imgs, batch_mode=True)
    assert full_batch_e2e_res == e2e_res
    batch_e2e_res = readtext(
        toy_imgs, batch_mode=True, recog_batch_size=2, det_batch_size=2)
    assert batch_e2e_res == full_batch_e2e_res

    # Batch mode test with DBNet only
    full_batch_det_res = mmocr_det.readtext(toy_imgs, batch_mode=True)
    det_res = mmocr_det.readtext(toy_imgs)
    batch_det_res = mmocr_det.readtext(
        toy_imgs, batch_mode=True, single_batch_size=2)
    assert len(full_batch_det_res) == len(det_res)
    assert len(batch_det_res) == len(det_res)
    assert all([
        np.allclose(full_batch_det_res[i]['boundary_result'],
                    det_res[i]['boundary_result'])
        for i in range(len(full_batch_det_res))
    ])
    assert all([
        np.allclose(batch_det_res[i]['boundary_result'],
                    det_res[i]['boundary_result'])
        for i in range(len(batch_det_res))
    ])

    # Batch mode test with CRNN_TPS only (CRNN doesn't support batch inference)
    full_batch_recog_res = mmocr_recog.readtext(toy_imgs, batch_mode=True)
    recog_res = mmocr_recog.readtext(toy_imgs)
    batch_recog_res = mmocr_recog.readtext(
        toy_imgs, batch_mode=True, single_batch_size=2)
    full_batch_recog_res.sort(key=lambda x: x['text'])
    batch_recog_res.sort(key=lambda x: x['text'])
    recog_res.sort(key=lambda x: x['text'])
    assert np.all([
        np.allclose(full_batch_recog_res[i]['score'], recog_res[i]['score'])
        for i in range(len(full_batch_recog_res))
    ])
    assert np.all([
        np.allclose(batch_recog_res[i]['score'], recog_res[i]['score'])
        for i in range(len(full_batch_recog_res))
    ])

    # Test export
    with tempfile.TemporaryDirectory() as tmpdirname:
        mmocr.readtext(toy_imgs, export=tmpdirname)
        assert len(os.listdir(tmpdirname)) == len(toy_imgs)
    with tempfile.TemporaryDirectory() as tmpdirname:
        mmocr_det.readtext(toy_imgs, export=tmpdirname)
        assert len(os.listdir(tmpdirname)) == len(toy_imgs)
    with tempfile.TemporaryDirectory() as tmpdirname:
        mmocr_recog.readtext(toy_imgs, export=tmpdirname)
        assert len(os.listdir(tmpdirname)) == len(toy_imgs)

    # Test output
    # Single image
    with tempfile.TemporaryDirectory() as tmpdirname:
        tmp_output = os.path.join(tmpdirname, '1.jpg')
        mmocr.readtext(toy_imgs[0], output=tmp_output)
        assert os.path.exists(tmp_output)
    # Multiple images
    with tempfile.TemporaryDirectory() as tmpdirname:
        mmocr.readtext(toy_imgs, output=tmpdirname)
        assert len(os.listdir(tmpdirname)) == len(toy_imgs)

    # Test imshow
    with mock.patch('mmocr.utils.ocr.mmcv.imshow') as mock_imshow:
        mmocr.readtext(toy_img1_path, imshow=True)
        mock_imshow.assert_called_once()
        mock_imshow.reset_mock()
        mmocr.readtext(toy_imgs, imshow=True)
        assert mock_imshow.call_count == len(toy_imgs)

    # Test print_result
    with io.StringIO() as capturedOutput:
        sys.stdout = capturedOutput
        res = mmocr.readtext(toy_imgs, print_result=True)
        assert json.loads('[%s]' % capturedOutput.getvalue().strip().replace(
            '\n\n', ',').replace("'", '"')) == res
        sys.stdout = sys.__stdout__
    with io.StringIO() as capturedOutput:
        sys.stdout = capturedOutput
        res = mmocr.readtext(toy_imgs, details=True, print_result=True)
        assert json.loads('[%s]' % capturedOutput.getvalue().strip().replace(
            '\n\n', ',').replace("'", '"')) == res
        sys.stdout = sys.__stdout__

    # Test merge
    with mock.patch('mmocr.utils.ocr.stitch_boxes_into_lines') as mock_merge:
        mmocr_det_recog.readtext(toy_imgs, merge=True)
        assert mock_merge.call_count == len(toy_imgs)