diff --git a/ppocr/modeling/architectures/base_model.py b/ppocr/modeling/architectures/base_model.py index c1bdaaafbf7f973df4dcdce18024966171f1e0b4..03fbcee8465df9c8bb7845ea62fc0ac04917caa0 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 d9b78b841b3c31ea349cfbf4e767328b12f39aa7..91bfa615a8206b5ec0f993429ccae990a05d0b9b 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 83e7a5ebfe131ed209b7dd2d4b5a324605be8370..f76cb34d37af7d81b5e628d06c1a4cfe126f8bb4 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 9d0c3c4cf83adb018fcc368374cbe305658e07a9..004eb5d7bb9a134d1a84f980e37e5336dc43a29a 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 263b28672299e733369938fa03952dca7685fabe..7a88a2db6c29c8c4fa1ee94d27bd0701cdbc90f8 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 0da9de7580a0ceb473f971b2246c966497026a5d..274e1cdac5172f45590c9f7d7b50522c74db6750 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 481f93e47e58f8267b23e632df1a1e80733d5944..86e81677544ebfd8a0852f5791d40f93d63d06a1 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 61dacd355cd67c7000487268ecfa717361714273..155f036d15673135eae9e5ee493648603609535d 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 450a83dd0d6c45878f0ab2b7c42a7ee668a50903..f743d87540f7fd64157a808db156c9f62a042d9c 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']