ultralytics 8.1.29__py3-none-any.whl → 8.3.63__py3-none-any.whl

ultralytics/models/fastsam/predict.py

@@ -1,86 +1,150 @@
-# Ultralytics YOLO 🚀, AGPL-3.0 license
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
 import torch
+from PIL import Image
 
-from ultralytics.engine.results import Results
-from ultralytics.models.fastsam.utils import bbox_iou
-from ultralytics.models.yolo.detect.predict import DetectionPredictor
-from ultralytics.utils import DEFAULT_CFG, ops
+from ultralytics.models.yolo.segment import SegmentationPredictor
+from ultralytics.utils import DEFAULT_CFG, checks
+from ultralytics.utils.metrics import box_iou
+from ultralytics.utils.ops import scale_masks
 
+from .utils import adjust_bboxes_to_image_border
 
-class FastSAMPredictor(DetectionPredictor):
+
+class FastSAMPredictor(SegmentationPredictor):
     """
     FastSAMPredictor is specialized for fast SAM (Segment Anything Model) segmentation prediction tasks in Ultralytics
     YOLO framework.
 
-    This class extends the DetectionPredictor, customizing the prediction pipeline specifically for fast SAM.
-    It adjusts post-processing steps to incorporate mask prediction and non-max suppression while optimizing
-    for single-class segmentation.
-
-    Attributes:
-        cfg (dict): Configuration parameters for prediction.
-        overrides (dict, optional): Optional parameter overrides for custom behavior.
-        _callbacks (dict, optional): Optional list of callback functions to be invoked during prediction.
+    This class extends the SegmentationPredictor, customizing the prediction pipeline specifically for fast SAM. It
+    adjusts post-processing steps to incorporate mask prediction and non-max suppression while optimizing for single-
+    class segmentation.
     """
 
     def __init__(self, cfg=DEFAULT_CFG, overrides=None, _callbacks=None):
+        """Initializes a FastSAMPredictor for fast SAM segmentation tasks in Ultralytics YOLO framework."""
+        super().__init__(cfg, overrides, _callbacks)
+        self.prompts = {}
+
+    def postprocess(self, preds, img, orig_imgs):
+        """Applies box postprocess for FastSAM predictions."""
+        bboxes = self.prompts.pop("bboxes", None)
+        points = self.prompts.pop("points", None)
+        labels = self.prompts.pop("labels", None)
+        texts = self.prompts.pop("texts", None)
+        results = super().postprocess(preds, img, orig_imgs)
+        for result in results:
+            full_box = torch.tensor(
+                [0, 0, result.orig_shape[1], result.orig_shape[0]], device=preds[0].device, dtype=torch.float32
+            )
+            boxes = adjust_bboxes_to_image_border(result.boxes.xyxy, result.orig_shape)
+            idx = torch.nonzero(box_iou(full_box[None], boxes) > 0.9).flatten()
+            if idx.numel() != 0:
+                result.boxes.xyxy[idx] = full_box
+
+        return self.prompt(results, bboxes=bboxes, points=points, labels=labels, texts=texts)
+
+    def prompt(self, results, bboxes=None, points=None, labels=None, texts=None):
         """
-        Initializes the FastSAMPredictor class, inheriting from DetectionPredictor and setting the task to 'segment'.
+        Internal function for image segmentation inference based on cues like bounding boxes, points, and masks.
+        Leverages SAM's specialized architecture for prompt-based, real-time segmentation.
 
         Args:
-            cfg (dict): Configuration parameters for prediction.
-            overrides (dict, optional): Optional parameter overrides for custom behavior.
-            _callbacks (dict, optional): Optional list of callback functions to be invoked during prediction.
+            results (Results | List[Results]): The original inference results from FastSAM models without any prompts.
+            bboxes (np.ndarray | List, optional): Bounding boxes with shape (N, 4), in XYXY format.
+            points (np.ndarray | List, optional): Points indicating object locations with shape (N, 2), in pixels.
+            labels (np.ndarray | List, optional): Labels for point prompts, shape (N, ). 1 = foreground, 0 = background.
+            texts (str | List[str], optional): Textual prompts, a list contains string objects.
+
+        Returns:
+            (List[Results]): The output results determined by prompts.
         """
