Add some url variables and model params that might be changeable in future iterations

python_utils/get_model.py

@@ -1,9 +1,14 @@
 import urllib.request
 import tempfile
 
-
+## Urls and model variables that might change.
 OPTIMAL_NMS_THRESHOLD = 0.7
-_model_config_url = "https://huggingface.co/TZProject/final_tz_segmentor/resolve/main/final_model_config.yaml"
+model_page = "https://huggingface.co/TZProject/final_tz_segmentor"
+_model_config_url = model_page+"/resolve/main/final_model_config.yaml"
+MODEL_VERSION = "v1.0"
+discussion_url = 'https://huggingface.co/spaces/TZProject/TZSeedApp/discussions'
+github_repo_url = 'https://github.com/JATamura/TZSegmenting'
+
 def get_set_up():
     import torch
     TORCH_VERSION = ".".join(torch.__version__.split(".")[:2])
@@ -14,7 +19,20 @@ def get_set_up():
 
     # print("detectron2:", detectron2.__version__)
 
-def load_model():
+def load_model(using_final_model:bool=True):
+    """
+    Load and configure a Detectron2 model predictor. The method creates a configuration
+    object, merges it with a specified configuration file fetched from a remote URL,
+    and initializes the model using the `DefaultPredictor` from Detectron2. The model
+    will be set up to run on either a GPU or CPU depending on the system's capabilities.
+
+    :param using_final_model: A flag to indicate whether the final model should be used with
+        specific configurations. When this is set to True, adjustments are made to suppress
+        settings that are relevant during training but cause runtime errors during inference.
+        This includes disabling certain loss calculations that depend on training data.
+
+    :return: A Detectron2 predictor object configured and ready for inference.
+    """
     # return None
     import torch
     from detectron2.engine import DefaultPredictor
@@ -32,12 +50,13 @@ def load_model():
     else:
         cfg.MODEL.DEVICE = 'cuda'
 
-    ## when rerouting to use the final model (final_tz_segmentor) USE_FED_LOSS has to be set to False
-    ## this setting requires the training data to calculate class imbalance that the app will not have access to and cause a runtime error
-    ## some messages will appear when using the model that certain weights are not being used 
-    ## but these are used during training and not inference and shouldn't affect the model performance
-    ## code below
-    cfg.MODEL.ROI_BOX_HEAD.USE_FED_LOSS = False
+    if using_final_model:
+        ## when rerouting to use the final model (final_tz_segmentor) USE_FED_LOSS has to be set to False
+        ## this setting requires the training data to calculate class imbalance that the app will not have access to and cause a runtime error
+        ## some messages will appear when using the model that certain weights are not being used
+        ## but these are used during training and not inference and shouldn't affect the model performance
+        ## code below
+        cfg.MODEL.ROI_BOX_HEAD.USE_FED_LOSS = False
     
     predictor = DefaultPredictor(cfg)
 
@@ -78,6 +97,25 @@ def mask_nms(masks, scores, nms_threshold=OPTIMAL_NMS_THRESHOLD):
     return masks_kept
 
 def apply_nms(prediction, mask=False, cls_agnostic_nms=OPTIMAL_NMS_THRESHOLD):
+    """
+    Applies Non-Maximum Suppression (NMS) to filter redundant bounding boxes
+    from the prediction output produced by an object detection model. The method is compatible
+    with instances containing bounding boxes, scores, classes, and optionally mask predictions.
+    The NMS operation is performed based on a specific IoU threshold, configurable via the
+    `cls_agnostic_nms` parameter. If mask mode is enabled, a mask-based NMS step will be applied
+    to filter out redundant masks among the predictions.
+
+    :param prediction: The dictionary containing detection results, where the
+        "instances" key points to an `Instances` object. This object should
+        contain attributes including bounding boxes, scores, predicted classes,
+        and optionally predicted masks.
+    :type prediction: Dict[str, Instances]
+    :param mask: A boolean flag indicating whether NMS should additionally be
+        applied to instance masks. Defaults to ``False``.
+    :type mask: bool
+    :param cls_agnostic_nms: An IoU threshold for NMS, reflecting the
+        level of overlap above which boxes are considered for suppression.
+        Defaults to ``"""
     from torchvision.ops import nms
     from detectron2.structures import Instances