Spaces:
Runtime error
Runtime error
| # 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 numpy as np | |
| import cv2 | |
| import math | |
| import paddle | |
| from arch import style_text_rec | |
| from utils.sys_funcs import check_gpu | |
| from utils.logging import get_logger | |
| class StyleTextRecPredictor(object): | |
| def __init__(self, config): | |
| algorithm = config['Predictor']['algorithm'] | |
| assert algorithm in ["StyleTextRec" | |
| ], "Generator {} not supported.".format(algorithm) | |
| use_gpu = config["Global"]['use_gpu'] | |
| check_gpu(use_gpu) | |
| paddle.set_device('gpu' if use_gpu else 'cpu') | |
| self.logger = get_logger() | |
| self.generator = getattr(style_text_rec, algorithm)(config) | |
| self.height = config["Global"]["image_height"] | |
| self.width = config["Global"]["image_width"] | |
| self.scale = config["Predictor"]["scale"] | |
| self.mean = config["Predictor"]["mean"] | |
| self.std = config["Predictor"]["std"] | |
| self.expand_result = config["Predictor"]["expand_result"] | |
| def reshape_to_same_height(self, img_list): | |
| h = img_list[0].shape[0] | |
| for idx in range(1, len(img_list)): | |
| new_w = round(1.0 * img_list[idx].shape[1] / | |
| img_list[idx].shape[0] * h) | |
| img_list[idx] = cv2.resize(img_list[idx], (new_w, h)) | |
| return img_list | |
| def predict_single_image(self, style_input, text_input): | |
| style_input = self.rep_style_input(style_input, text_input) | |
| tensor_style_input = self.preprocess(style_input) | |
| tensor_text_input = self.preprocess(text_input) | |
| style_text_result = self.generator.forward(tensor_style_input, | |
| tensor_text_input) | |
| fake_fusion = self.postprocess(style_text_result["fake_fusion"]) | |
| fake_text = self.postprocess(style_text_result["fake_text"]) | |
| fake_sk = self.postprocess(style_text_result["fake_sk"]) | |
| fake_bg = self.postprocess(style_text_result["fake_bg"]) | |
| bbox = self.get_text_boundary(fake_text) | |
| if bbox: | |
| left, right, top, bottom = bbox | |
| fake_fusion = fake_fusion[top:bottom, left:right, :] | |
| fake_text = fake_text[top:bottom, left:right, :] | |
| fake_sk = fake_sk[top:bottom, left:right, :] | |
| fake_bg = fake_bg[top:bottom, left:right, :] | |
| # fake_fusion = self.crop_by_text(img_fake_fusion, img_fake_text) | |
| return { | |
| "fake_fusion": fake_fusion, | |
| "fake_text": fake_text, | |
| "fake_sk": fake_sk, | |
| "fake_bg": fake_bg, | |
| } | |
| def predict(self, style_input, text_input_list): | |
| if not isinstance(text_input_list, (tuple, list)): | |
| return self.predict_single_image(style_input, text_input_list) | |
| synth_result_list = [] | |
| for text_input in text_input_list: | |
| synth_result = self.predict_single_image(style_input, text_input) | |
| synth_result_list.append(synth_result) | |
| for key in synth_result: | |
| res = [r[key] for r in synth_result_list] | |
| res = self.reshape_to_same_height(res) | |
| synth_result[key] = np.concatenate(res, axis=1) | |
| return synth_result | |
| def preprocess(self, img): | |
| img = (img.astype('float32') * self.scale - self.mean) / self.std | |
| img_height, img_width, channel = img.shape | |
| assert channel == 3, "Please use an rgb image." | |
| ratio = img_width / float(img_height) | |
| if math.ceil(self.height * ratio) > self.width: | |
| resized_w = self.width | |
| else: | |
| resized_w = int(math.ceil(self.height * ratio)) | |
| img = cv2.resize(img, (resized_w, self.height)) | |
| new_img = np.zeros([self.height, self.width, 3]).astype('float32') | |
| new_img[:, 0:resized_w, :] = img | |
| img = new_img.transpose((2, 0, 1)) | |
| img = img[np.newaxis, :, :, :] | |
| return paddle.to_tensor(img) | |
| def postprocess(self, tensor): | |
| img = tensor.numpy()[0] | |
| img = img.transpose((1, 2, 0)) | |
| img = (img * self.std + self.mean) / self.scale | |
| img = np.maximum(img, 0.0) | |
| img = np.minimum(img, 255.0) | |
| img = img.astype('uint8') | |
| return img | |
| def rep_style_input(self, style_input, text_input): | |
| rep_num = int(1.2 * (text_input.shape[1] / text_input.shape[0]) / | |
| (style_input.shape[1] / style_input.shape[0])) + 1 | |
| style_input = np.tile(style_input, reps=[1, rep_num, 1]) | |
| max_width = int(self.width / self.height * style_input.shape[0]) | |
| style_input = style_input[:, :max_width, :] | |
| return style_input | |
| def get_text_boundary(self, text_img): | |
| img_height = text_img.shape[0] | |
| img_width = text_img.shape[1] | |
| bounder = 3 | |
| text_canny_img = cv2.Canny(text_img, 10, 20) | |
| edge_num_h = text_canny_img.sum(axis=0) | |
| no_zero_list_h = np.where(edge_num_h > 0)[0] | |
| edge_num_w = text_canny_img.sum(axis=1) | |
| no_zero_list_w = np.where(edge_num_w > 0)[0] | |
| if len(no_zero_list_h) == 0 or len(no_zero_list_w) == 0: | |
| return None | |
| left = max(no_zero_list_h[0] - bounder, 0) | |
| right = min(no_zero_list_h[-1] + bounder, img_width) | |
| top = max(no_zero_list_w[0] - bounder, 0) | |
| bottom = min(no_zero_list_w[-1] + bounder, img_height) | |
| return [left, right, top, bottom] | |