-        super().__init__(cfg, overrides, _callbacks)
-        self.args.task = "segment"
+        if bboxes is None and points is None and texts is None:
+            return results
+        prompt_results = []
+        if not isinstance(results, list):
+            results = [results]
+        for result in results:
+            if len(result) == 0:
+                prompt_results.append(result)
+                continue
+            masks = result.masks.data
+            if masks.shape[1:] != result.orig_shape:
+                masks = scale_masks(masks[None], result.orig_shape)[0]
+            # bboxes prompt
+            idx = torch.zeros(len(result), dtype=torch.bool, device=self.device)
+            if bboxes is not None:
+                bboxes = torch.as_tensor(bboxes, dtype=torch.int32, device=self.device)
+                bboxes = bboxes[None] if bboxes.ndim == 1 else bboxes
+                bbox_areas = (bboxes[:, 3] - bboxes[:, 1]) * (bboxes[:, 2] - bboxes[:, 0])
+                mask_areas = torch.stack([masks[:, b[1] : b[3], b[0] : b[2]].sum(dim=(1, 2)) for b in bboxes])
+                full_mask_areas = torch.sum(masks, dim=(1, 2))
 
-    def postprocess(self, preds, img, orig_imgs):
+                union = bbox_areas[:, None] + full_mask_areas - mask_areas
+                idx[torch.argmax(mask_areas / union, dim=1)] = True
+            if points is not None:
+                points = torch.as_tensor(points, dtype=torch.int32, device=self.device)
+                points = points[None] if points.ndim == 1 else points
+                if labels is None:
+                    labels = torch.ones(points.shape[0])
+                labels = torch.as_tensor(labels, dtype=torch.int32, device=self.device)
+                assert len(labels) == len(points), (
+                    f"Excepted `labels` got same size as `point`, but got {len(labels)} and {len(points)}"
+                )
+                point_idx = (
+                    torch.ones(len(result), dtype=torch.bool, device=self.device)
+                    if labels.sum() == 0  # all negative points
+                    else torch.zeros(len(result), dtype=torch.bool, device=self.device)
+                )
+                for point, label in zip(points, labels):
+                    point_idx[torch.nonzero(masks[:, point[1], point[0]], as_tuple=True)[0]] = bool(label)
+                idx |= point_idx
+            if texts is not None:
+                if isinstance(texts, str):
+                    texts = [texts]
+                crop_ims, filter_idx = [], []
+                for i, b in enumerate(result.boxes.xyxy.tolist()):
+                    x1, y1, x2, y2 = (int(x) for x in b)
+                    if masks[i].sum() <= 100:
+                        filter_idx.append(i)
+                        continue
+                    crop_ims.append(Image.fromarray(result.orig_img[y1:y2, x1:x2, ::-1]))
+                similarity = self._clip_inference(crop_ims, texts)
+                text_idx = torch.argmax(similarity, dim=-1)  # (M, )
+                if len(filter_idx):
+                    text_idx += (torch.tensor(filter_idx, device=self.device)[None] <= int(text_idx)).sum(0)
+                idx[text_idx] = True
+
+            prompt_results.append(result[idx])
+
+        return prompt_results
+
+    def _clip_inference(self, images, texts):
         """
-        Perform post-processing steps on predictions, including non-max suppression and scaling boxes to original image
-        size, and returns the final results.
+        CLIP Inference process.
 
         Args:
-            preds (list): The raw output predictions from the model.
-            img (torch.Tensor): The processed image tensor.
-            orig_imgs (list | torch.Tensor): The original image or list of images.
+            images (List[PIL.Image]): A list of source images and each of them should be PIL.Image type with RGB channel order.
+            texts (List[str]): A list of prompt texts and each of them should be string object.
 
         Returns:
-            (list): A list of Results objects, each containing processed boxes, masks, and other metadata.
+            (torch.Tensor): The similarity between given images and texts.
         """
