update

app.py

@@ -238,30 +238,62 @@ class AsyncPredictor:
 
 
 detectron2_model_list = {
+    # Cityscapes
+    "Cityscapes/mask_rcnn_R_50_FPN":{
+        "config_file": "configs/Cityscapes/mask_rcnn_R_50_FPN.yaml",
+        "ckpts": "detectron2://Cityscapes/"
+    },
+    # COCO-Detection
     "COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x":{
         "config_file": "configs/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml",
         "ckpts": "detectron2://COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x/137849600/model_final_f10217.pkl"
     },
+    "COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_1x":{
+        "config_file": "configs/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_1x.yaml",
+        "ckpts": "detectron2://COCO-InstanceSegmentation/"
+    },
+    "COCO-InstanceSegmentation/mask_rcnn_R_50_DC5_1x":{
+        "config_file": "configs/COCO-InstanceSegmentation/mask_rcnn_R_50_DC5_1x.yaml",
+        "ckpts": "detectron2://COCO-InstanceSegmentation/"
+    },
+    "COCO-InstanceSegmentation/mask_rcnn_R_50_DC5_3x": {
+        "config_file": "configs/COCO-InstanceSegmentation/mask_rcnn_R_50_DC5_3x.yaml",
+        "ckpts": "detectron2://COCO-InstanceSegmentation/"
+    },
+    "COCO-InstanceSegmentation/mask_rcnn_R_50_C4_1x": {
+        "config_file": "configs/COCO-InstanceSegmentation/mask_rcnn_R_50_C4_1x.yaml",
+        "ckpts": "detectron2://COCO-InstanceSegmentation/"
+    },
+    "COCO-InstanceSegmentation/mask_rcnn_R_50_C4_3x": {
+        "config_file": "configs/COCO-InstanceSegmentation/mask_rcnn_R_50_C4_3x.yaml",
+        "ckpts": "detectron2://COCO-InstanceSegmentation/"
+    },
+    "COCO-InstanceSegmentation/mask_rcnn_R_101_C4_3x": {
+        "config_file": "configs/COCO-InstanceSegmentation/mask_rcnn_R_101_C4_3x.yaml",
+        "ckpts": "detectron2://COCO-InstanceSegmentation/"
+    },
+    "COCO-InstanceSegmentation/mask_rcnn_R_101_DC5_3x": {
+        "config_file": "configs/COCO-InstanceSegmentation/mask_rcnn_R_101_DC5_3x.yaml",
+        "ckpts": "detectron2://COCO-InstanceSegmentation/"
+    },
+    "COCO-InstanceSegmentation/mask_rcnn_R_101_FPN_3x": {
+        "config_file": "configs/COCO-InstanceSegmentation/mask_rcnn_R_101_FPN_3x.yaml",
+        "ckpts": "detectron2://COCO-InstanceSegmentation/"
+    },
+    "COCO-InstanceSegmentation/mask_rcnn_X_101_32x8d_FPN_3x": {
+        "config_file": "configs/COCO-InstanceSegmentation/mask_rcnn_X_101_32x8d_FPN_3x.yaml",
+        "ckpts": "detectron2://COCO-InstanceSegmentation/"
+    },
+    # COCO-Detection
+    "COCO-Detection/mask_rcnn_X_101_32x8d_FPN_3x": {
+        "config_file": "configs/COCO-Detection/mask_rcnn_X_101_32x8d_FPN_3x.yaml",
+        "ckpts": "detectron2://COCO-Detection/"
+    },
 }
 
 
-# def dtectron2_instance_inference(image, config_file, ckpts, device):
-#     cfg = get_cfg()
-#     cfg.merge_from_file(config_file)
-#     cfg.MODEL.WEIGHTS = ckpts
-#     cfg.MODEL.DEVICE = "cpu"
-#     cfg.output = "output_img.jpg"
-#     visualization_demo = VisualizationDemo(cfg, device=device)
-#     if image:
-#         intput_path = "intput_img.jpg"
-#         image.save(intput_path)
-#         image = read_image(intput_path, format="BGR")
-#         predictions, vis_output = visualization_demo.run_on_image(image)
-#         output_image = PIL.Image.fromarray(vis_output.get_image())
-#         # print("predictions: ", predictions)
-#         return output_image
-
-def dtectron2_instance_inference(image, input_model_name, device):
+
+def dtectron2_instance_inference(image, input_model_name, confidence_threshold, device):
     cfg = get_cfg()
     config_file = detectron2_model_list[input_model_name]["config_file"]
     ckpts = detectron2_model_list[input_model_name]["ckpts"]
@@ -269,6 +301,9 @@ def dtectron2_instance_inference(image, input_model_name, device):
     cfg.MODEL.WEIGHTS = ckpts
     cfg.MODEL.DEVICE = "cpu"
     cfg.output = "output_img.jpg"
+    cfg.MODEL.RETINANET.SCORE_THRESH_TEST = confidence_threshold
+    cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = confidence_threshold
+    cfg.MODEL.PANOPTIC_FPN.COMBINE.INSTANCES_CONFIDENCE_THRESH = confidence_threshold
     visualization_demo = VisualizationDemo(cfg, device=device)
     if image:
         intput_path = "intput_img.jpg"
@@ -280,15 +315,20 @@ def dtectron2_instance_inference(image, input_model_name, device):
         return output_image
 
 def download_test_img():
+    import shutil
+    torch.hub.download_url_to_file(
+        'https://github.com/isLinXu/issues/files/12643351/configs.zip',
+        'configs.zip')
     # Images
     torch.hub.download_url_to_file(
         'https://user-images.githubusercontent.com/59380685/268517006-d8d4d3b3-964a-4f4d-8458-18c7eb75a4f2.jpg',
         '000000502136.jpg')
-
+    shutil.unpack_archive('configs.zip', 'configs', 'zip')
 
 if __name__ == '__main__':
     input_image = gr.inputs.Image(type='pil', label='Input Image')
     input_model_name = gr.inputs.Dropdown(list(detectron2_model_list.keys()), label="Model Name", default="COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x")
+    input_prediction_threshold = gr.inputs.Slider(minimum=0.0, maximum=1.0, step=0.01, default=0.25, label="Confidence Threshold")
     input_device = gr.inputs.Dropdown(["cpu", "cuda"], label="Devices", default="cpu")
     output_image = gr.outputs.Image(type='pil', label='Output Image')
     output_predictions = gr.outputs.Textbox(type='text', label='Output Predictions')
@@ -301,8 +341,8 @@ if __name__ == '__main__':
               "<p style='text-align: center'><a href='https://github.com/facebookresearch/detectron2'>gradio build by gatilin</a></a></p>"
     download_test_img()
 
-    examples = [["000000502136.jpg", "COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x", "cpu"]]
+    examples = [["000000502136.jpg", "COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x", 0.25, "cpu"]]
     gr.Interface(fn=dtectron2_instance_inference,
-                 inputs=[input_image, input_model_name, input_device],
+                 inputs=[input_image, input_model_name, input_prediction_threshold, input_device],
                  outputs=output_image,examples=examples,
                  title=title, description=description, article=article).launch()