Spaces:
Sleeping
Sleeping
File size: 4,569 Bytes
3a5d9c5 df08103 3a5d9c5 3bee6cc 3a5d9c5 3bee6cc 3a5d9c5 3bee6cc 3a5d9c5 df08103 3a5d9c5 df08103 3a5d9c5 8b70a7d 3a5d9c5 3bee6cc 3a5d9c5 b2972d7 3a5d9c5 b2972d7 3a5d9c5 3bee6cc 3a5d9c5 df08103 3a5d9c5 df08103 3bee6cc 3a5d9c5 df08103 3a5d9c5 b2972d7 3a5d9c5 df08103 3bee6cc 3a5d9c5 b2972d7 df08103 3a5d9c5 b2972d7 |
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 137 138 139 140 141 142 143 144 145 146 147 |
import os
import glob
import time
import json
import random
import numpy as np
import gradio as gr
from matplotlib import pyplot as plt
from PIL import Image, ImageDraw, ImageFilter
def get_template(path):
# RGB to HSV
img = Image.open(path).convert('HSV')
copy = img.convert('L').copy()
# where is white
copy = np.where(np.array(copy) == 255, 255, 0).astype(np.uint8)
# convert numpy array to image
copy = Image.fromarray(copy, "L")
# pick the H channel
img = img.split()[0]
# divide by 16
img = np.array(img) + 8
img = Image.fromarray(img // 16)
return img, copy
# 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 json file
def load_config(path):
with open(path, "r") as f:
config = json.load(f)
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, dot = 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
plateText = ""
for idx, plate in enumerate(plates):
for char in plate:
count += 1
plateText += char
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)
# add the dot img on the mask
mask = Image.composite(dot, mask, dot)
mask, _ = resize_height(mask, 512)
mask_wo_border = mask.copy()
mask_wo_border = mask_wo_border.filter(ImageFilter.MaxFilter(3))
mask = mask.filter(ImageFilter.MaxFilter(5)).filter(ImageFilter.GaussianBlur(1.5))
ord_, _ = resize_height(ord_, 512, nearest=True)
plate_bg = get_background(name)
mask_wo_border = mask_wo_border.resize(plate_bg.size)
# make the mask as alpha channel
bg = Image.new("RGBA", plate_bg.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 plateText, bg, mask, ord_
# 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):
cfg = load_config(os.path.join(config_path, name, "config.json"))
return get_mask(cfg) # text, bg, mask, ord_
demo = gr.Interface(
generate,
[
gr.Dropdown(plates, value="cn_truck", label="Plate Style"),
],
[
gr.Text(label="Plate Text"),
gr.Image(label="Background", show_label=False, format="png"),
gr.Image(label="Mask", show_label=False, format="png"),
gr.Image(label="Ord", show_label=False, format="png"),
],
allow_flagging=False,
)
if __name__ == "__main__":
demo.launch(debug=True)
|