-        p = ops.non_max_suppression(
-            preds[0],
-            self.args.conf,
-            self.args.iou,
-            agnostic=self.args.agnostic_nms,
-            max_det=self.args.max_det,
-            nc=1,  # set to 1 class since SAM has no class predictions
-            classes=self.args.classes,
-        )
-        full_box = torch.zeros(p[0].shape[1], device=p[0].device)
-        full_box[2], full_box[3], full_box[4], full_box[6:] = img.shape[3], img.shape[2], 1.0, 1.0
-        full_box = full_box.view(1, -1)
-        critical_iou_index = bbox_iou(full_box[0][:4], p[0][:, :4], iou_thres=0.9, image_shape=img.shape[2:])
-        if critical_iou_index.numel() != 0:
-            full_box[0][4] = p[0][critical_iou_index][:, 4]
-            full_box[0][6:] = p[0][critical_iou_index][:, 6:]
-            p[0][critical_iou_index] = full_box
-
-        if not isinstance(orig_imgs, list):  # input images are a torch.Tensor, not a list
-            orig_imgs = ops.convert_torch2numpy_batch(orig_imgs)
-
-        results = []
-        proto = preds[1][-1] if len(preds[1]) == 3 else preds[1]  # second output is len 3 if pt, but only 1 if exported
-        for i, pred in enumerate(p):
-            orig_img = orig_imgs[i]
-            img_path = self.batch[0][i]
-            if not len(pred):  # save empty boxes
-                masks = None
-            elif self.args.retina_masks:
-                pred[:, :4] = ops.scale_boxes(img.shape[2:], pred[:, :4], orig_img.shape)
-                masks = ops.process_mask_native(proto[i], pred[:, 6:], pred[:, :4], orig_img.shape[:2])  # HWC
-            else:
-                masks = ops.process_mask(proto[i], pred[:, 6:], pred[:, :4], img.shape[2:], upsample=True)  # HWC
-                pred[:, :4] = ops.scale_boxes(img.shape[2:], pred[:, :4], orig_img.shape)
-            results.append(Results(orig_img, path=img_path, names=self.model.names, boxes=pred[:, :6], masks=masks))
-        return results
+        try:
+            import clip
+        except ImportError:
+            checks.check_requirements("git+https://github.com/ultralytics/CLIP.git")
+            import clip
+        if (not hasattr(self, "clip_model")) or (not hasattr(self, "clip_preprocess")):
+            self.clip_model, self.clip_preprocess = clip.load("ViT-B/32", device=self.device)
+        images = torch.stack([self.clip_preprocess(image).to(self.device) for image in images])
+        tokenized_text = clip.tokenize(texts).to(self.device)
+        image_features = self.clip_model.encode_image(images)
+        text_features = self.clip_model.encode_text(tokenized_text)
+        image_features /= image_features.norm(dim=-1, keepdim=True)  # (N, 512)
+        text_features /= text_features.norm(dim=-1, keepdim=True)  # (M, 512)
+        return (image_features * text_features[:, None]).sum(-1)  # (M, N)
+
+    def set_prompts(self, prompts):
+        """Set prompts in advance."""
+        self.prompts = prompts

ultralytics/models/fastsam/utils.py

@@ -1,6 +1,4 @@
-# Ultralytics YOLO 🚀, AGPL-3.0 license
-
-import torch
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
 
 def adjust_bboxes_to_image_border(boxes, image_shape, threshold=20):
