Add application file

app.py

@@ -0,0 +1,134 @@
+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 = 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()