[Feature] ABCNet train (#1610)

* abcnet train

* fix comment

* updata link

* fix lint

* fix name

Author: Harold-lkk

mmocr/models/textdet/heads/base.py

@@ -108,7 +108,7 @@ def loss_and_predict(self, x: Tuple[Tensor], data_samples: DetSampleList
         outs = self(x, data_samples)
         losses = self.module_loss(outs, data_samples)
 
-        predictions = self.postprocessor(outs, data_samples)
+        predictions = self.postprocessor(outs, data_samples, self.training)
         return losses, predictions
 
     def predict(self, x: torch.Tensor,

projects/ABCNet/abcnet/model/__init__.py

@@ -1,17 +1,19 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 from .abcnet import ABCNet
 from .abcnet_det_head import ABCNetDetHead
+from .abcnet_det_module_loss import ABCNetDetModuleLoss
 from .abcnet_det_postprocessor import ABCNetDetPostprocessor
 from .abcnet_postprocessor import ABCNetPostprocessor
 from .abcnet_rec import ABCNetRec
 from .abcnet_rec_backbone import ABCNetRecBackbone
 from .abcnet_rec_decoder import ABCNetRecDecoder
 from .abcnet_rec_encoder import ABCNetRecEncoder
 from .bezier_roi_extractor import BezierRoIExtractor
-from .only_rec_roi_head import OnlyRecRoIHead
+from .rec_roi_head import RecRoIHead
 
 __all__ = [
     'ABCNetDetHead', 'ABCNetDetPostprocessor', 'ABCNetRecBackbone',
     'ABCNetRecDecoder', 'ABCNetRecEncoder', 'ABCNet', 'ABCNetRec',
-    'BezierRoIExtractor', 'OnlyRecRoIHead', 'ABCNetPostprocessor'
+    'BezierRoIExtractor', 'RecRoIHead', 'ABCNetPostprocessor',
+    'ABCNetDetModuleLoss'
 ]

projects/ABCNet/abcnet/model/abcnet_det_module_loss.py

@@ -1,6 +1,7 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 from typing import Tuple
 
+from mmengine.structures import LabelData
 from torch import Tensor
 
 from mmocr.registry import MODELS, TASK_UTILS
@@ -10,7 +11,7 @@
 
 
 @MODELS.register_module()
-class OnlyRecRoIHead(BaseRoIHead):
+class RecRoIHead(BaseRoIHead):
     """Simplest base roi head including one bbox head and one mask head."""
 
     def __init__(self,
@@ -39,8 +40,17 @@ def loss(self, inputs: Tuple[Tensor], data_samples: DetSampleList) -> dict:
         Returns:
             dict[str, Tensor]: A dictionary of loss components
         """
+        proposals = [
+            ds.gt_instances[~ds.gt_instances.ignored] for ds in data_samples
+        ]
 
-        pass
+        proposals = [p for p in proposals if len(p) > 0]
+        bbox_feats = self.roi_extractor(inputs, proposals)
+        rec_data_samples = [
+            TextRecogDataSample(gt_text=LabelData(item=text))
+            for proposal in proposals for text in proposal.texts
+        ]
+        return self.rec_head.loss(bbox_feats, rec_data_samples)
 
     def predict(self, inputs: Tuple[Tensor],
                 data_samples: DetSampleList) -> RecSampleList:

projects/ABCNet/abcnet/model/only_rec_roi_head.py renamed to projects/ABCNet/abcnet/model/rec_roi_head.py

@@ -70,7 +70,14 @@ def extract_feat(self, img):
 
     def loss(self, inputs: torch.Tensor,
              data_samples: OptDetSampleList) -> Dict:
-        pass
+        losses = dict()
+        inputs = self.extract_feat(inputs)
+        det_loss, data_samples = self.det_head.loss_and_predict(
+            inputs, data_samples)
+        roi_losses = self.roi_head.loss(inputs, data_samples)
+        losses.update(det_loss)
+        losses.update(roi_losses)
+        return losses
 
     def predict(self, inputs: torch.Tensor,
                 data_samples: OptDetSampleList) -> OptDetSampleList:

projects/ABCNet/abcnet/model/two_stage_text_spotting.py

@@ -0,0 +1,12 @@
+# optimizer
+optim_wrapper = dict(
+    type='OptimWrapper',
+    optimizer=dict(type='SGD', lr=0.001, momentum=0.9, weight_decay=0.0001),
+    clip_grad=dict(type='value', clip_value=1))
+train_cfg = dict(type='EpochBasedTrainLoop', max_epochs=500, val_interval=20)
+val_cfg = dict(type='ValLoop')
+test_cfg = dict(type='TestLoop')
+# learning policy
+param_scheduler = [
+    dict(type='LinearLR', end=1000, start_factor=0.001, by_epoch=False),
+]

projects/ABCNet/config/_base_/schedules/schedule_sgd_500e.py

@@ -67,21 +67,37 @@
                 std=0.01,
                 bias=-4.59511985013459),  # -log((1-p)/p) where p=0.01
         ),
-        module_loss=None,
+        module_loss=dict(
+            type='ABCNetDetModuleLoss',
+            num_classes=num_classes,
+            strides=strides,
+            center_sampling=True,
+            center_sample_radius=1.5,
+            bbox_coder=bbox_coder,
+            norm_on_bbox=norm_on_bbox,
+            loss_cls=dict(
+                type='mmdet.FocalLoss',
+                use_sigmoid=use_sigmoid_cls,
+                gamma=2.0,
+                alpha=0.25,
+                loss_weight=1.0),
+            loss_bbox=dict(type='mmdet.GIoULoss', loss_weight=1.0),
+            loss_centerness=dict(
+                type='mmdet.CrossEntropyLoss',
+                use_sigmoid=True,
+                loss_weight=1.0)),
         postprocessor=dict(
             type='ABCNetDetPostprocessor',
-            # rescale_fields=['polygons', 'bboxes'],
             use_sigmoid_cls=use_sigmoid_cls,
             strides=[8, 16, 32, 64, 128],
             bbox_coder=dict(type='mmdet.DistancePointBBoxCoder'),
             with_bezier=True,
             test_cfg=dict(
-                # rescale_fields=['polygon', 'bboxes', 'bezier'],
                 nms_pre=1000,
                 nms=dict(type='nms', iou_threshold=0.5),
                 score_thr=0.3))),
     roi_head=dict(
-        type='OnlyRecRoIHead',
+        type='RecRoIHead',
         roi_extractor=dict(
             type='BezierRoIExtractor',
             roi_layer=dict(
@@ -95,7 +111,14 @@
             decoder=dict(
                 type='ABCNetRecDecoder',
                 dictionary=dictionary,
-                postprocessor=dict(type='AttentionPostprocessor'),
+                postprocessor=dict(
+                    type='AttentionPostprocessor',
+                    ignore_chars=['padding', 'unknown']),
+                module_loss=dict(
+                    type='CEModuleLoss',
+                    ignore_first_char=False,
+                    ignore_char=-1,
+                    reduction='mean'),
                 max_seq_len=25))),
     postprocessor=dict(
         type='ABCNetPostprocessor',
@@ -118,3 +141,32 @@
         type='PackTextDetInputs',
         meta_keys=('img_path', 'ori_shape', 'img_shape', 'scale_factor'))
 ]
+
+train_pipeline = [
+    dict(
+        type='LoadImageFromFile',
+        file_client_args=file_client_args,
+        color_type='color_ignore_orientation'),
+    dict(
+        type='LoadOCRAnnotations',
+        with_polygon=True,
+        with_bbox=True,
+        with_label=True,
+        with_text=True),
+    dict(type='RemoveIgnored'),
+    dict(type='RandomCrop', min_side_ratio=0.1),
+    dict(
+        type='RandomRotate',
+        max_angle=30,
+        pad_with_fixed_color=True,
+        use_canvas=True),
+    dict(
+        type='RandomChoiceResize',
+        scales=[(980, 2900), (1044, 2900), (1108, 2900), (1172, 2900),
+                (1236, 2900), (1300, 2900), (1364, 2900), (1428, 2900),
+                (1492, 2900)],
+        keep_ratio=True),
+    dict(
+        type='PackTextDetInputs',
+        meta_keys=('img_path', 'ori_shape', 'img_shape', 'scale_factor'))
+]

projects/ABCNet/config/abcnet/_base_abcnet-det_resnet50_fpn.py renamed to projects/ABCNet/config/abcnet/_base_abcnet_resnet50_fpn.py

@@ -0,0 +1,37 @@
+_base_ = [
+    '_base_abcnet_resnet50_fpn.py',
+    '../_base_/datasets/icdar2015.py',
+    '../_base_/default_runtime.py',
+    '../_base_/schedules/schedule_sgd_500e.py',
+]
+
+# dataset settings
+icdar2015_textspotting_train = _base_.icdar2015_textspotting_train
+icdar2015_textspotting_train.pipeline = _base_.train_pipeline
+icdar2015_textspotting_test = _base_.icdar2015_textspotting_test
+icdar2015_textspotting_test.pipeline = _base_.test_pipeline
+
+train_dataloader = dict(
+    batch_size=2,
+    num_workers=8,
+    persistent_workers=True,
+    sampler=dict(type='DefaultSampler', shuffle=True),
+    dataset=icdar2015_textspotting_train)
+
+val_dataloader = dict(
+    batch_size=1,
+    num_workers=4,
+    persistent_workers=True,
+    sampler=dict(type='DefaultSampler', shuffle=False),
+    dataset=icdar2015_textspotting_test)
+
+test_dataloader = val_dataloader
+
+val_cfg = dict(type='ValLoop')
+test_cfg = dict(type='TestLoop')
+
+custom_imports = dict(imports=['abcnet'], allow_failed_imports=False)
+
+load_from = 'https://download.openmmlab.com/mmocr/textspotting/abcnet/abcnet_resnet50_fpn_500e_icdar2015/abcnet_resnet50_fpn_pretrain-d060636c.pth'  # noqa
+
+find_unused_parameters = True