@@ -15,7 +13,6 @@ def adjust_bboxes_to_image_border(boxes, image_shape, threshold=20):
     Returns:
         adjusted_boxes (torch.Tensor): adjusted bounding boxes
     """
-
     # Image dimensions
     h, w = image_shape
 
@@ -25,43 +22,3 @@ def adjust_bboxes_to_image_border(boxes, image_shape, threshold=20):
     boxes[boxes[:, 2] > w - threshold, 2] = w  # x2
     boxes[boxes[:, 3] > h - threshold, 3] = h  # y2
     return boxes
-
-
-def bbox_iou(box1, boxes, iou_thres=0.9, image_shape=(640, 640), raw_output=False):
-    """
-    Compute the Intersection-Over-Union of a bounding box with respect to an array of other bounding boxes.
-
-    Args:
-        box1 (torch.Tensor): (4, )
-        boxes (torch.Tensor): (n, 4)
-        iou_thres (float): IoU threshold
-        image_shape (tuple): (height, width)
-        raw_output (bool): If True, return the raw IoU values instead of the indices
-
-    Returns:
-        high_iou_indices (torch.Tensor): Indices of boxes with IoU > thres
-    """
-    boxes = adjust_bboxes_to_image_border(boxes, image_shape)
-    # Obtain coordinates for intersections
-    x1 = torch.max(box1[0], boxes[:, 0])
-    y1 = torch.max(box1[1], boxes[:, 1])
-    x2 = torch.min(box1[2], boxes[:, 2])
-    y2 = torch.min(box1[3], boxes[:, 3])
-
-    # Compute the area of intersection
-    intersection = (x2 - x1).clamp(0) * (y2 - y1).clamp(0)
-
-    # Compute the area of both individual boxes
-    box1_area = (box1[2] - box1[0]) * (box1[3] - box1[1])
-    box2_area = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
-
-    # Compute the area of union
-    union = box1_area + box2_area - intersection
-
-    # Compute the IoU
-    iou = intersection / union  # Should be shape (n, )
-    if raw_output:
-        return 0 if iou.numel() == 0 else iou
-
-    # return indices of boxes with IoU > thres
-    return torch.nonzero(iou > iou_thres).flatten()

ultralytics/models/fastsam/val.py

@@ -1,4 +1,4 @@
-# Ultralytics YOLO 🚀, AGPL-3.0 license
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
 from ultralytics.models.yolo.segment import SegmentationValidator
 from ultralytics.utils.metrics import SegmentMetrics

ultralytics/models/nas/__init__.py

@@ -1,4 +1,4 @@
-# Ultralytics YOLO 🚀, AGPL-3.0 license
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
 from .model import NAS
 from .predict import NASPredictor

ultralytics/models/nas/model.py

@@ -1,4 +1,4 @@
-# Ultralytics YOLO 🚀, AGPL-3.0 license
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 """
 YOLO-NAS model interface.
 
