Add text recognition model

main.py

@@ -3,8 +3,11 @@ import gradio as gr
 import cv2
 import sys
 
-ocr = MMOCR(det_config='model/config.py',
-            det_ckpt='model/model.pth', device='cpu')
+ocr = MMOCR(det_config='model/det/config.py',
+            det_ckpt='model/det/model.pth',
+            recog_config='model/recog/config.py',
+            recog_ckpt='model/det/model.pth',
+            device='cpu')
 
 def get_rec(points):
     xs = []
@@ -34,13 +37,13 @@ def predict(image_input, score_threshold):
         p0, p1 = get_rec([int(i) for i in polygon])
         draw_img = cv2.rectangle(draw_img, p0, p1, (255,0,0), 2)
 
-    return draw_img
+    return draw_img, output
 
 def run():
     demo = gr.Interface(
         fn=predict,
         inputs=[gr.components.Image(), gr.Slider(0, 1, 0.8)],
-        outputs=gr.components.Image(),
+        outputs=[gr.components.Image(), gr.JSON()],
     )
 
     demo.launch(server_name="0.0.0.0", server_port=7860)

model/recog/config.py

@@ -0,0 +1,298 @@
+thvl_textrecog_data_root = 'data/recog/synTH'
+thvl_textrecog_train = dict(
+    type='OCRDataset',
+    data_root='data/recog/synTH',
+    ann_file='textrecog_train.json',
+    pipeline=None)
+thvl_textrecog_test = dict(
+    type='OCRDataset',
+    data_root='data/recog/synTH',
+    ann_file='textrecog_test.json',
+    test_mode=True,
+    pipeline=None)
+default_scope = 'mmocr'
+env_cfg = dict(
+    cudnn_benchmark=True,
+    mp_cfg=dict(mp_start_method='fork', opencv_num_threads=0),
+    dist_cfg=dict(backend='nccl'))
+randomness = dict(seed=None)
+default_hooks = dict(
+    timer=dict(type='IterTimerHook'),
+    logger=dict(type='LoggerHook', interval=100),
+    param_scheduler=dict(type='ParamSchedulerHook'),
+    checkpoint=dict(type='CheckpointHook', interval=1),
+    sampler_seed=dict(type='DistSamplerSeedHook'),
+    sync_buffer=dict(type='SyncBuffersHook'),
+    visualization=dict(
+        type='VisualizationHook',
+        interval=1,
+        enable=False,
+        show=False,
+        draw_gt=False,
+        draw_pred=False))
+log_level = 'INFO'
+log_processor = dict(type='LogProcessor', window_size=10, by_epoch=True)
+load_from = None
+resume = False
+val_evaluator = dict(
+    type='MultiDatasetsEvaluator',
+    metrics=[
+        dict(
+            type='WordMetric',
+            mode=['exact', 'ignore_case', 'ignore_case_symbol']),
+        dict(type='CharMetric')
+    ],
+    dataset_prefixes=None)
+test_evaluator = dict(
+    type='MultiDatasetsEvaluator',
+    metrics=[
+        dict(
+            type='WordMetric',
+            mode=['exact', 'ignore_case', 'ignore_case_symbol']),
+        dict(type='CharMetric')
+    ],
+    dataset_prefixes=None)
+vis_backends = [dict(type='LocalVisBackend')]
+visualizer = dict(
+    type='TextRecogLocalVisualizer',
+    name='visualizer',
+    vis_backends=[dict(type='LocalVisBackend')])
+optim_wrapper = dict(
+    type='OptimWrapper', optimizer=dict(type='Adam', lr=0.0003))
+train_cfg = dict(type='EpochBasedTrainLoop', max_epochs=50, val_interval=1)
+val_cfg = dict(type='ValLoop')
+test_cfg = dict(type='TestLoop')
+param_scheduler = [dict(type='MultiStepLR', milestones=[3, 4], end=6)]
+file_client_args = dict(backend='disk')
+dictionary = dict(
+    type='Dictionary',
+    dict_file=
+    '/content/mmocr/configs/textrecog/nrtr/../../../dicts/thai_digits_symbols.txt',
+    with_padding=True,
+    with_unknown=True,
+    same_start_end=True,
+    with_start=True,
+    with_end=True)
+model = dict(
+    type='NRTR',
+    backbone=dict(type='NRTRModalityTransform'),
+    encoder=dict(type='NRTREncoder', n_layers=12),
+    decoder=dict(
+        type='NRTRDecoder',
+        module_loss=dict(
+            type='CEModuleLoss', ignore_first_char=True, flatten=True),
+        postprocessor=dict(type='AttentionPostprocessor'),
+        dictionary=dict(
+            type='Dictionary',
+            dict_file=
+            '/content/mmocr/configs/textrecog/nrtr/../../../dicts/thai_digits_symbols.txt',
+            with_padding=True,
+            with_unknown=True,
+            same_start_end=True,
+            with_start=True,
+            with_end=True),
+        max_seq_len=30),
+    data_preprocessor=dict(
+        type='TextRecogDataPreprocessor',
+        mean=[123.675, 116.28, 103.53],
+        std=[58.395, 57.12, 57.375]))
+train_pipeline = [
+    dict(
+        type='LoadImageFromFile',
+        file_client_args=dict(backend='disk'),
+        ignore_empty=True,
+        min_size=2),
+    dict(type='LoadOCRAnnotations', with_text=True),
+    dict(
+        type='RescaleToHeight',
+        height=32,
+        min_width=32,
+        max_width=160,
+        width_divisor=4),
+    dict(type='PadToWidth', width=160),
+    dict(
+        type='PackTextRecogInputs',
+        meta_keys=('img_path', 'ori_shape', 'img_shape', 'valid_ratio'))
+]
+test_pipeline = [
+    dict(type='LoadImageFromFile', file_client_args=dict(backend='disk')),
+    dict(
+        type='RescaleToHeight',
+        height=32,
+        min_width=32,
+        max_width=160,
+        width_divisor=16),
+    dict(type='PadToWidth', width=160),
+    dict(type='LoadOCRAnnotations', with_text=True),
+    dict(
+        type='PackTextRecogInputs',
+        meta_keys=('img_path', 'ori_shape', 'img_shape', 'valid_ratio'))
+]
+train_list = [
+    dict(
+        type='OCRDataset',
+        data_root='data/recog/synTH',
+        ann_file='textrecog_train.json',
+        pipeline=None)
+]
+test_list = [
+    dict(
+        type='OCRDataset',
+        data_root='data/recog/synTH',
+        ann_file='textrecog_test.json',
+        test_mode=True,
+        pipeline=None)
+]
+train_dataset = dict(
+    type='ConcatDataset',
+    datasets=[
+        dict(
+            type='OCRDataset',
+            data_root='data/recog/synTH',
+            ann_file='textrecog_train.json',
+            pipeline=None)
+    ],
+    pipeline=[
+        dict(
+            type='LoadImageFromFile',
+            file_client_args=dict(backend='disk'),
+            ignore_empty=True,
+            min_size=2),
+        dict(type='LoadOCRAnnotations', with_text=True),
+        dict(
+            type='RescaleToHeight',
+            height=32,
+            min_width=32,
+            max_width=160,
+            width_divisor=4),
+        dict(type='PadToWidth', width=160),
+        dict(
+            type='PackTextRecogInputs',
+            meta_keys=('img_path', 'ori_shape', 'img_shape', 'valid_ratio'))
+    ])
+test_dataset = dict(
+    type='ConcatDataset',
+    datasets=[
+        dict(
+            type='OCRDataset',
+            data_root='data/recog/synTH',
+            ann_file='textrecog_test.json',
+            test_mode=True,
+            pipeline=None)
+    ],
+    pipeline=[
+        dict(type='LoadImageFromFile', file_client_args=dict(backend='disk')),
+        dict(
+            type='RescaleToHeight',
+            height=32,
+            min_width=32,
+            max_width=160,
+            width_divisor=16),
+        dict(type='PadToWidth', width=160),
+        dict(type='LoadOCRAnnotations', with_text=True),
+        dict(
+            type='PackTextRecogInputs',
+            meta_keys=('img_path', 'ori_shape', 'img_shape', 'valid_ratio'))
+    ])
+train_dataloader = dict(
+    batch_size=384,
+    num_workers=24,
+    persistent_workers=True,
+    sampler=dict(type='DefaultSampler', shuffle=True),
+    dataset=dict(
+        type='ConcatDataset',
+        datasets=[
+            dict(
+                type='OCRDataset',
+                data_root='data/recog/synTH',
+                ann_file='textrecog_train.json',
+                pipeline=None)
+        ],
+        pipeline=[
+            dict(
+                type='LoadImageFromFile',
+                file_client_args=dict(backend='disk'),
+                ignore_empty=True,
+                min_size=2),
+            dict(type='LoadOCRAnnotations', with_text=True),
+            dict(
+                type='RescaleToHeight',
+                height=32,
+                min_width=32,
+                max_width=160,
+                width_divisor=4),
+            dict(type='PadToWidth', width=160),
+            dict(
+                type='PackTextRecogInputs',
+                meta_keys=('img_path', 'ori_shape', 'img_shape',
+                           'valid_ratio'))
+        ]))
+test_dataloader = dict(
+    batch_size=1,
+    num_workers=4,
+    persistent_workers=True,
+    drop_last=False,
+    sampler=dict(type='DefaultSampler', shuffle=False),
+    dataset=dict(
+        type='ConcatDataset',
+        datasets=[
+            dict(
+                type='OCRDataset',
+                data_root='data/recog/synTH',
+                ann_file='textrecog_test.json',
+                test_mode=True,
+                pipeline=None)
+        ],
+        pipeline=[
+            dict(
+                type='LoadImageFromFile',
+                file_client_args=dict(backend='disk')),
+            dict(
+                type='RescaleToHeight',
+                height=32,
+                min_width=32,
+                max_width=160,
+                width_divisor=16),
+            dict(type='PadToWidth', width=160),
+            dict(type='LoadOCRAnnotations', with_text=True),
+            dict(
+                type='PackTextRecogInputs',
+                meta_keys=('img_path', 'ori_shape', 'img_shape',
+                           'valid_ratio'))
+        ]))
+val_dataloader = dict(
+    batch_size=1,
+    num_workers=4,
+    persistent_workers=True,
+    drop_last=False,
+    sampler=dict(type='DefaultSampler', shuffle=False),
+    dataset=dict(
+        type='ConcatDataset',
+        datasets=[
+            dict(
+                type='OCRDataset',
+                data_root='data/recog/synTH',
+                ann_file='textrecog_test.json',
+                test_mode=True,
+                pipeline=None)
+        ],
+        pipeline=[
+            dict(
+                type='LoadImageFromFile',
+                file_client_args=dict(backend='disk')),
+            dict(
+                type='RescaleToHeight',
+                height=32,
+                min_width=32,
+                max_width=160,
+                width_divisor=16),
+            dict(type='PadToWidth', width=160),
+            dict(type='LoadOCRAnnotations', with_text=True),
+            dict(
+                type='PackTextRecogInputs',
+                meta_keys=('img_path', 'ori_shape', 'img_shape',
+                           'valid_ratio'))
+        ]))
+auto_scale_lr = dict(base_batch_size=384)
+launcher = 'none'
+work_dir = './work_dirs/nrtr_modality-transform_50e_thvl'

model/recog/model.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5b8611b397bd03159632e0e2b63f30f0f0ce0089107255fdf934dc94a063f59b
+size 365107614