support local dataset

README.md

@@ -12,3 +12,20 @@ short_description: Wireframe Dataset (for Line Segment Detection) Space
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+
+# Install
+```shell
+git clone https://huggingface.co/spaces/lh9171338/Wireframe
+cd Wireframe
+python3 -m pip install -r requirements.txt
+```
+
+# Run
+
+```shell
+# using Huggingface online dataset
+python3 app.py -d lh9171338/Wireframe -s all
+
+# using local dataset
+python3 app.py -d <DATASET_PATH> -l True -s all
+```

app.py

@@ -18,7 +18,8 @@ from datasets import load_dataset
 
 ds = None
 DATASET_NAME = None
-FAST = True
+LOCAL = False
+FAST = False
 SPLIT = "all"
 
 
@@ -34,7 +35,10 @@ def get_dataset():
     """
     global ds
     if ds is None:
-        ds = load_dataset(DATASET_NAME)
+        if LOCAL:
+            ds = load_dataset("imagefolder", data_dir=DATASET_NAME)
+        else:
+            ds = load_dataset(DATASET_NAME)
     return ds
 
 
@@ -146,11 +150,13 @@ def main():
 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("-l", "--local", type=bool, help="use local data or not", default=False)
+    argparser.add_argument("-f", "--fast", type=bool, help="whether to use fast drawing method", default=False)
     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
+    LOCAL = args.local
     FAST = args.fast
     SPLIT = args.split
     main()