refine

ppocr/modeling/architectures/base_model.py

@@ -69,7 +69,7 @@ class BaseModel(nn.Layer):
 
         self.return_all_feats = config.get("return_all_feats", False)
 
-    def forward(self, x, data=None, mode='Train'):
+    def forward(self, x, data=None):
         y = dict()
         if self.use_transform:
             x = self.transform(x)
@@ -78,13 +78,7 @@ class BaseModel(nn.Layer):
         if self.use_neck:
             x = self.neck(x)
         y["neck_out"] = x
-        if data is None:
+        x = self.head(x, targets=data)
-            x = self.head(x)
-        else:
-            if mode == 'Eval' or mode == 'Test':
-                x = self.head(x, targets=data, mode=mode)
-            else:
-                x = self.head(x, targets=data)
         y["head_out"] = x
         if self.return_all_feats:
             return y

ppocr/modeling/heads/cls_head.py

@@ -43,7 +43,7 @@ class ClsHead(nn.Layer):
                 initializer=nn.initializer.Uniform(-stdv, stdv)),
             bias_attr=ParamAttr(name="fc_0.b_0"), )
 
-    def forward(self, x):
+    def forward(self, x, targets=None):
         x = self.pool(x)
         x = paddle.reshape(x, shape=[x.shape[0], x.shape[1]])
         x = self.fc(x)

ppocr/modeling/heads/det_db_head.py

@@ -106,7 +106,7 @@ class DBHead(nn.Layer):
     def step_function(self, x, y):
         return paddle.reciprocal(1 + paddle.exp(-self.k * (x - y)))
 
-    def forward(self, x):
+    def forward(self, x, targets=None):
         shrink_maps = self.binarize(x)
         if not self.training:
             return {'maps': shrink_maps}

ppocr/modeling/heads/det_east_head.py

@@ -109,7 +109,7 @@ class EASTHead(nn.Layer):
             act=None,
             name="f_geo")
 
-    def forward(self, x):
+    def forward(self, x, targets=None):
         f_det = self.det_conv1(x)
         f_det = self.det_conv2(f_det)
         f_score = self.score_conv(f_det)

ppocr/modeling/heads/det_sast_head.py

@@ -116,7 +116,7 @@ class SASTHead(nn.Layer):
         self.head1 = SAST_Header1(in_channels)
         self.head2 = SAST_Header2(in_channels)
 
-    def forward(self, x):
+    def forward(self, x, targets=None):
         f_score, f_border = self.head1(x)
         f_tvo, f_tco = self.head2(x)
 

ppocr/modeling/heads/e2e_pg_head.py

@@ -220,7 +220,7 @@ class PGHead(nn.Layer):
             weight_attr=ParamAttr(name="conv_f_direc{}".format(4)),
             bias_attr=False)
 
-    def forward(self, x):
+    def forward(self, x, targets=None):
         f_score = self.conv_f_score1(x)
         f_score = self.conv_f_score2(f_score)
         f_score = self.conv_f_score3(f_score)

ppocr/modeling/heads/rec_ctc_head.py

@@ -44,7 +44,7 @@ class CTCHead(nn.Layer):
             bias_attr=bias_attr)
         self.out_channels = out_channels
 
-    def forward(self, x, labels=None):
+    def forward(self, x, targets=None):
         predicts = self.fc(x)
         if not self.training:
             predicts = F.softmax(predicts, axis=2)

ppocr/modeling/heads/table_att_head.py

@@ -53,7 +53,7 @@ class TableAttentionHead(nn.Layer):
         input_ont_hot = F.one_hot(input_char, onehot_dim)
         return input_ont_hot
 
-    def forward(self, inputs, targets=None, mode='Train'):
+    def forward(self, inputs, targets=None):
         # if and else branch are both needed when you want to assign a variable
         # if you modify the var in just one branch, then the modification will not work.
         fea = inputs[-1]
@@ -67,7 +67,7 @@ class TableAttentionHead(nn.Layer):
         
         hidden = paddle.zeros((batch_size, self.hidden_size))
         output_hiddens = []
-        if mode == 'Train' and targets is not None:
+        if self.training and targets is not None:
             structure = targets[0]
             for i in range(self.max_elem_length+1):
                 elem_onehots = self._char_to_onehot(

tools/infer_table.py

@@ -81,7 +81,7 @@ def main(config, device, logger, vdl_writer):
         batch = transform(data, ops)
         images = np.expand_dims(batch[0], axis=0)
         images = paddle.to_tensor(images)
-        preds = model(images, data=None, mode='Test')
+        preds = model(images)
         post_result = post_process_class(preds)
         res_html_code = post_result['res_html_code']
         res_loc = post_result['res_loc']