Move config file to model repository

python_utils/get_model.py

@@ -1,5 +1,7 @@
 import urllib.request
 
+OPTIMAL_NMS_THRESHOLD = 0.7
+_model_config_url = "https://huggingface.co/TZProject/final_tz_segmentor/resolve/main/final_model_config.yaml"
 def get_set_up():
     import torch
     TORCH_VERSION = ".".join(torch.__version__.split(".")[:2])
@@ -18,9 +20,9 @@ def load_model():
 
     ## define relevant parameters
     cfg = get_cfg()
-    config_url = "https://huggingface.co/TZProject/final_tz_segmentor/resolve/main/final_model_config.yaml"
-    config_filename = "final_model_config.yaml"
-    urllib.request.urlretrieve(config_url, filename=config_filename)
+
+    config_filename = "./final_model_config.yaml"
+    urllib.request.urlretrieve(_model_config_url, filename=config_filename)
     cfg.merge_from_file(config_filename)
     if not torch.cuda.is_available():
         cfg.MODEL.DEVICE = "cpu"
@@ -38,7 +40,7 @@ def load_model():
 
     return predictor
 
-def mask_nms(masks, scores, nms_threshold=0.5):
+def mask_nms(masks, scores, nms_threshold=OPTIMAL_NMS_THRESHOLD):
     """
     Runs class agnostic NMS on masks/segmentations instead of the bounding boxes.
     :param      masks: (list float) List of coordinates that make up the mask output from the model.
@@ -72,7 +74,7 @@ def mask_nms(masks, scores, nms_threshold=0.5):
                 order.remove(j)
     return masks_kept
 
-def apply_nms(prediction, mask=False, cls_agnostic_nms=0.7):
+def apply_nms(prediction, mask=False, cls_agnostic_nms=OPTIMAL_NMS_THRESHOLD):
     from torchvision.ops import nms
     from detectron2.structures import Instances