mirror of
https://github.com/PaddlePaddle/PaddleOCR.git
synced 2025-10-31 09:49:30 +00:00
78 lines
3.3 KiB
Python
78 lines
3.3 KiB
Python
# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
import os
|
|
import numpy as np
|
|
import cv2
|
|
|
|
from utils.config import ArgsParser, load_config, override_config
|
|
from utils.logging import get_logger
|
|
from engine import style_samplers, corpus_generators, text_drawers, predictors, writers
|
|
|
|
|
|
class ImageSynthesiser(object):
|
|
def __init__(self):
|
|
self.FLAGS = ArgsParser().parse_args()
|
|
self.config = load_config(self.FLAGS.config)
|
|
self.config = override_config(self.config, options=self.FLAGS.override)
|
|
self.output_dir = self.config["Global"]["output_dir"]
|
|
if not os.path.exists(self.output_dir):
|
|
os.mkdir(self.output_dir)
|
|
self.logger = get_logger(
|
|
log_file='{}/predict.log'.format(self.output_dir))
|
|
|
|
self.text_drawer = text_drawers.StdTextDrawer(self.config)
|
|
|
|
predictor_method = self.config["Predictor"]["method"]
|
|
assert predictor_method is not None
|
|
self.predictor = getattr(predictors, predictor_method)(self.config)
|
|
|
|
def synth_image(self, corpus, style_input, language="en"):
|
|
corpus_list, text_input_list = self.text_drawer.draw_text(
|
|
corpus, language, style_input_width=style_input.shape[1])
|
|
synth_result = self.predictor.predict(style_input, text_input_list)
|
|
return synth_result
|
|
|
|
|
|
class DatasetSynthesiser(ImageSynthesiser):
|
|
def __init__(self):
|
|
super(DatasetSynthesiser, self).__init__()
|
|
self.tag = self.FLAGS.tag
|
|
self.output_num = self.config["Global"]["output_num"]
|
|
corpus_generator_method = self.config["CorpusGenerator"]["method"]
|
|
self.corpus_generator = getattr(corpus_generators,
|
|
corpus_generator_method)(self.config)
|
|
|
|
style_sampler_method = self.config["StyleSampler"]["method"]
|
|
assert style_sampler_method is not None
|
|
self.style_sampler = style_samplers.DatasetSampler(self.config)
|
|
self.writer = writers.SimpleWriter(self.config, self.tag)
|
|
|
|
def synth_dataset(self):
|
|
for i in range(self.output_num):
|
|
style_data = self.style_sampler.sample()
|
|
style_input = style_data["image"]
|
|
corpus_language, text_input_label = self.corpus_generator.generate()
|
|
text_input_label_list, text_input_list = self.text_drawer.draw_text(
|
|
text_input_label,
|
|
corpus_language,
|
|
style_input_width=style_input.shape[1])
|
|
|
|
text_input_label = "".join(text_input_label_list)
|
|
|
|
synth_result = self.predictor.predict(style_input, text_input_list)
|
|
fake_fusion = synth_result["fake_fusion"]
|
|
self.writer.save_image(fake_fusion, text_input_label)
|
|
self.writer.save_label()
|
|
self.writer.merge_label()
|