parser.add_argument("-e", "--epochs", type=int, default=100, help="Num of epochs to train.") args = parser.parse_args() print('args', args) train_dl = OCRDataLoader(args.train_annotation_paths, args.train_parse_funcs, args.image_width, args.table_path, args.batch_size, True) with open(r'E:\tsl_file\python_project\CRNN.tf2\example\table.txt', "r") as f: inv_table = [char.strip() for char in f] decoder = Decoder(inv_table) print("Num of training samples: {}".format(len(train_dl))) localtime = time.strftime("%Y-%m-%d-%H-%M-%S", time.localtime()) # saved_model_path = ("saved_models/{}/".format(localtime) + # "{epoch:03d}_{word_accuracy:.4f}.h5") if args.val_annotation_paths: val_dl = OCRDataLoader(args.val_annotation_paths, args.val_parse_funcs, args.image_width, args.table_path, args.batch_size) print("Num of val samples: {}".format(len(val_dl))) # saved_model_path = ("saved_models/{}/".format(localtime) + # "{epoch:03d}_{word_accuracy:.4f}_{val_word_accuracy:.4f}.h5") else: val_dl = lambda: None print("Start at {}".format(localtime))
def read_img_and_preprocess(path): img = tf.io.read_file(path) img = tf.io.decode_jpeg(img, channels=args.img_channels) img = tf.image.convert_image_dtype(img, tf.float32) img = tf.image.resize(img, (32, args.img_width)) return img if os.path.isdir(args.images): img_paths = os.listdir(args.images) img_paths = [os.path.join(args.images, path) for path in img_paths] imgs = list(map(read_img_and_preprocess, img_paths)) imgs = tf.stack(imgs) else: img_paths = [args.images] img = read_img_and_preprocess(args.images) imgs = tf.expand_dims(img, 0) with open(args.table_path, 'r') as f: inv_table = [char.strip() for char in f] model = keras.models.load_model(args.model, compile=False) decoder = Decoder(inv_table) y_pred = model.predict(imgs) for path, g_pred, b_pred in zip(img_paths, decoder.decode(y_pred, method='greedy'), decoder.decode(y_pred, method='beam_search')): print('Path: {}, greedy: {}, beam search: {}'.format(path, g_pred, b_pred))