fiber

app.py

@@ -42,8 +42,6 @@ from maskrcnn_benchmark.engine.predictor_glip import GLIPDemo
 config_file = "configs/pretrain_new/desco_glip.yaml"
 weight_file = "MODEL/desco_glip_tiny.pth"
 
-config_file = "configs/pretrain_new/desco_fiber.yaml"
-weight_file = "MODEL/desco_fiber_base.pth"
 
 # Use this command if you want to try the GLIP-L model
 # ! wget https://penzhanwu2bbs.blob.core.windows.net/data/GLIPv1_Open/models/glip_large_model.pth -O MODEL/glip_large_model.pth
@@ -65,6 +63,16 @@ glip_demo = GLIPDemo(
     show_mask_heatmaps=False
 )
 
+config_file = "configs/pretrain_new/desco_fiber.yaml"
+weight_file = "MODEL/desco_fiber_base.pth"
+cfg.merge_from_list(["MODEL.WEIGHT", weight_file])
+cfg.merge_from_list(["MODEL.DEVICE", "cuda"])
+fiber_demo = GLIPDemo(
+    cfg,
+    min_image_size=800,
+    confidence_threshold=0.7,
+    show_mask_heatmaps=False
+)
 
 def predict(image, text, specified_tokens=""):
     if specified_tokens.strip() == "":
@@ -72,7 +80,8 @@ def predict(image, text, specified_tokens=""):
     else:
         specified_tokens = specified_tokens.strip().split(";")
     result, _ = glip_demo.run_on_web_image(image[:, :, [2, 1, 0]], text, 0.5, specified_tokens)
-    return result[:, :, [2, 1, 0]]
+    fiber_result, _ = fiber_demo.run_on_web_image(image[:, :, [2, 1, 0]], text, 0.5, specified_tokens)
+    return result[:, :, [2, 1, 0]], fiber_result[:, :, [2, 1, 0]]
 
 
 image = gr.inputs.Image()
@@ -86,6 +95,10 @@ gr.Interface(
             type="pil",
             # label="grounding results"
         ),
+        gr.outputs.Image(
+            type="pil",
+            # label="grounding results"
+        ),
     ],
     examples=[
         #["./flickr_9472793441.jpg", "bobble heads on top of the shelf ."],