diff --git a/adet/data/augmentation.py b/adet/data/augmentation.py
index 626385f..cd2e1df 100644
--- a/adet/data/augmentation.py
+++ b/adet/data/augmentation.py
@@ -1,14 +1,12 @@
 import random
-from typing import Tuple
-import sys
-from PIL import Image
 import numpy as np
+
 from fvcore.transforms import transform as T
+from fvcore.transforms.transform import Transform, NoOpTransform
+
 from detectron2.data.transforms import RandomCrop, StandardAugInput
 from detectron2.structures import BoxMode
-import torch
 from detectron2.data.transforms import Augmentation, PadTransform
-from fvcore.transforms.transform import Transform, NoOpTransform
 
 
 def gen_crop_transform_with_instance(crop_size, image_size, instances, crop_box=True):
@@ -130,4 +128,4 @@ def get_transform(self, img):
         # pad_h, pad_w = 32 - ori_h % 32, 32 - ori_w % 32
         return PadTransform(
             0, 0, pad_w, pad_h, pad_value=0
-        )
\ No newline at end of file
+        )
diff --git a/adet/data/builtin.py b/adet/data/builtin.py
index 024e72e..30a2aee 100644
--- a/adet/data/builtin.py
+++ b/adet/data/builtin.py
@@ -75,24 +75,24 @@ def register_all_coco(root="datasets", voc_size_cfg=37, num_pts_cfg=25):
             num_pts_cfg
         )
 
-
-# get the vocabulary size and number of point queries in each instance
-# to eliminate blank text and sample gt according to Bezier control points
-parser = default_argument_parser()
-# add the following argument to avoid some errors while running demo/demo.py
-parser.add_argument("--input", nargs="+", help="A list of space separated input images")
-parser.add_argument(
-    "--output",
-    help="A file or directory to save output visualizations. "
-    "If not given, will show output in an OpenCV window.",
-)
-parser.add_argument(
-    "--opts",
-    help="Modify config options using the command-line 'KEY VALUE' pairs",
-    default=[],
-    nargs=argparse.REMAINDER,
+if __name__ == "__main__":
+    # get the vocabulary size and number of point queries in each instance
+    # to eliminate blank text and sample gt according to Bezier control points
+    parser = default_argument_parser()
+    # add the following argument to avoid some errors while running demo/demo.py
+    parser.add_argument("--input", nargs="+", help="A list of space separated input images")
+    parser.add_argument(
+        "--output",
+        help="A file or directory to save output visualizations. "
+        "If not given, will show output in an OpenCV window.",
     )
-args = parser.parse_args()
-cfg = get_cfg()
-cfg.merge_from_file(args.config_file)
-register_all_coco(voc_size_cfg=cfg.MODEL.TRANSFORMER.VOC_SIZE, num_pts_cfg=cfg.MODEL.TRANSFORMER.NUM_POINTS)
+    parser.add_argument(
+        "--opts",
+        help="Modify config options using the command-line 'KEY VALUE' pairs",
+        default=[],
+        nargs=argparse.REMAINDER,
+        )
+    args = parser.parse_args()
+    cfg = get_cfg()
+    cfg.merge_from_file(args.config_file)
+    register_all_coco(voc_size_cfg=cfg.MODEL.TRANSFORMER.VOC_SIZE, num_pts_cfg=cfg.MODEL.TRANSFORMER.NUM_POINTS)