Merge pull request #203 from Rhcsky/master

Correct demo.py error when running with CTC
This commit is contained in:
gwkrsrch 2020-07-30 12:06:32 +09:00 committed by GitHub
commit 5e277f057b
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 2 additions and 2 deletions

View File

@ -57,8 +57,8 @@ def demo(opt):
# Select max probabilty (greedy decoding) then decode index to character # Select max probabilty (greedy decoding) then decode index to character
preds_size = torch.IntTensor([preds.size(1)] * batch_size) preds_size = torch.IntTensor([preds.size(1)] * batch_size)
_, preds_index = preds.max(2) _, preds_index = preds.max(2)
preds_index = preds_index.view(-1) # preds_index = preds_index.view(-1)
preds_str = converter.decode(preds_index.data, preds_size.data) preds_str = converter.decode(preds_index, preds_size)
else: else:
preds = model(image, text_for_pred, is_train=False) preds = model(image, text_for_pred, is_train=False)