Create app.py

app.py

@@ -0,0 +1,56 @@
+import argparse
+
+import cv2
+from ditod import add_vit_config
+from detectron2.config import get_cfg
+from detectron2.utils.visualizer import ColorMode, Visualizer
+from detectron2.data import MetadataCatalog
+from detectron2.engine import DefaultPredictor
+
+
+# Step 1: instantiate config
+cfg = get_cfg()
+add_vit_config(cfg)
+cfg.merge_from_file("cascade_dit_base.yaml")
+
+# Step 2: add model weights URL to config
+cfg.MODEL.WEIGHTS = https://layoutlm.blob.core.windows.net/dit/dit-fts/publaynet_dit-b_mrcnn.pth
+
+# Step 3: set device
+# TODO also support GPU
+cfg.MODEL.DEVICE='cpu'
+
+# Step 4: define model
+predictor = DefaultPredictor(cfg)
+
+
+def analyze_image(img):
+    md = MetadataCatalog.get(cfg.DATASETS.TEST[0])
+    if cfg.DATASETS.TEST[0]=='icdar2019_test':
+        md.set(thing_classes=["table"])
+    else:
+        md.set(thing_classes=["text","title","list","table","figure"])
+    
+    output = predictor(img)["instances"]
+    v = Visualizer(img[:, :, ::-1],
+                    md,
+                    scale=1.0,
+                    instance_mode=ColorMode.SEGMENTATION)
+    result = v.draw_instance_predictions(output.to("cpu"))
+    result_image = result.get_image()[:, :, ::-1]
+    
+    return result_image
+    
+title = "Interactive demo: Document Layout Analysis with DiT"
+description = "This is a demo for Microsoft's Document Image Transformer (DiT)."
+examples =[['document.png']]
+
+iface = gr.Interface(fn=analyze_image, 
+                     inputs=gr.inputs.Image(type="numpy"), 
+                     outputs=gr.outputs.Image(type="numpy", label="analyzed image"),
+                     title=title,
+                     description=description,
+                     article=article,
+                     examples=examples,
+                     enable_queue=True)
+iface.launch(debug=True)