update app.py

app.py

@@ -8,14 +8,18 @@
 @Contact :   2909171338@qq.com
 """
 
+import argparse
 import gradio as gr
-from PIL import Image
+from PIL import Image, ImageDraw
 import io
 import matplotlib.pyplot as plt
 import numpy as np
 from datasets import load_dataset
 
 ds = None
+DATASET_NAME = None
+FAST = True
+SPLIT = "all"
 
 
 def get_dataset():
@@ -30,7 +34,7 @@ def get_dataset():
     """
     global ds
     if ds is None:
-        ds = load_dataset("lh9171338/Wireframe")
+        ds = load_dataset(DATASET_NAME)
     return ds
 
 
@@ -63,25 +67,35 @@ def draw_lines(image, lines):
     Returns:
         image (PIL.Image): drawn image
     """
-    height, width = image.shape[:2]
-    fig = plt.figure()
-    fig.set_size_inches(width / height, 1, forward=False)
-    ax = plt.Axes(fig, [0.0, 0.0, 1.0, 1.0])
-    ax.set_axis_off()
-    fig.add_axes(ax)
-    plt.xlim([-0.5, width - 0.5])
-    plt.ylim([height - 0.5, -0.5])
-    plt.imshow(image)
-    for pts in lines:
-        pts = pts - 0.5
-        plt.plot(pts[:, 0], pts[:, 1], color="orange", linewidth=0.5)
-        plt.scatter(pts[[0, -1], 0], pts[[0, -1], 1], color="#33FFFF", s=1.2, edgecolors="none", zorder=5)
-
-    buf = io.BytesIO()
-    fig.savefig(buf, format="png", dpi=height, bbox_inches=0)
-    buf.seek(0)
-    plt.close(fig)
-    image = Image.open(buf)
+    if FAST:
+        image = Image.fromarray(image)
+        draw = ImageDraw.Draw(image)
+        for pts in lines:
+            pts = pts - 0.5
+            pts_list = [tuple(p) for p in pts]
+            draw.line(pts_list, fill="orange", width=2)
+            draw.circle(pts_list[0], 3, fill="#33FFFF")
+            draw.circle(pts_list[-1], 3, fill="#33FFFF")
+    else:
+        height, width = image.shape[:2]
+        fig = plt.figure()
+        fig.set_size_inches(width / height, 1, forward=False)
+        ax = plt.Axes(fig, [0.0, 0.0, 1.0, 1.0])
+        ax.set_axis_off()
+        fig.add_axes(ax)
+        plt.xlim([-0.5, width - 0.5])
+        plt.ylim([height - 0.5, -0.5])
+        plt.imshow(image)
+        for pts in lines:
+            pts = pts - 0.5
+            plt.plot(pts[:, 0], pts[:, 1], color="orange", linewidth=0.5)
+            plt.scatter(pts[[0, -1], 0], pts[[0, -1], 1], color="#33FFFF", s=1.2, edgecolors="none", zorder=5)
+
+        buf = io.BytesIO()
+        fig.savefig(buf, format="png", dpi=height, bbox_inches=0)
+        buf.seek(0)
+        plt.close(fig)
+        image = Image.open(buf)
     return image
 
 
@@ -104,9 +118,22 @@ def show_image(split, index):
     return image
 
 
-if __name__ == "__main__":
+def main():
+    """
+    main
+
+    Args:
+        None
+
+    Returns:
+        None
+    """
     with gr.Blocks() as demo:
-        split_selector = gr.Dropdown(["train", "test"], label="Split", value="train")
+        if SPLIT == "all":
+            choices = ["train", "test"]
+        else:
+            choices = [SPLIT]
+        split_selector = gr.Dropdown(choices, label="Split", value=choices[0])
         index_slider = gr.Slider(0, 1, step=1, label="Index", value=0)
         output = gr.Image()
 
@@ -114,3 +141,16 @@ if __name__ == "__main__":
         index_slider.change(show_image, [split_selector, index_slider], output)
         demo.load(selector_change_callback, split_selector, [index_slider, output])
         demo.launch(share=False)
+
+
+if __name__ == "__main__":
+    argparser = argparse.ArgumentParser()
+    argparser.add_argument("-n", "--dataset_name", type=str, help="dataset name", default="lh9171338/Wireframe")
+    argparser.add_argument("-f", "--fast", type=bool, help="whether to use fast drawing method", default=True)
+    argparser.add_argument("-s", "--split", type=str, help="split", default="all", choices=["all", "train", "test"])
+    args = argparser.parse_args()
+    print(args)
+    DATASET_NAME = args.dataset_name
+    FAST = args.fast
+    SPLIT = args.split
+    main()