Merge pull request #196 from SamitHuang/main

DBNet output changed  in inference mode for faster speed - only binary map is output

Author: SamitHuang

mindocr/models/heads/det_db_head.py

@@ -7,7 +7,7 @@ def __init__(self, in_channels: int, k=50, adaptive=False, bias=False, weight_in
         self.adaptive = adaptive
 
         self.segm = self._init_heatmap(in_channels, in_channels // 4, weight_init, bias)
-        if adaptive:
+        if self.adaptive:
             self.thresh = self._init_heatmap(in_channels, in_channels // 4, weight_init, bias)
             self.k = k
             self.diff_bin = nn.Sigmoid()
@@ -29,11 +29,20 @@ def _init_heatmap(in_channels, inter_channels, weight_init, bias):
         ])
 
     def construct(self, features):
+        '''
+        Args:
+            features (Tensor): features output by backbone
+        Returns:
+            pred (dict):
+                - binary: predicted binary map
+                - thresh: predicted threshold  map, only used if adaptive is True in training
+                - thres_binary: differentiable binary map, only if adaptive is True in training
+        '''
         pred = {'binary': self.segm(features)}
 
-        if self.adaptive:
+        if self.adaptive and self.training:
+            # only use binary map to derive polygons in inference
             pred['thresh'] = self.thresh(features)
             pred['thresh_binary'] = self.diff_bin(
                 self.k * (pred['binary'] - pred['thresh']))  # Differentiable Binarization
-
         return pred

mindocr/postprocess/det_postprocess.py

@@ -30,8 +30,12 @@ def __call__(self, pred):
                 polygons: np.ndarray of shape (N, K, 4, 2) for the polygons of objective regions if region_type is 'quad'
                 scores: np.ndarray of shape (N, K), score for each box
         """
-        pred = (pred[self._name] if isinstance(pred, dict) else pred[self._names[self._name]]).squeeze(1)
-        pred = pred.asnumpy()
+        if isinstance(pred, dict):
+            pred = pred[self._name]
+        elif isinstance(pred, tuple):
+            pred = pred[self._names[self._name]]
+        pred = pred.squeeze(1).asnumpy()
+        
         segmentation = pred >= self._binary_thresh
 
         # FIXME: dest_size is supposed to be the original image shape (pred.shape -> batch['shape'])