Update app.py

app.py

@@ -16,11 +16,17 @@ cfg = get_cfg()
 cfg.merge_from_file(config_file)
 cfg.MODEL.DEVICE="cpu"
 cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5
-cfg.MODEL.WEIGHTS = "checkpoints_model_final_imagenet_40k_synthetic.pth"
+# cfg.MODEL.WEIGHTS = "checkpoints_model_final_imagenet_40k_synthetic.pth"
 
 def predict(
-    img
+    img,
+    model="40k_synthetic"
 ):
+    if model=="40k_synthetic":
+        weights = "checkpoints_model_final_imagenet_40k_synthetic.pth"
+    else:
+        weights = "checkpoints_model_final_imagenet_40k_synthetic.pth"
+    cfg.MODEL.WEIGHTS=model
     predictor = DefaultPredictor(cfg)
     im = cv2.imread(img.name)
     output = predictor(im)
@@ -32,7 +38,7 @@ description = "Demo for Indoor Pet Detection"
 examples = [['example.jpg']]
 
 
-gr.Interface(predict, inputs=gr.inputs.Image(type="file"), outputs=gr.outputs.Image(type="pil"),enable_queue=True, title=title,
+gr.Interface(predict, [inputs=gr.inputs.Image(type="file"), gr.inputs.Dropdown(["40k_synthetic", "100k_synthetic"])], outputs=gr.outputs.Image(type="pil"),enable_queue=True, title=title,
     description=description,
     # article=article,
     examples=examples).launch()