Spaces:
Sleeping
Sleeping
File size: 4,074 Bytes
3a5d9c5 1ed13ab 3a5d9c5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 |
import os
import sys
import time
import glob
import yaml
import random
import numpy as np
import gradio as gr
from PIL import Image, ImageDraw, ImageFilter
def get_template(path):
# RGB to HSV
img = Image.open(path).convert('HSV')
# pick the H channel
img = img.split()[0]
# divide by 16
img = np.array(img) + 8
img = Image.fromarray(img // 16)
return img
# resize via height, keep the aspect ratio
def resize_height(img, height, nearest=False):
width, old_height = img.size
new_width = int(width * height / old_height)
return img.resize((new_width, height), Image.NEAREST if nearest else Image.BICUBIC), new_width
# load config yaml file
def load_config(path):
with open(path, "r") as f:
config = yaml.load(f, Loader=yaml.FullLoader)
return config
def parse_config(config):
def random_choice(char):
if char == "#":
return random.choice(config["plateNums"] + config["plateChars"])
elif char == "A":
return random.choice(config["plateChars"])
elif char == "0":
return random.choice(config["plateNums"])
elif char == "@":
return random.choice(config["plateCities"])
return char
plate_patch = [(random_choice(char) for char in plate) for plate in config["plateFormat"]]
return config["name"], plate_patch
def get_background(name):
bgs = os.path.join("config", name, "background")
bg_path = random.choice(glob.glob(f"{bgs}/*.*"))
bg = Image.open(bg_path).convert("RGB")
return bg
# get mask and paste the file from folder to generate the mask
def get_mask(cfg):
name, plates = parse_config(cfg)
root_dir = os.path.join("config", name)
# if none key, return default value
plateTextColor = cfg.get("plateTextColor", "#000000").lstrip("#")
img = get_template(os.path.join(root_dir, "template.png"))
mask = Image.new("L", img.size, 0)
ord_ = Image.new("L", img.size, 0)
count = 0
for idx, plate in enumerate(plates):
for char in plate:
count += 1
m_char_path = random.choice(glob.glob(f"{root_dir}/{idx}/{char}_*.png"))
mask_rect = np.where(np.array(img) == count)
h_min, w_min = np.min(mask_rect, axis=1)
h_max, w_max = np.max(mask_rect, axis=1)
m_char = Image.open(m_char_path).convert("RGBA").split()[-1]
m_char, new_width = resize_height(m_char, h_max - h_min)
x = (w_min + w_max - new_width) // 2
mask.paste(m_char, (x, h_min), m_char)
m_char = np.array(np.array(m_char) > 128, dtype=np.uint8) * 255
m_char = Image.fromarray(m_char)
ord_.paste(count*16, (x, h_min), m_char)
mask, _ = resize_height(mask, 512)
mask_wo_border = mask.copy()
mask_wo_border = mask_wo_border.filter(ImageFilter.MaxFilter(3)).filter(ImageFilter.GaussianBlur(1.5))
mask = mask.filter(ImageFilter.MaxFilter(5)).filter(ImageFilter.GaussianBlur(1.5))
ord_, _ = resize_height(ord_, 512, nearest=True)
plate_bg = get_background(name).resize(mask.size)
# make the mask as alpha channel
bg = Image.new("RGBA", mask.size, (0, 0, 0, 0))
bg.paste(plate_bg, (0, 0))
# hex to rgb
plateTextColor = tuple(int(plateTextColor[i:i + 2], 16) for i in (0, 2, 4))
bg.paste(plateTextColor, (0, 0), mask_wo_border)
return mask, ord_, bg
# get all folders in the config folder
config_path = os.path.join(os.path.dirname(__file__), "config")
plates = [f for f in os.listdir(config_path) if os.path.isdir(os.path.join(config_path, f))]
def generate(name):
print(name)
cfg = load_config(os.path.join(config_path, name, "config.yaml"))
mask, ord_, bg = get_mask(cfg)
return bg, mask, ord_
demo = gr.Interface(
generate,
[
gr.Dropdown(plates, value="cn_truck", label="Plate Style"),
],
[
gr.Image(label="Background"),
gr.Image(label="Mask"),
gr.Image(label="Ord"),
]
)
if __name__ == "__main__":
demo.launch()
|