@@ -6,8 +6,8 @@ Example:
     ```python
     from ultralytics import NAS
 
-    model = NAS('yolo_nas_s')
-    results = model.predict('ultralytics/assets/bus.jpg')
+    model = NAS("yolo_nas_s")
+    results = model.predict("ultralytics/assets/bus.jpg")
     ```
 """
 
@@ -16,7 +16,9 @@ from pathlib import Path
 import torch
 
 from ultralytics.engine.model import Model
-from ultralytics.utils.torch_utils import model_info, smart_inference_mode
+from ultralytics.utils.downloads import attempt_download_asset
+from ultralytics.utils.torch_utils import model_info
+
 from .predict import NASPredictor
 from .val import NASValidator
 
@@ -32,8 +34,8 @@ class NAS(Model):
         ```python
         from ultralytics import NAS
 
-        model = NAS('yolo_nas_s')
-        results = model.predict('ultralytics/assets/bus.jpg')
+        model = NAS("yolo_nas_s")
+        results = model.predict("ultralytics/assets/bus.jpg")
         ```
 
     Attributes:
@@ -45,19 +47,28 @@ class NAS(Model):
 
     def __init__(self, model="yolo_nas_s.pt") -> None:
         """Initializes the NAS model with the provided or default 'yolo_nas_s.pt' model."""
-        assert Path(model).suffix not in (".yaml", ".yml"), "YOLO-NAS models only support pre-trained models."
+        assert Path(model).suffix not in {".yaml", ".yml"}, "YOLO-NAS models only support pre-trained models."
         super().__init__(model, task="detect")
 
-    @smart_inference_mode()
-    def _load(self, weights: str, task: str):
+    def _load(self, weights: str, task=None) -> None:
         """Loads an existing NAS model weights or creates a new NAS model with pretrained weights if not provided."""
         import super_gradients
 
         suffix = Path(weights).suffix
         if suffix == ".pt":
-            self.model = torch.load(weights)
+            self.model = torch.load(attempt_download_asset(weights))
+
         elif suffix == "":
             self.model = super_gradients.training.models.get(weights, pretrained_weights="coco")
+
+        # Override the forward method to ignore additional arguments
+        def new_forward(x, *args, **kwargs):
+            """Ignore additional __call__ arguments."""
+            return self.model._original_forward(x)
+
+        self.model._original_forward = self.model.forward
+        self.model.forward = new_forward
+
         # Standardize model
         self.model.fuse = lambda verbose=True: self.model
         self.model.stride = torch.tensor([32])

ultralytics/models/nas/predict.py

@@ -1,4 +1,4 @@
-# Ultralytics YOLO 🚀, AGPL-3.0 license
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
 import torch
 
@@ -22,7 +22,7 @@ class NASPredictor(BasePredictor):
         ```python
         from ultralytics import NAS
 
-        model = NAS('yolo_nas_s')
+        model = NAS("yolo_nas_s")
         predictor = model.predictor
         # Assumes that raw_preds, img, orig_imgs are available
         results = predictor.postprocess(raw_preds, img, orig_imgs)
@@ -34,7 +34,6 @@ class NASPredictor(BasePredictor):
 
     def postprocess(self, preds_in, img, orig_imgs):
         """Postprocess predictions and returns a list of Results objects."""
-
         # Cat boxes and class scores
         boxes = ops.xyxy2xywh(preds_in[0][0])
         preds = torch.cat((boxes, preds_in[0][1]), -1).permute(0, 2, 1)
@@ -52,9 +51,7 @@ class NASPredictor(BasePredictor):
             orig_imgs = ops.convert_torch2numpy_batch(orig_imgs)
 
         results = []
-        for i, pred in enumerate(preds):
-            orig_img = orig_imgs[i]
+        for pred, orig_img, img_path in zip(preds, orig_imgs, self.batch[0]):
             pred[:, :4] = ops.scale_boxes(img.shape[2:], pred[:, :4], orig_img.shape)
-            img_path = self.batch[0][i]
             results.append(Results(orig_img, path=img_path, names=self.model.names, boxes=pred))
         return results

ultralytics/models/nas/val.py

@@ -1,4 +1,4 @@
-# Ultralytics YOLO 🚀, AGPL-3.0 license
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
 import torch
 
@@ -17,14 +17,14 @@ class NASValidator(DetectionValidator):
     ultimately producing the final detections.
 
     Attributes:
-        args (Namespace): Namespace containing various configurations for post-processing, such as confidence and IoU thresholds.
+        args (Namespace): Namespace containing various configurations for post-processing, such as confidence and IoU.
         lb (torch.Tensor): Optional tensor for multilabel NMS.
 
     Example:
         ```python
         from ultralytics import NAS
 
-        model = NAS('yolo_nas_s')
+        model = NAS("yolo_nas_s")
         validator = model.validator
         # Assumes that raw_preds are available
         final_preds = validator.postprocess(raw_preds)
@@ -44,7 +44,7 @@ class NASValidator(DetectionValidator):
             self.args.iou,
             labels=self.lb,
             multi_label=False,
-            agnostic=self.args.single_cls,
+            agnostic=self.args.single_cls or self.args.agnostic_nms,
             max_det=self.args.max_det,
             max_time_img=0.5,
         )

ultralytics/models/rtdetr/__init__.py

@@ -1,4 +1,4 @@
-# Ultralytics YOLO 🚀, AGPL-3.0 license
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
 from .model import RTDETR
 from .predict import RTDETRPredictor

ultralytics/models/rtdetr/model.py

@@ -1,4 +1,4 @@
-# Ultralytics YOLO 🚀, AGPL-3.0 license
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 """
 Interface for Baidu's RT-DETR, a Vision Transformer-based real-time object detector. RT-DETR offers real-time
 performance and high accuracy, excelling in accelerated backends like CUDA with TensorRT. It features an efficient

ultralytics/models/rtdetr/predict.py

@@ -1,4 +1,4 @@
-# Ultralytics YOLO 🚀, AGPL-3.0 license
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
 import torch
 
@@ -21,7 +21,7 @@ class RTDETRPredictor(BasePredictor):
         from ultralytics.utils import ASSETS
         from ultralytics.models.rtdetr import RTDETRPredictor
 
