abcnetv2 train

Author: Harold-lkk

projects/ABCNet/abcnet/model/__init__.py

@@ -12,6 +12,7 @@
 from .bifpn import BiFPN
 from .coordinate_head import CoordinateHead
 from .rec_roi_head import RecRoIHead
+from .task_utils import *  # noqa: F401,F403
 
 __all__ = [
     'ABCNetDetHead', 'ABCNetDetPostprocessor', 'ABCNetRecBackbone',

projects/ABCNet/abcnet/model/abcnet_det_head.py

@@ -181,8 +181,8 @@ def forward_single(self, x, scale, stride):
         # float to avoid overflow when enabling FP16
         if self.use_scale:
             bbox_pred = scale(bbox_pred).float()
-        else:
-            bbox_pred = bbox_pred.float()
+        # else:
+        #     bbox_pred = bbox_pred.float()
         if self.norm_on_bbox:
             # bbox_pred needed for gradient computation has been modified
             # by F.relu(bbox_pred) when run with PyTorch 1.10. So replace

projects/ABCNet/abcnet/model/abcnet_det_postprocessor.py

@@ -216,8 +216,8 @@ def __call__(self, pred_results, data_samples, training: bool = False):
         Returns:
             list[TextDetDataSample]: Batch of post-processed datasamples.
         """
-        if training:
-            return data_samples
+        # if training:
+        #     return data_samples
         cfg = self.train_cfg if training else self.test_cfg
         if cfg is None:
             cfg = {}

projects/ABCNet/abcnet/model/bezier_roi_extractor.py

@@ -88,7 +88,7 @@ def forward(self, feats: Tuple[Tensor],
         # convert fp32 to fp16 when amp is on
         rois = rois.type_as(feats[0])
         out_size = self.roi_layers[0].output_size
-        feats = feats[:3]
+        # feats = feats[:3]
         num_levels = len(feats)
         roi_feats = feats[0].new_zeros(
             rois.size(0), self.out_channels, *out_size)

projects/ABCNet/abcnet/model/bifpn.py

@@ -170,10 +170,12 @@ def __init__(self,
         self.bifpn_convs = nn.ModuleList()
         # weighted
         self.weight_two_nodes = nn.Parameter(
-            torch.Tensor(2, levels).fill_(init))
+            torch.Tensor(2, levels).fill_(init), requires_grad=True)
+
         self.weight_three_nodes = nn.Parameter(
-            torch.Tensor(3, levels - 2).fill_(init))
-        self.relu = nn.ReLU()
+            torch.Tensor(3, levels - 2).fill_(init), requires_grad=True)
+
+        # self.relu = nn.ReLU(inplace=False)
         for _ in range(2):
             for _ in range(self.levels - 1):  # 1,2,3
                 fpn_conv = nn.Sequential(
@@ -193,9 +195,10 @@ def forward(self, inputs):
         # build top-down and down-top path with stack
         levels = self.levels
         # w relu
-        w1 = self.relu(self.weight_two_nodes)
-        w1 /= torch.sum(w1, dim=0) + self.eps  # normalize
-        w2 = self.relu(self.weight_three_nodes)
+
+        _w1 = F.relu(self.weight_two_nodes)
+        w1 = _w1 / (torch.sum(_w1, dim=0) + self.eps)  # normalize
+        w2 = F.relu(self.weight_three_nodes)
         # w2 /= torch.sum(w2, dim=0) + self.eps  # normalize
         # build top-down
         idx_bifpn = 0

projects/ABCNet/abcnet/model/rec_roi_head.py

@@ -1,5 +1,5 @@
 # Copyright (c) OpenMMLab. All rights reserved.
-from typing import Tuple
+from typing import Optional, Sequence, Tuple
 
 from mmengine.structures import LabelData
 from torch import Tensor
@@ -15,14 +15,19 @@ class RecRoIHead(BaseRoIHead):
     """Simplest base roi head including one bbox head and one mask head."""
 
     def __init__(self,
-                 neck=None,
+                 inputs_indices: Optional[Sequence] = None,
+                 neck: OptMultiConfig = None,
+                 assigner: OptMultiConfig = None,
                  sampler: OptMultiConfig = None,
                  roi_extractor: OptMultiConfig = None,
                  rec_head: OptMultiConfig = None,
                  init_cfg=None):
         super().__init__(init_cfg)
-        if sampler is not None:
-            self.sampler = TASK_UTILS.build(sampler)
+        self.inputs_indices = inputs_indices
+        self.assigner = assigner
+        if assigner is not None:
+            self.assigner = TASK_UTILS.build(assigner)
+        self.sampler = TASK_UTILS.build(sampler)
         if neck is not None:
             self.neck = MODELS.build(neck)
         self.roi_extractor = MODELS.build(roi_extractor)
@@ -43,11 +48,39 @@ 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
-        ]
+
+        if self.inputs_indices is not None:
+            inputs = [inputs[i] for i in self.inputs_indices]
+        # proposals = [
+        #     ds.gt_instances[~ds.gt_instances.ignored] for ds in data_samples
+        # ]
+        proposals = list()
+        for ds in data_samples:
+            pred_instances = ds.pred_instances
+            gt_instances = ds.gt_instances
+            # # assign
+            # gt_beziers = gt_instances.beziers
+            # pred_beziers = pred_instances.beziers
+            # assign_index = [
+            #     int(
+            #         torch.argmin(
+            #             torch.abs(gt_beziers - pred_beziers[i]).sum(dim=1)))
+            #     for i in range(len(pred_beziers))
+            # ]
+            # proposal = InstanceData()
+            # proposal.texts = gt_instances.texts + gt_instances[
+            #     assign_index].texts
+            # proposal.beziers = torch.cat(
+            #     [gt_instances.beziers, pred_instances.beziers], dim=0)
+            if self.assigner:
+                gt_instances, pred_instances = self.assigner.assign(
+                    gt_instances, pred_instances)
+            proposal = self.sampler.sample(gt_instances, pred_instances)
+            proposals.append(proposal)
 
         proposals = [p for p in proposals if len(p) > 0]
+        if hasattr(self, 'neck') and self.neck is not None:
+            inputs = self.neck(inputs)
         bbox_feats = self.roi_extractor(inputs, proposals)
         rec_data_samples = [
             TextRecogDataSample(gt_text=LabelData(item=text))
@@ -57,6 +90,7 @@ def loss(self, inputs: Tuple[Tensor], data_samples: DetSampleList) -> dict:
 
     def predict(self, inputs: Tuple[Tensor],
                 data_samples: DetSampleList) -> RecSampleList:
+        inputs = inputs[:3]
         if hasattr(self, 'neck') and self.neck is not None:
             inputs = self.neck(inputs)
         pred_instances = [ds.pred_instances for ds in data_samples]

projects/ABCNet/abcnet/model/task_utils/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+from .assigner import L1DistanceAssigner
+from .sampler import ConcatSampler, OnlyGTSampler
+
+__all__ = ['L1DistanceAssigner', 'ConcatSampler', 'OnlyGTSampler']

projects/ABCNet/abcnet/model/task_utils/assigner.py

@@ -0,0 +1,20 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import torch
+
+from mmocr.registry import TASK_UTILS
+
+
+@TASK_UTILS.register_module()
+class L1DistanceAssigner:
+
+    def assign(self, gt_instances, pred_instances):
+        gt_beziers = gt_instances.beziers
+        pred_beziers = pred_instances.beziers
+        assign_index = [
+            int(
+                torch.argmin(
+                    torch.abs(gt_beziers - pred_beziers[i]).sum(dim=1)))
+            for i in range(len(pred_beziers))
+        ]
+        pred_instances.assign_index = assign_index
+        return gt_instances, pred_instances

projects/ABCNet/abcnet/model/task_utils/sampler.py

@@ -0,0 +1,26 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import torch
+from mmengine.structures import InstanceData
+
+from mmocr.registry import TASK_UTILS
+
+
+@TASK_UTILS.register_module()
+class ConcatSampler:
+
+    def sample(self, gt_instances, pred_instances):
+        if len(pred_instances) == 0:
+            return gt_instances
+        proposals = InstanceData()
+        proposals.texts = gt_instances.texts + gt_instances[
+            pred_instances.assign_index].texts
+        proposals.beziers = torch.cat(
+            [gt_instances.beziers, pred_instances.beziers], dim=0)
+        return proposals
+
+
+@TASK_UTILS.register_module()
+class OnlyGTSampler:
+
+    def sample(self, gt_instances, pred_instances):
+        return gt_instances[~gt_instances.ignored]

projects/ABCNet/config/_base_/default_runtime.py

@@ -1,6 +1,6 @@
 default_scope = 'mmocr'
 env_cfg = dict(
-    cudnn_benchmark=True,
+    cudnn_benchmark=False,
     mp_cfg=dict(mp_start_method='fork', opencv_num_threads=0),
     dist_cfg=dict(backend='nccl'),
 )