Update test.py

2019-05-08 15:16:45 +09:00 · 2019-05-08 15:16:45 +09:00 · 3a1ebfadc3
parent 3772d75f40
commit 3a1ebfadc3
1 changed files with 4 additions and 0 deletions
--- a/test.py
+++ b/test.py
@ -135,6 +135,10 @@ def test(opt):
        converter = CTCLabelConverter(opt.character)
    else:
        converter = AttnLabelConverter(opt.character)
+        
+    if opt.rgb:
+        opt.input_channel = 3
+    
    opt.num_class = len(converter.character)
    model = Model(opt)
    print('model input parameters', opt.imgH, opt.imgW, opt.num_fiducial, opt.input_channel, opt.output_channel,