SAM: Fix labels not being used in predict mode. (#16642)

Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
2024-10-03 02:43:46 +08:00 · 2024-10-03 02:43:46 +08:00 · 7d6528a41a
commit 7d6528a41a
parent c773f27653
1 changed files with 1 additions and 0 deletions
--- a/ultralytics/models/sam/predict.py
+++ b/ultralytics/models/sam/predict.py
@ -196,6 +196,7 @@ class Predictor(BasePredictor):
        bboxes = self.prompts.pop("bboxes", bboxes)
        points = self.prompts.pop("points", points)
        masks = self.prompts.pop("masks", masks)
+        labels = self.prompts.pop("labels", labels)

        if all(i is None for i in [bboxes, points, masks]):
            return self.generate(im, *args, **kwargs)