feat: add plate text input field to UI

app.py

@@ -48,7 +48,6 @@ def parse_config(config):
         return char
 
     plate_patch = [(random_choice(char) for char in plate) for plate in config["plateFormat"]]
-
     return config["name"], plate_patch
 
 
@@ -62,6 +61,7 @@ def get_background(name):
 # 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("#")
@@ -70,10 +70,12 @@ def get_mask(cfg):
     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)
@@ -103,7 +105,7 @@ def get_mask(cfg):
     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
+    return plateText, bg, mask, ord_
 
 
 # get all folders in the config folder
@@ -112,10 +114,9 @@ plates = [f for f in os.listdir(config_path) if os.path.isdir(os.path.join(confi
 
 
 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_
+    text, bg, mask, ord_ = get_mask(cfg)
+    return text, bg, mask, ord_
 
 
 demo = gr.Interface(
@@ -124,6 +125,7 @@ demo = gr.Interface(
         gr.Dropdown(plates, value="cn_truck", label="Plate Style"),
     ],
     [
+        gr.Text(label="Plate Text"),
         gr.Image(label="Background"),
         gr.Image(label="Mask"),
         gr.Image(label="Ord"),
@@ -132,4 +134,4 @@ demo = gr.Interface(
 
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(debug=True)