Update app.py

app.py

@@ -26,10 +26,12 @@ import numpy as np
 
 model_path = "model_final.pth"
 
-cfg = get_cfg()
-cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7
+ccfg=get_cfg()
+cfg.merge_from_file(model_zoo.get_config_file("COCO-Detection/faster_rcnn_R_50_FPN_1x.yaml")) # load  the default configuration
+cfg.MODEL.WEIGHTS = 'model_final.pth'
+cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.8
 cfg.MODEL.ROI_HEADS.NUM_CLASSES = 2
-cfg.MODEL.WEIGHTS = model_path
+
 st.write("""
          # Gun Detection
          """
@@ -38,8 +40,9 @@ file = st.file_uploader("Please upload an image file(JPG/PNG/JPEG format)", type
 
 st.set_option('deprecation.showfileUploaderEncoding', False)
 
-gun_metadata = MetadataCatalog.get("gun_data")
-gun_metadata.thing_classes = ['guns','Gun']
+dataset_name="guns"
+classes=['guns','Gun']
+MetadataCatalog.get(dataset_name).set(thing_classes=classes)
 
 if not torch.cuda.is_available():
     cfg.MODEL.DEVICE='cpu'