-        args = dict(model='rtdetr-l.pt', source=ASSETS)
+        args = dict(model="rtdetr-l.pt", source=ASSETS)
         predictor = RTDETRPredictor(overrides=args)
         predictor.predict_cli()
         ```
@@ -56,18 +56,16 @@ class RTDETRPredictor(BasePredictor):
             orig_imgs = ops.convert_torch2numpy_batch(orig_imgs)
 
         results = []
-        for i, bbox in enumerate(bboxes):  # (300, 4)
+        for bbox, score, orig_img, img_path in zip(bboxes, scores, orig_imgs, self.batch[0]):  # (300, 4)
             bbox = ops.xywh2xyxy(bbox)
-            score, cls = scores[i].max(-1, keepdim=True)  # (300, 1)
-            idx = score.squeeze(-1) > self.args.conf  # (300, )
+            max_score, cls = score.max(-1, keepdim=True)  # (300, 1)
+            idx = max_score.squeeze(-1) > self.args.conf  # (300, )
             if self.args.classes is not None:
                 idx = (cls == torch.tensor(self.args.classes, device=cls.device)).any(1) & idx
-            pred = torch.cat([bbox, score, cls], dim=-1)[idx]  # filter
-            orig_img = orig_imgs[i]
+            pred = torch.cat([bbox, max_score, cls], dim=-1)[idx]  # filter
             oh, ow = orig_img.shape[:2]
             pred[..., [0, 2]] *= ow
             pred[..., [1, 3]] *= oh
-            img_path = self.batch[0][i]
             results.append(Results(orig_img, path=img_path, names=self.model.names, boxes=pred))
         return results
 

ultralytics/models/rtdetr/train.py

@@ -1,4 +1,4 @@
-# Ultralytics YOLO 🚀, AGPL-3.0 license
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
 from copy import copy
 
@@ -7,6 +7,7 @@ import torch
 from ultralytics.models.yolo.detect import DetectionTrainer
 from ultralytics.nn.tasks import RTDETRDetectionModel
 from ultralytics.utils import RANK, colorstr
+
 from .val import RTDETRDataset, RTDETRValidator
 
 
@@ -24,7 +25,7 @@ class RTDETRTrainer(DetectionTrainer):
         ```python
         from ultralytics.models.rtdetr.train import RTDETRTrainer
 
-        args = dict(model='rtdetr-l.yaml', data='coco8.yaml', imgsz=640, epochs=3)
+        args = dict(model="rtdetr-l.yaml", data="coco8.yaml", imgsz=640, epochs=3)
         trainer = RTDETRTrainer(overrides=args)
         trainer.train()
         ```
@@ -67,8 +68,11 @@ class RTDETRTrainer(DetectionTrainer):
             hyp=self.args,
             rect=False,
             cache=self.args.cache or None,
+            single_cls=self.args.single_cls or False,
             prefix=colorstr(f"{mode}: "),
+            classes=self.args.classes,
             data=self.data,
+            fraction=self.args.fraction if mode == "train" else 1.0,
         )
 
     def get_validator(self):

ultralytics/models/rtdetr/val.py

@@ -1,4 +1,4 @@
-# Ultralytics YOLO 🚀, AGPL-3.0 license
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
 import torch
 
@@ -62,7 +62,7 @@ class RTDETRValidator(DetectionValidator):
         ```python
         from ultralytics.models.rtdetr import RTDETRValidator
 
-        args = dict(model='rtdetr-l.pt', data='coco8.yaml')
+        args = dict(model="rtdetr-l.pt", data="coco8.yaml")
         validator = RTDETRValidator(args=args)
         validator()
         ```
@@ -125,7 +125,7 @@ class RTDETRValidator(DetectionValidator):
             bbox = ops.xywh2xyxy(bbox)  # target boxes
             bbox[..., [0, 2]] *= ori_shape[1]  # native-space pred
             bbox[..., [1, 3]] *= ori_shape[0]  # native-space pred
-        return dict(cls=cls, bbox=bbox, ori_shape=ori_shape, imgsz=imgsz, ratio_pad=ratio_pad)
+        return {"cls": cls, "bbox": bbox, "ori_shape": ori_shape, "imgsz": imgsz, "ratio_pad": ratio_pad}
 
     def _prepare_pred(self, pred, pbatch):
         """Prepares and returns a batch with transformed bounding boxes and class labels."""

ultralytics/models/sam/__init__.py

@@ -1,6 +1,6 @@
-# Ultralytics YOLO 🚀, AGPL-3.0 license
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
 from .model import SAM
-from .predict import Predictor
+from .predict import Predictor, SAM2Predictor, SAM2VideoPredictor
 
-__all__ = "SAM", "Predictor"  # tuple or list
+__all__ = "SAM", "Predictor", "SAM2Predictor", "SAM2VideoPredictor"  # tuple or list