from pathlib import Path
from typing import Dict, List, Union, Tuple
from omegaconf import OmegaConf
import numpy as np
import torch
from torch import nn
from PIL import Image, ImageDraw, ImageFont
import models
GENERATOR_PREFIX = "networks.g."
WHITE = 255
EXAMPLE_CHARACTERS = ['A', 'B', 'C', 'D', 'E']
class InferenceServicer:
def __init__(self, hp, checkpoint_path, content_image_dir, imsize=64, gpu_id='0') -> None:
self.hp = hp
self.imsize = imsize
if gpu_id is None:
self.device = torch.device(f'cuda:0') if torch.cuda.is_available() else 'cpu'
else:
self.device = torch.device(f'cuda:{gpu_id}')
model_config = self.hp.models.G
self.model: nn.Module = models.Generator(model_config)
# Load Generator model weight
model_state_dict_pl = torch.load(checkpoint_path, map_location='cpu')
generator_state_dict = self.convert_generator_state_dict(model_state_dict_pl)
self.model.load_state_dict(generator_state_dict)
self.model.to(device=self.device)
self.model.eval()
# Setting Content font files
self.content_character_dict = self.load_content_character_dict(Path(content_image_dir))
@staticmethod
def convert_generator_state_dict(model_state_dict_pl):
generator_prefix = GENERATOR_PREFIX
generator_state_dict = {}
for module_name, module_state in model_state_dict_pl['state_dict'].items():
if module_name.startswith(generator_prefix):
generator_state_dict[module_name[len(generator_prefix):]] = module_state
return generator_state_dict
@staticmethod
def load_content_character_dict(content_image_dir: Path) -> Dict[str, Path]:
content_character_dict = {}
for filepath in content_image_dir.glob("**/*.png"):
content_character_dict[filepath.stem] = filepath
return content_character_dict
@staticmethod
def center_align(bg_img: Image.Image, item_img: Image.Image, fit=False) -> Image.Image:
bg_img = bg_img.copy()
item_img = item_img.copy()
item_w, item_h = item_img.size
W, H = bg_img.size
if fit:
item_ratio = item_w / item_h
bg_ratio = W / H
if bg_ratio > item_ratio:
# height fitting
resize_ratio = H / item_h
else:
# width fitting
resize_ratio = W / item_w
item_img = item_img.resize((int(item_w * resize_ratio), int(item_h * resize_ratio)))
item_w, item_h = item_img.size
bg_img.paste(item_img, ((W - item_w) // 2, (H - item_h) // 2))
return bg_img
def set_image(self, image: Union[Path, Image.Image]) -> Image.Image:
if isinstance(image, (str, Path)):
image = Image.open(image)
assert isinstance(image, Image.Image)
bg_img = Image.new('RGB', (self.imsize, self.imsize), color='white')
blend_img = self.center_align(bg_img, image, fit=True)
return blend_img
@staticmethod
def pil_image_to_array(blend_img: Image.Image) -> np.ndarray:
normalized_array = np.mean(np.array(blend_img, dtype=np.float32), axis=-1) / WHITE # L-only image normalized to [0, 1]
return normalized_array
def get_images_from_fontfile(self, font_file_path: Path, imgmode: str = 'RGB', position: tuple = (0, 0), font_size: int = 128, padding: int = 100) -> List[Image.Image]:
imagefont = ImageFont.truetype(str(font_file_path), size=font_size)
example_characters = EXAMPLE_CHARACTERS
font_images: List[Image.Image] = []
for character in example_characters:
x, y, _, _ = imagefont.getbbox(character)
img = Image.new(imgmode, (x + padding, y + padding), color='white')
draw = ImageDraw.Draw(img)
# bbox = draw.textbbox((0,0), character, font=imagefont)
# w = bbox[2] - bbox[0]
# h = bbox[3] - bbox[1]
w, h = draw.textsize(character, font=imagefont)
img = Image.new(imgmode, (w + padding, h + padding), color='white')
draw = ImageDraw.Draw(img)
draw.text(position, text=character, font=imagefont, fill='black')
img = img.convert(imgmode)
font_images.append(img)
return font_images
@staticmethod
def get_hex_from_char(char: str) -> str:
assert len(char) == 1
return f"{ord(char):04X}".upper() # 4-digit hex string
@torch.no_grad()
def inference(self, content_char: str, style_font: Union[str, Path]) -> Tuple[Image.Image, List[Image.Image], Image.Image]:
assert len(content_char) > 0
content_char = content_char[:1] # only get the first character if the length > 1
char_hex = self.get_hex_from_char(content_char)
if char_hex not in self.content_character_dict:
raise ValueError(f"The character {content_char} (hex: {char_hex}) is not supported in this model!")
content_image = self.set_image(self.content_character_dict[char_hex])
style_images: List[Image.Image] = self.get_images_from_fontfile(Path(style_font))
style_images: List[Image.Image] = [self.set_image(image) for image in style_images]
content_image_array = self.pil_image_to_array(content_image)[np.newaxis, np.newaxis, ...] # 1 x C(=1) x H x W
style_images_array: np.ndarray = np.array([self.pil_image_to_array(image) for image in style_images])[np.newaxis, ...] # 1 x C(=5, # shots) x H x W, k-shots goes to batch
content_input_tensor = torch.from_numpy(content_image_array).to(self.device)
style_input_tensor = torch.from_numpy(style_images_array).to(self.device)
generated_images: torch.Tensor = self.model((content_input_tensor, style_input_tensor))
generated_images = torch.clip(generated_images, 0, 1)
assert generated_images.size(0) == 1
generated_image_numpy = (generated_images[0].cpu().numpy() * 255).astype(np.uint8)[0, ...] # H x W
return content_image, style_images, Image.fromarray(generated_image_numpy, mode='L')
if __name__ == '__main__':
hp = OmegaConf.load("config/models/google-font.yaml")
checkpoint_path = "epoch=199-step=257400.ckpt"
content_image_dir = "../DATA/NotoSans"
servicer = InferenceServicer(hp, checkpoint_path, content_image_dir)
style_font = "example_fonts/MaShanZheng-Regular.ttf"
content_image, style_images, result = servicer.inference("7", style_font)
content_image.save("result_content.png")
for idx, style_image in enumerate(style_images):
style_image.save(f"result_style_{idx:02d}.png")
result.save("result_generated.png")