Make labels unique

luxonis · Mar 13, 2024 · f1767d0 · f1767d0
1 parent ef99c46
commit f1767d0
Showing 1 changed file with 9 additions and 5 deletions.
diff --git a/datadreamer/dataset_annotation/clip_annotator.py b/datadreamer/dataset_annotation/clip_annotator.py
@@ -91,11 +91,15 @@ def annotate_batch(
         if synonym_dict is not None:
             for prob in probs:
                 labels.append(
-                    np.array(
-                        [
-                            synonym_dict_rev[label.item()]
-                            for label in torch.where(prob > conf_threshold)[0].numpy()
-                        ]
+                    np.unique(
+                        np.array(
+                            [
+                                synonym_dict_rev[label.item()]
+                                for label in torch.where(prob > conf_threshold)[
+                                    0
+                                ].numpy()
+                            ]
+                        )
                     )
                 )
         else: