Update app.py

app.py

@@ -1,8 +1,14 @@
 import os
 os.system('pip install detectron2 -f https://dl.fbaipublicfiles.com/detectron2/wheels/cu102/torch1.9/index.html')
+os.system("git clone https://github.com/microsoft/unilm.git")
+
+import sys
+sys.path.append("unilm")
 
 import cv2
 
+from unilm.dit.object_detection.ditod import add_vit_config
+
 from detectron2.config import CfgNode as CN
 from detectron2.config import get_cfg
 from detectron2.utils.visualizer import ColorMode, Visualizer
@@ -12,37 +18,6 @@ from detectron2.engine import DefaultPredictor
 import gradio as gr
 
 
-def add_vit_config(cfg):
-    """
-    Add config for VIT.
-    """
-    _C = cfg
-
-    _C.MODEL.VIT = CN()
-
-    # CoaT model name.
-    _C.MODEL.VIT.NAME = ""
-
-    # Output features from CoaT backbone.
-    _C.MODEL.VIT.OUT_FEATURES = ["layer3", "layer5", "layer7", "layer11"]
-
-    _C.MODEL.VIT.IMG_SIZE = [224, 224]
-
-    _C.MODEL.VIT.POS_TYPE = "shared_rel"
-
-    _C.MODEL.VIT.DROP_PATH = 0.
-
-    _C.MODEL.VIT.MODEL_KWARGS = "{}"
-
-    _C.SOLVER.OPTIMIZER = "ADAMW"
-
-    _C.SOLVER.BACKBONE_MULTIPLIER = 1.0
-
-    _C.AUG = CN()
-
-    _C.AUG.DETR = False
-
-
 # Step 1: instantiate config
 cfg = get_cfg()
 add_vit_config(cfg)