mirrors/mmocr
1
0
Fork 0
mirror of https://github.com/open-mmlab/mmocr.git synced 2025-06-03 21:54:47 +08:00
Code Issues Projects Releases Wiki Activity

[Fix]: Fix dim mismatch in batch test/inference of DBNet (#383)

Browse Source 
Modification
1.change the dbnet output dim
2.change the input dim for DBloss
This commit is contained in:
liukuikun 2021-07-24 09:14:03 +08:00 committed by GitHub
parent 02d657d141
commit 200dfe5fe2
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 2 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
.gitignore vendored
View File

@ -138,3 +138,4 @@ workspace.code-workspace
results results
mmocr/core/font.TTF mmocr/core/font.TTF
workdirs/ workdirs/
.history/

2
mmocr/models/textdet/dense_heads/db_head.py
View File

@ -66,7 +66,7 @@ class DBHead(HeadMixin, BaseModule):
thr_map = self.threshold(inputs) thr_map = self.threshold(inputs)
binary_map = self.diff_binarize(prob_map, thr_map, k=50) binary_map = self.diff_binarize(prob_map, thr_map, k=50)
outputs = torch.cat((prob_map, thr_map, binary_map), dim=1) outputs = torch.cat((prob_map, thr_map, binary_map), dim=1)
return (outputs, ) return outputs
def _init_thr(self, inner_channels, bias=False): def _init_thr(self, inner_channels, bias=False):
in_channels = inner_channels in_channels = inner_channels

2
mmocr/models/textdet/losses/db_loss.py
View File

@ -133,8 +133,6 @@ class DBLoss(nn.Module):
assert isinstance(gt_thr, list) assert isinstance(gt_thr, list)
assert isinstance(gt_thr_mask, list) assert isinstance(gt_thr_mask, list)
preds = preds[0]
pred_prob = preds[:, 0, :, :] pred_prob = preds[:, 0, :, :]
pred_thr = preds[:, 1, :, :] pred_thr = preds[:, 1, :, :]
pred_db = preds[:, 2, :, :] pred_db = preds[:, 2, :, :]