From c0492e02c73cc2ac6acbeb88da34b79d2783558f Mon Sep 17 00:00:00 2001 From: MissPenguin Date: Tue, 22 Jun 2021 03:32:00 +0000 Subject: [PATCH] refine --- ppocr/modeling/architectures/base_model.py | 10 ++-------- ppocr/modeling/heads/cls_head.py | 2 +- ppocr/modeling/heads/det_db_head.py | 2 +- ppocr/modeling/heads/det_east_head.py | 2 +- ppocr/modeling/heads/det_sast_head.py | 2 +- ppocr/modeling/heads/e2e_pg_head.py | 2 +- ppocr/modeling/heads/rec_ctc_head.py | 2 +- ppocr/modeling/heads/table_att_head.py | 4 ++-- tools/infer_table.py | 2 +- 9 files changed, 11 insertions(+), 17 deletions(-) diff --git a/ppocr/modeling/architectures/base_model.py b/ppocr/modeling/architectures/base_model.py index c1bdaaaf..03fbcee8 100644 --- a/ppocr/modeling/architectures/base_model.py +++ b/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) - else: - if mode == 'Eval' or mode == 'Test': - x = self.head(x, targets=data, mode=mode) - else: - x = self.head(x, targets=data) + x = self.head(x, targets=data) y["head_out"] = x if self.return_all_feats: return y diff --git a/ppocr/modeling/heads/cls_head.py b/ppocr/modeling/heads/cls_head.py index d9b78b84..91bfa615 100644 --- a/ppocr/modeling/heads/cls_head.py +++ b/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) diff --git a/ppocr/modeling/heads/det_db_head.py b/ppocr/modeling/heads/det_db_head.py index 83e7a5eb..f76cb34d 100644 --- a/ppocr/modeling/heads/det_db_head.py +++ b/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} diff --git a/ppocr/modeling/heads/det_east_head.py b/ppocr/modeling/heads/det_east_head.py index 9d0c3c4c..004eb5d7 100644 --- a/ppocr/modeling/heads/det_east_head.py +++ b/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) diff --git a/ppocr/modeling/heads/det_sast_head.py b/ppocr/modeling/heads/det_sast_head.py index 263b2867..7a88a2db 100644 --- a/ppocr/modeling/heads/det_sast_head.py +++ b/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) diff --git a/ppocr/modeling/heads/e2e_pg_head.py b/ppocr/modeling/heads/e2e_pg_head.py index 0da9de75..274e1cda 100644 --- a/ppocr/modeling/heads/e2e_pg_head.py +++ b/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) diff --git a/ppocr/modeling/heads/rec_ctc_head.py b/ppocr/modeling/heads/rec_ctc_head.py index 481f93e4..86e81677 100755 --- a/ppocr/modeling/heads/rec_ctc_head.py +++ b/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) diff --git a/ppocr/modeling/heads/table_att_head.py b/ppocr/modeling/heads/table_att_head.py index 61dacd35..155f036d 100644 --- a/ppocr/modeling/heads/table_att_head.py +++ b/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( diff --git a/tools/infer_table.py b/tools/infer_table.py index 450a83dd..f743d875 100644 --- a/tools/infer_table.py +++ b/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']