chore: Refactor get_mask function to improve readability and performance

app.py

@@ -6,18 +6,24 @@ 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
+    return img, copy
 
 
 # resize via height, keep the aspect ratio
@@ -65,7 +71,7 @@ def get_mask(cfg):
     # if none key, return default value
     plateTextColor = cfg.get("plateTextColor", "#000000").lstrip("#")
 
-    img = get_template(os.path.join(root_dir, "template.png"))
+    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
@@ -89,6 +95,8 @@ def get_mask(cfg):
             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))
@@ -98,7 +106,7 @@ def get_mask(cfg):
 
     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))
@@ -116,8 +124,7 @@ plates = [f for f in os.listdir(config_path) if os.path.isdir(os.path.join(confi
 
 def generate(name):
     cfg = load_config(os.path.join(config_path, name, "config.json"))
-    text, bg, mask, ord_ = get_mask(cfg)
-    return text, bg, mask, ord_
+    return get_mask(cfg) # text, bg, mask, ord_
 
 
 demo = gr.Interface(

config/cn_car/config.json

@@ -41,6 +41,6 @@
     "plateCities": "京津冀晋蒙辽吉黑沪苏浙皖闽赣鲁豫鄂湘粤桂琼渝川贵云藏陕甘青宁新",
     "plateTextColor": "#FFFFFF",
     "plateFormat": [
-        "@A·#####"
+        "@A#####"
     ]
 }

config/cn_car/config.yaml

@@ -9,4 +9,4 @@ plateTextColor: "#FFFFFF"
 # A representing any letter
 # 0 representing any number
 # @ representing any city
-plateFormat: ["@A·#####"]
+plateFormat: ["@A#####"]

config/cn_truck/config.json

@@ -7,7 +7,7 @@
     "plateCities": "京津冀晋蒙辽吉黑沪苏浙皖闽赣鲁豫鄂湘粤桂琼渝川贵云藏陕甘青宁新",
     "plateTextColor": "#111111",
     "plateFormat": [
-        "@·A",
+        "@A",
         "#0000"
     ]
 }

config/cn_truck/config.yaml

@@ -11,7 +11,7 @@ plateTextColor: "#111111"
 # 0 representing any number (if multi-line, 0 means any number in the first line)
 # 1 representing any number (if multi-line, 1 means any number in the second line)
 # @ representing any city
-plateFormat: ["@·A", "#0000"]
+plateFormat: ["@A", "#0000"]
 
 # single-line plate
 # +----------+