Fixed TypeError in YOLO model by ensuring setText accepts only string…

… arguments (#327)
CVHub520 · Mar 20, 2024 · 6f4447a · 6f4447a
1 parent ea435b0
commit 6f4447a
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/anylabeling/services/auto_labeling/__base__/yolo.py b/anylabeling/services/auto_labeling/__base__/yolo.py
@@ -311,7 +311,7 @@ def predict_shapes(self, image, image_path=None):
                 shape.fill_color = "#000000"
                 shape.line_color = "#000000"
                 shape.line_width = 1
-                shape.label = self.classes[int(class_id)]
+                shape.label = str(self.classes[int(class_id)])
                 shape.selected = False
                 shapes.append(shape)
             if self.task == "seg":
@@ -339,7 +339,7 @@ def predict_shapes(self, image, image_path=None):
                 shape.fill_color = "#000000"
                 shape.line_color = "#000000"
                 shape.line_width = 1
-                shape.label = self.classes[int(class_id)]
+                shape.label = str(self.classes[int(class_id)])
                 shape.selected = False
                 shapes.append(shape)
             if self.task == "obb":
@@ -360,7 +360,7 @@ def predict_shapes(self, image, image_path=None):
                 shape.fill_color = "#000000"
                 shape.line_color = "#000000"
                 shape.line_width = 1
-                shape.label = self.classes[int(class_id)]
+                shape.label = str(self.classes[int(class_id)])
                 shape.selected = False
                 shapes.append(shape)
 

diff --git a/anylabeling/services/auto_labeling/rtmdet_pose.py b/anylabeling/services/auto_labeling/rtmdet_pose.py
@@ -107,7 +107,7 @@ def predict_shapes(self, image, image_path=None):
                 if score < self.kpt_thr:
                     continue
                 point_shape = Shape(
-                    label=self.kpt_classes[j],
+                    label=str(self.kpt_classes[j]),
                     shape_type="point",
                     group_id=int(i),
                 )