fix: RND-117: YOLO score threshold name (#624)

HumanSignal · Sep 8, 2024 · f45a92a · f45a92a
1 parent 4828d56
commit f45a92a
Show file tree

Hide file tree

Showing 14 changed files with 27 additions and 27 deletions.
diff --git a/label_studio_ml/examples/yolo/control_models/base.py b/label_studio_ml/examples/yolo/control_models/base.py
@@ -12,7 +12,7 @@
 
 # use matplotlib plots for debug
 DEBUG_PLOT = os.getenv("DEBUG_PLOT", "false").lower() in ["1", "true"]
-SCORE_THRESHOLD = float(os.getenv("SCORE_THRESHOLD", 0.5))
+MODEL_SCORE_THRESHOLD = float(os.getenv("MODEL_SCORE_THRESHOLD", 0.5))
 DEFAULT_MODEL_ROOT = os.path.join(os.path.dirname(os.path.dirname(__file__)), "models")
 MODEL_ROOT = os.getenv("MODEL_ROOT", DEFAULT_MODEL_ROOT)
 os.makedirs(MODEL_ROOT, exist_ok=True)
@@ -40,7 +40,7 @@ class ControlModel(BaseModel):
         value (str): The value name from the object that this control operates on, e.g., an image or text field.
         model (object): The model instance (e.g., YOLO) used to generate predictions for this control.
         model_path (str): Path to the YOLO model file.
-        score_threshold (float): Threshold for prediction scores; predictions below this value will be ignored.
+        model_score_threshold (float): Threshold for prediction scores; predictions below this value will be ignored.
         label_map (Optional[Dict[str, str]]): A mapping of model labels to Label Studio labels.
     """
 
@@ -51,7 +51,7 @@ class ControlModel(BaseModel):
     value: str
     model: YOLO
     model_path: ClassVar[str]
-    score_threshold: float = 0.5
+    model_score_threshold: float = 0.5
     label_map: Optional[Dict[str, str]] = {}
     label_studio_ml_backend: LabelStudioMLBase
 
@@ -90,8 +90,8 @@ def create(cls, mlbackend: LabelStudioMLBase, control: ControlTag):
                 f"Skipping control tag '{control.tag}' with name '{from_name}', model_skip=true found"
             )
             return None
-        # read `score_threshold` attribute from the control tag, e.g.: <RectangleLabels score_threshold="0.5">
-        score_threshold = float(control.attr.get("score_threshold") or SCORE_THRESHOLD)
+        # read `model_score_threshold` attribute from the control tag, e.g.: <RectangleLabels model_score_threshold="0.5">
+        model_score_threshold = float(control.attr.get("model_score_threshold") or MODEL_SCORE_THRESHOLD)
         # read `model_path` attribute from the control tag
         model_path = (
             ALLOW_CUSTOM_MODEL_PATH and control.attr.get("model_path")
@@ -111,7 +111,7 @@ def create(cls, mlbackend: LabelStudioMLBase, control: ControlTag):
             to_name=to_name,
             value=value,
             model=model,
-            score_threshold=score_threshold,
+            model_score_threshold=model_score_threshold,
             label_map=label_map,
             label_studio_ml_backend=mlbackend,
         )
@@ -154,7 +154,7 @@ def __str__(self):
         """Return a string with full representation of the control tag."""
         return (
             f"{self.type} from_name={self.from_name}, "
-            f"label_map={self.label_map}, score_threshold={self.score_threshold}"
+            f"label_map={self.label_map}, model_score_threshold={self.model_score_threshold}"
         )
 
     class Config:

diff --git a/label_studio_ml/examples/yolo/control_models/choices.py b/label_studio_ml/examples/yolo/control_models/choices.py
@@ -48,8 +48,8 @@ def create_choices(self, results, path):
             names = [model_names[index]]
         # multi
         else:
-            # get indexes of data where data >= self.score_threshold
-            indexes = np.where(data >= self.score_threshold)
+            # get indexes of data where data >= self.model_score_threshold
+            indexes = np.where(data >= self.model_score_threshold)
             probs = data[indexes].tolist()
             names = [self.model.names[int(i)] for i in indexes[0]]
 
@@ -67,7 +67,7 @@ def create_choices(self, results, path):
             f"names > {names}\n"
         )
 
-        if score < self.score_threshold:
+        if score < self.model_score_threshold:
             logger.debug(f"Score is too low for single choice: {names[0]} = {probs[0]}")
             return []
 

diff --git a/label_studio_ml/examples/yolo/control_models/keypoint_labels.py b/label_studio_ml/examples/yolo/control_models/keypoint_labels.py
@@ -93,7 +93,7 @@ def create_keypoints(self, results, path):
             )
 
             # bbox score is too low
-            if bbox_conf < self.score_threshold:
+            if bbox_conf < self.model_score_threshold:
                 continue
 
             # There is no mapping between model label and LS label

diff --git a/label_studio_ml/examples/yolo/control_models/polygon_labels.py b/label_studio_ml/examples/yolo/control_models/polygon_labels.py
@@ -48,7 +48,7 @@ def create_polygons(self, results, path):
             )
 
             # bbox score is too low
-            if score < self.score_threshold:
+            if score < self.model_score_threshold:
                 continue
 
             # there is no mapping between model label and LS label

diff --git a/label_studio_ml/examples/yolo/control_models/rectangle_labels.py b/label_studio_ml/examples/yolo/control_models/rectangle_labels.py
@@ -68,7 +68,7 @@ def create_rectangles(self, results, path):
             )
 
             # bbox score is too low
-            if score < self.score_threshold:
+            if score < self.model_score_threshold:
                 continue
 
             # there is no mapping between model label and LS label

diff --git a/label_studio_ml/examples/yolo/control_models/rectangle_labels_obb.py b/label_studio_ml/examples/yolo/control_models/rectangle_labels_obb.py
@@ -67,7 +67,7 @@ def create_rotated_rectangles(self, results, path):
             )
 
             # bbox score is too low
-            if score < self.score_threshold:
+            if score < self.model_score_threshold:
                 continue
 
             # there is no mapping between model label and LS label

diff --git a/label_studio_ml/examples/yolo/control_models/timeline_labels.py b/label_studio_ml/examples/yolo/control_models/timeline_labels.py
@@ -42,7 +42,7 @@ def create_timelines(self, frame_results, video_path):
             for label_index, prob in enumerate(probs):
                 label = self.model.names[label_index]
 
-                if prob >= self.score_threshold:
+                if prob >= self.model_score_threshold:
                     # Start or continue a segment for this label
                     if ongoing_segments[label] is None:
                         # Start a new segment

diff --git a/label_studio_ml/examples/yolo/tests/test_choices.py b/label_studio_ml/examples/yolo/tests/test_choices.py
@@ -18,7 +18,7 @@
     """
     <View>
       <Image name="image" value="$image"/>
-      <Choices name="label" toName="image" score_threshold="0.53">
+      <Choices name="label" toName="image" model_score_threshold="0.53">
         <Choice value="Airplane" background="green"/>
         <Choice value="Car" background="blue" predicted_values="racer, cab"/>
       </Choices>
@@ -28,7 +28,7 @@
     """
     <View>
       <Image name="image" value="$image"/>
-      <Choices name="label" toName="image" choice="multiple" score_threshold="0.1">
+      <Choices name="label" toName="image" choice="multiple" model_score_threshold="0.1">
         <Choice value="Grille" background="green"/>
         <Choice value="Cab" background="blue" predicted_values="racer, cab"/>
       </Choices>
@@ -38,7 +38,7 @@
     """
     <View>
       <Image name="image" value="$image"/>
-      <Choices name="label" toName="image" choice="multiple" score_threshold="0.9">
+      <Choices name="label" toName="image" choice="multiple" model_score_threshold="0.9">
         <Choice value="Grille" background="green"/>
         <Choice value="Cab" background="blue" predicted_values="racer, cab"/>
       </Choices>

diff --git a/label_studio_ml/examples/yolo/tests/test_common.py b/label_studio_ml/examples/yolo/tests/test_common.py
@@ -40,7 +40,7 @@ def client():
     """
     <View>
       <Image name="image" value="$image"/>
-      <RectangleLabels name="label" toName="image" score_threshold="0.25">
+      <RectangleLabels name="label" toName="image" model_score_threshold="0.25">
         <Label value="Airplane" background="green"/>
         <Label value="Car" background="blue" predicted_values="car, truck"/>
       </RectangleLabels>

diff --git a/label_studio_ml/examples/yolo/tests/test_keypoint_labels.py b/label_studio_ml/examples/yolo/tests/test_keypoint_labels.py
@@ -18,7 +18,7 @@
     """
     <View>
       <KeyPointLabels name="keypoints" toName="image" 
-        score_threshold="0.85" model_point_threshold="0.95" model_add_bboxes="true" model_point_size="1"  
+        model_score_threshold="0.85" model_point_threshold="0.95" model_add_bboxes="true" model_point_size="1"  
       >
         <Label value="nose" predicted_values="person" model_index="0" background="red" />
 
@@ -57,7 +57,7 @@
     """
     <View>
       <KeyPointLabels name="keypoints" toName="image" 
-        score_threshold="0.85" model_point_threshold="0.96" model_add_bboxes="false" model_point_size="1"  
+        model_score_threshold="0.85" model_point_threshold="0.96" model_add_bboxes="false" model_point_size="1"  
       >
         <Label value="nose" predicted_values="person" model_index="0" background="red" />
 

diff --git a/label_studio_ml/examples/yolo/tests/test_polygon_labels.py b/label_studio_ml/examples/yolo/tests/test_polygon_labels.py
@@ -18,7 +18,7 @@
     """
     <View>
       <Image name="image" value="$image"/>
-      <PolygonLabels name="label" toName="image" score_threshold="0.6">
+      <PolygonLabels name="label" toName="image" model_score_threshold="0.6">
         <Label value="Person" background="green"/>
         <Label value="Truck" background="blue"/>
       </PolygonLabels>

diff --git a/label_studio_ml/examples/yolo/tests/test_rectangle_labels.py b/label_studio_ml/examples/yolo/tests/test_rectangle_labels.py
@@ -18,7 +18,7 @@
     """
     <View>
       <Image name="image" value="$image"/>
-      <RectangleLabels name="label" toName="image" score_threshold="0.25">
+      <RectangleLabels name="label" toName="image" model_score_threshold="0.25">
         <Label value="Airplane" background="green"/>
         <Label value="Car" background="blue" predicted_values="car, truck"/>
       </RectangleLabels>
@@ -28,13 +28,13 @@
     """
     <View>
       <Image name="image" value="$image"/>
-      <RectangleLabels name="label" toName="image" score_threshold="0.30">
+      <RectangleLabels name="label" toName="image" model_score_threshold="0.30">
         <Label value="Airplane" background="green"/>
         <Label value="Car" background="blue" predicted_values="car, truck"/>
       </RectangleLabels>
       
       <Image name="image2" value="$image2"/>
-      <RectangleLabels name="label2" toName="image2" score_threshold="0.90">
+      <RectangleLabels name="label2" toName="image2" model_score_threshold="0.90">
         <Label value="Person" background="green"/>
         <Label value="Animal" background="blue" predicted_values="cat,dog"/>
       </RectangleLabels>

diff --git a/label_studio_ml/examples/yolo/tests/test_rectangle_labels_obb.py b/label_studio_ml/examples/yolo/tests/test_rectangle_labels_obb.py
@@ -21,7 +21,7 @@
       <Image name="image" value="$image" zoom="true"/>
       
       <RectangleLabels name="rect" toName="image"
-                     score_threshold="0.1" model_obb="true">
+                     model_score_threshold="0.1" model_obb="true">
         <Label value="plane" background="red" 
             predicted_values="plane,helicopter"/>
         <Label value="vehicle" background="blue" 

diff --git a/label_studio_ml/model.py b/label_studio_ml/model.py
@@ -351,7 +351,7 @@ def build_label_map(self, tag_name: str, names: List[str]) -> Dict[str, str]:
 
             <View>
               <Image name="image" value="$image"/>
-              <RectangleLabels name="label" toName="image" score_threshold="0.25">
+              <RectangleLabels name="label" toName="image" model_score_threshold="0.25">
                 <Label value="Airplane" background="green"/>
                 <Label value="Car" background="blue" predicted_values="car, truck"/>
               </RectangleLabels>