dgenerate-ultralytics-headless 8.3.236__tar.gz → 8.3.237__tar.gz

{dgenerate_ultralytics_headless-8.3.236 → dgenerate_ultralytics_headless-8.3.237}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dgenerate-ultralytics-headless
-Version: 8.3.236
+Version: 8.3.237
 Summary: Automatically built Ultralytics package with python-opencv-headless dependency instead of python-opencv
 Author-email: Glenn Jocher <glenn.jocher@ultralytics.com>, Jing Qiu <jing.qiu@ultralytics.com>
 Maintainer-email: Ultralytics <hello@ultralytics.com>

{dgenerate_ultralytics_headless-8.3.236 → dgenerate_ultralytics_headless-8.3.237}/dgenerate_ultralytics_headless.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dgenerate-ultralytics-headless
-Version: 8.3.236
+Version: 8.3.237
 Summary: Automatically built Ultralytics package with python-opencv-headless dependency instead of python-opencv
 Author-email: Glenn Jocher <glenn.jocher@ultralytics.com>, Jing Qiu <jing.qiu@ultralytics.com>
 Maintainer-email: Ultralytics <hello@ultralytics.com>

{dgenerate_ultralytics_headless-8.3.236 → dgenerate_ultralytics_headless-8.3.237}/dgenerate_ultralytics_headless.egg-info/SOURCES.txt

@@ -168,6 +168,7 @@ ultralytics/models/rtdetr/val.py
 ultralytics/models/sam/__init__.py
 ultralytics/models/sam/amg.py
 ultralytics/models/sam/build.py
+ultralytics/models/sam/build_sam3.py
 ultralytics/models/sam/model.py
 ultralytics/models/sam/predict.py
 ultralytics/models/sam/modules/__init__.py
@@ -179,6 +180,18 @@ ultralytics/models/sam/modules/sam.py
 ultralytics/models/sam/modules/tiny_encoder.py
 ultralytics/models/sam/modules/transformer.py
 ultralytics/models/sam/modules/utils.py
+ultralytics/models/sam/sam3/__init__.py
+ultralytics/models/sam/sam3/decoder.py
+ultralytics/models/sam/sam3/encoder.py
+ultralytics/models/sam/sam3/geometry_encoders.py
+ultralytics/models/sam/sam3/maskformer_segmentation.py
+ultralytics/models/sam/sam3/model_misc.py
+ultralytics/models/sam/sam3/necks.py
+ultralytics/models/sam/sam3/sam3_image.py
+ultralytics/models/sam/sam3/text_encoder_ve.py
+ultralytics/models/sam/sam3/tokenizer_ve.py
+ultralytics/models/sam/sam3/vitdet.py
+ultralytics/models/sam/sam3/vl_combiner.py
 ultralytics/models/utils/__init__.py
 ultralytics/models/utils/loss.py
 ultralytics/models/utils/ops.py

{dgenerate_ultralytics_headless-8.3.236 → dgenerate_ultralytics_headless-8.3.237}/ultralytics/__init__.py

@@ -1,6 +1,6 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
-__version__ = "8.3.236"
+__version__ = "8.3.237"
 
 import importlib
 import os

{dgenerate_ultralytics_headless-8.3.236 → dgenerate_ultralytics_headless-8.3.237}/ultralytics/engine/exporter.py

@@ -392,8 +392,10 @@ class Exporter:
         if self.args.half and self.args.int8:
             LOGGER.warning("half=True and int8=True are mutually exclusive, setting half=False.")
             self.args.half = False
-        if self.args.half and (onnx or jit) and self.device.type == "cpu":
-            LOGGER.warning("half=True only compatible with GPU export, i.e. use device=0, setting half=False.")
+        if self.args.half and jit and self.device.type == "cpu":
+            LOGGER.warning(
+                "half=True only compatible with GPU export for TorchScript, i.e. use device=0, setting half=False."
+            )
             self.args.half = False
         self.imgsz = check_imgsz(self.args.imgsz, stride=model.stride, min_dim=2)  # check image size
         if self.args.optimize:
@@ -716,6 +718,16 @@ class Exporter:
             LOGGER.info(f"{prefix} limiting IR version {model_onnx.ir_version} to 10 for ONNXRuntime compatibility...")
             model_onnx.ir_version = 10
 
+        # FP16 conversion for CPU export (GPU exports are already FP16 from model.half() during tracing)
+        if self.args.half and self.device.type == "cpu":
+            try:
+                from onnxruntime.transformers import float16
+
+                LOGGER.info(f"{prefix} converting to FP16...")
+                model_onnx = float16.convert_float_to_float16(model_onnx, keep_io_types=True)
+            except Exception as e:
+                LOGGER.warning(f"{prefix} FP16 conversion failure: {e}")
+
         onnx.save(model_onnx, f)
         return f
 
@@ -1220,10 +1232,9 @@ class Exporter:
                 f"{sudo}mkdir -p /etc/apt/keyrings",
                 f"curl -fsSL https://packages.cloud.google.com/apt/doc/apt-key.gpg | {sudo}gpg --dearmor -o /etc/apt/keyrings/google.gpg",
                 f'echo "deb [signed-by=/etc/apt/keyrings/google.gpg] https://packages.cloud.google.com/apt coral-edgetpu-stable main" | {sudo}tee /etc/apt/sources.list.d/coral-edgetpu.list',
-                f"{sudo}apt-get update",
-                f"{sudo}apt-get install -y edgetpu-compiler",
             ):
                 subprocess.run(c, shell=True, check=True)
+            check_apt_requirements(["edgetpu-compiler"])
 
         ver = subprocess.run(cmd, shell=True, capture_output=True, check=True).stdout.decode().rsplit(maxsplit=1)[-1]
         LOGGER.info(f"\n{prefix} starting export with Edge TPU compiler {ver}...")
@@ -1301,16 +1312,12 @@ class Exporter:
             java_version = int(version_match.group(1)) if version_match else 0
             assert java_version >= 17, "Java version too old"
         except (FileNotFoundError, subprocess.CalledProcessError, AssertionError):
-            cmd = None
             if IS_UBUNTU or IS_DEBIAN_TRIXIE:
                 LOGGER.info(f"\n{prefix} installing Java 21 for Ubuntu...")
-                cmd = (["sudo"] if is_sudo_available() else []) + ["apt-get", "install", "-y", "openjdk-21-jre"]
+                check_apt_requirements(["openjdk-21-jre"])
             elif IS_RASPBERRYPI or IS_DEBIAN_BOOKWORM:
                 LOGGER.info(f"\n{prefix} installing Java 17 for Raspberry Pi or Debian ...")
-                cmd = (["sudo"] if is_sudo_available() else []) + ["apt-get", "install", "-y", "openjdk-17-jre"]
-
-            if cmd:
-                subprocess.run(cmd, check=True)
+                check_apt_requirements(["openjdk-17-jre"])
 
         return torch2imx(
             self.model,

{dgenerate_ultralytics_headless-8.3.236 → dgenerate_ultralytics_headless-8.3.237}/ultralytics/engine/predictor.py

@@ -244,14 +244,15 @@ class BasePredictor:
         for _ in gen:  # sourcery skip: remove-empty-nested-block, noqa
             pass
 
-    def setup_source(self, source):
+    def setup_source(self, source, stride: int | None = None):
         """Set up source and inference mode.
 
         Args:
             source (str | Path | list[str] | list[Path] | list[np.ndarray] | np.ndarray | torch.Tensor): Source for
                 inference.
+            stride (int, optional): Model stride for image size checking.
         """
-        self.imgsz = check_imgsz(self.args.imgsz, stride=self.model.stride, min_dim=2)  # check image size
+        self.imgsz = check_imgsz(self.args.imgsz, stride=stride or self.model.stride, min_dim=2)  # check image size
         self.dataset = load_inference_source(
             source=source,
             batch=self.args.batch,

{dgenerate_ultralytics_headless-8.3.236 → dgenerate_ultralytics_headless-8.3.237}/ultralytics/engine/trainer.py

@@ -812,6 +812,14 @@ class BaseTrainer:
                     "device",
                     "close_mosaic",
                     "augmentations",
+                    "save_period",
+                    "workers",
+                    "cache",
+                    "patience",
+                    "time",
+                    "freeze",
+                    "val",
+                    "plots",
                 ):  # allow arg updates to reduce memory or update device on resume
                     if k in overrides:
                         setattr(self.args, k, overrides[k])

{dgenerate_ultralytics_headless-8.3.236 → dgenerate_ultralytics_headless-8.3.237}/ultralytics/models/rtdetr/val.py

@@ -85,7 +85,7 @@ class RTDETRDataset(YOLODataset):
             transforms = v8_transforms(self, self.imgsz, hyp, stretch=True)
         else:
             # transforms = Compose([LetterBox(new_shape=(self.imgsz, self.imgsz), auto=False, scale_fill=True)])
-            transforms = Compose([lambda x: {**x, **{"ratio_pad": [x["ratio_pad"], [0, 0]]}}])
+            transforms = Compose([])
         transforms.append(
             Format(
                 bbox_format="xywh",
@@ -150,6 +150,10 @@ class RTDETRValidator(DetectionValidator):
             data=self.data,
         )
 
+    def scale_preds(self, predn: dict[str, torch.Tensor], pbatch: dict[str, Any]) -> dict[str, torch.Tensor]:
+        """Scales predictions to the original image size."""
+        return predn
+
     def postprocess(
         self, preds: torch.Tensor | list[torch.Tensor] | tuple[torch.Tensor]
     ) -> list[dict[str, torch.Tensor]]:

dgenerate_ultralytics_headless-8.3.237/ultralytics/models/sam/__init__.py

@@ -0,0 +1,25 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+from .model import SAM
+from .predict import (
+    Predictor,
+    SAM2DynamicInteractivePredictor,
+    SAM2Predictor,
+    SAM2VideoPredictor,
+    SAM3Predictor,
+    SAM3SemanticPredictor,
+    SAM3VideoPredictor,
+    SAM3VideoSemanticPredictor,
+)
+
+__all__ = (
+    "SAM",
+    "Predictor",
+    "SAM2DynamicInteractivePredictor",
+    "SAM2Predictor",
+    "SAM2VideoPredictor",
+    "SAM3Predictor",
+    "SAM3SemanticPredictor",
+    "SAM3VideoPredictor",
+    "SAM3VideoSemanticPredictor",
+)  # tuple or list of exportable items

{dgenerate_ultralytics_headless-8.3.236 → dgenerate_ultralytics_headless-8.3.237}/ultralytics/models/sam/build.py

@@ -21,6 +21,21 @@ from .modules.tiny_encoder import TinyViT
 from .modules.transformer import TwoWayTransformer
 
 
+def _load_checkpoint(model, checkpoint):
+    """Load checkpoint into model from file path."""
+    if checkpoint is None:
+        return model
+
+    checkpoint = attempt_download_asset(checkpoint)
+    with open(checkpoint, "rb") as f:
+        state_dict = torch_load(f)
+    # Handle nested "model" key
+    if "model" in state_dict and isinstance(state_dict["model"], dict):
+        state_dict = state_dict["model"]
+    model.load_state_dict(state_dict)
+    return model
+
+
 def build_sam_vit_h(checkpoint=None):
     """Build and return a Segment Anything Model (SAM) h-size model with specified encoder parameters."""
     return _build_sam(
@@ -205,10 +220,7 @@ def _build_sam(
         pixel_std=[58.395, 57.12, 57.375],
     )
     if checkpoint is not None:
-        checkpoint = attempt_download_asset(checkpoint)
-        with open(checkpoint, "rb") as f:
-            state_dict = torch_load(f)
-        sam.load_state_dict(state_dict)
+        sam = _load_checkpoint(sam, checkpoint)
     sam.eval()
     return sam
 
@@ -299,10 +311,7 @@ def _build_sam2(
     )
 
     if checkpoint is not None:
-        checkpoint = attempt_download_asset(checkpoint)
-        with open(checkpoint, "rb") as f:
-            state_dict = torch_load(f)["model"]
-        sam2.load_state_dict(state_dict)
+        sam2 = _load_checkpoint(sam2, checkpoint)
     sam2.eval()
     return sam2
 

dgenerate_ultralytics_headless-8.3.237/ultralytics/models/sam/build_sam3.py

@@ -0,0 +1,374 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+
+import torch.nn as nn
+
+from ultralytics.nn.modules.transformer import MLP
+from ultralytics.utils.patches import torch_load
+
+from .modules.blocks import PositionEmbeddingSine, RoPEAttention
+from .modules.encoders import MemoryEncoder
+from .modules.memory_attention import MemoryAttention, MemoryAttentionLayer
+from .modules.sam import SAM3Model
+from .sam3.decoder import TransformerDecoder, TransformerDecoderLayer
+from .sam3.encoder import TransformerEncoderFusion, TransformerEncoderLayer
+from .sam3.geometry_encoders import SequenceGeometryEncoder
+from .sam3.maskformer_segmentation import PixelDecoder, UniversalSegmentationHead
+from .sam3.model_misc import DotProductScoring, TransformerWrapper
+from .sam3.necks import Sam3DualViTDetNeck
+from .sam3.sam3_image import SAM3SemanticModel
+from .sam3.text_encoder_ve import VETextEncoder
+from .sam3.tokenizer_ve import SimpleTokenizer
+from .sam3.vitdet import ViT
+from .sam3.vl_combiner import SAM3VLBackbone
+
+
+def _create_vision_backbone(compile_mode=None, enable_inst_interactivity=True) -> Sam3DualViTDetNeck:
+    """Create SAM3 visual backbone with ViT and neck."""
+    # Position encoding
+    position_encoding = PositionEmbeddingSine(
+        num_pos_feats=256,
+        normalize=True,
+        scale=None,
+        temperature=10000,
+    )
+
+    # ViT backbone
+    vit_backbone = ViT(
+        img_size=1008,
+        pretrain_img_size=336,
+        patch_size=14,
+        embed_dim=1024,
+        depth=32,
+        num_heads=16,
+        mlp_ratio=4.625,
+        norm_layer="LayerNorm",
+        drop_path_rate=0.1,
+        qkv_bias=True,
+        use_abs_pos=True,
+        tile_abs_pos=True,
+        global_att_blocks=(7, 15, 23, 31),
+        rel_pos_blocks=(),
+        use_rope=True,
+        use_interp_rope=True,
+        window_size=24,
+        pretrain_use_cls_token=True,
+        retain_cls_token=False,
+        ln_pre=True,
+        ln_post=False,
+        return_interm_layers=False,
+        bias_patch_embed=False,
+        compile_mode=compile_mode,
+    )
+    return Sam3DualViTDetNeck(
+        position_encoding=position_encoding,
+        d_model=256,
+        scale_factors=[4.0, 2.0, 1.0, 0.5],
+        trunk=vit_backbone,
+        add_sam2_neck=enable_inst_interactivity,
+    )
+
+
+def _create_sam3_transformer() -> TransformerWrapper:
+    """Create SAM3 detector encoder and decoder."""
+    encoder: TransformerEncoderFusion = TransformerEncoderFusion(
+        layer=TransformerEncoderLayer(
+            d_model=256,
+            dim_feedforward=2048,
+            dropout=0.1,
+            pos_enc_at_attn=True,
+            pos_enc_at_cross_attn_keys=False,
+            pos_enc_at_cross_attn_queries=False,
+            pre_norm=True,
+            self_attention=nn.MultiheadAttention(
+                num_heads=8,
+                dropout=0.1,
+                embed_dim=256,
+                batch_first=True,
+            ),
+            cross_attention=nn.MultiheadAttention(
+                num_heads=8,
+                dropout=0.1,
+                embed_dim=256,
+                batch_first=True,
+            ),
+        ),
+        num_layers=6,
+        d_model=256,
+        num_feature_levels=1,
+        frozen=False,
+        use_act_checkpoint=True,
+        add_pooled_text_to_img_feat=False,
+        pool_text_with_mask=True,
+    )
+    decoder: TransformerDecoder = TransformerDecoder(
+        layer=TransformerDecoderLayer(
+            d_model=256,
+            dim_feedforward=2048,
+            dropout=0.1,
+            cross_attention=nn.MultiheadAttention(
+                num_heads=8,
+                dropout=0.1,
+                embed_dim=256,
+            ),
+            n_heads=8,
+            use_text_cross_attention=True,
+        ),
+        num_layers=6,
+        num_queries=200,
+        return_intermediate=True,
+        box_refine=True,
+        num_o2m_queries=0,
+        dac=True,
+        boxRPB="log",
+        d_model=256,
+        frozen=False,
+        interaction_layer=None,
+        dac_use_selfatt_ln=True,
+        use_act_checkpoint=True,
+        presence_token=True,
+    )
+
+    return TransformerWrapper(encoder=encoder, decoder=decoder, d_model=256)
+
+
+def build_sam3_image_model(
+    checkpoint_path: str, bpe_path: str, enable_segmentation: bool = True, compile: bool = False
+):
+    """Build SAM3 image model.
+
+    Args:
+        checkpoint_path: Optional path to model checkpoint
+        bpe_path: Path to the BPE tokenizer vocabulary
+        enable_segmentation: Whether to enable segmentation head
+        compile: To enable compilation, set to "default"
+
+    Returns:
+        A SAM3 image model
+    """
+    # Create visual components
+    compile_mode = "default" if compile else None
+    vision_encoder = _create_vision_backbone(compile_mode=compile_mode, enable_inst_interactivity=True)
+
+    # Create text components
+    text_encoder = VETextEncoder(
+        tokenizer=SimpleTokenizer(bpe_path=bpe_path),
+        d_model=256,
+        width=1024,
+        heads=16,
+        layers=24,
+    )
+
+    # Create visual-language backbone
+    backbone = SAM3VLBackbone(visual=vision_encoder, text=text_encoder, scalp=1)
+
+    # Create transformer components
+    transformer = _create_sam3_transformer()
+
+    # Create dot product scoring
+    dot_prod_scoring = DotProductScoring(
+        d_model=256,
+        d_proj=256,
+        prompt_mlp=MLP(
+            input_dim=256,
+            hidden_dim=2048,
+            output_dim=256,
+            num_layers=2,
+            residual=True,
+            out_norm=nn.LayerNorm(256),
+        ),
+    )
+
+    # Create segmentation head if enabled
+    segmentation_head = (
+        UniversalSegmentationHead(
+            hidden_dim=256,
+            upsampling_stages=3,
+            aux_masks=False,
+            presence_head=False,
+            dot_product_scorer=None,
+            act_ckpt=True,
+            cross_attend_prompt=nn.MultiheadAttention(
+                num_heads=8,
+                dropout=0,
+                embed_dim=256,
+            ),
+            pixel_decoder=PixelDecoder(
+                num_upsampling_stages=3,
+                interpolation_mode="nearest",
+                hidden_dim=256,
+                compile_mode=compile_mode,
+            ),
+        )
+        if enable_segmentation
+        else None
+    )
+
+    # Create geometry encoder
+    input_geometry_encoder = SequenceGeometryEncoder(
+        pos_enc=PositionEmbeddingSine(
+            num_pos_feats=256,
+            normalize=True,
+            scale=None,
+            temperature=10000,
+        ),
+        encode_boxes_as_points=False,
+        boxes_direct_project=True,
+        boxes_pool=True,
+        boxes_pos_enc=True,
+        d_model=256,
+        num_layers=3,
+        layer=TransformerEncoderLayer(
+            d_model=256,
+            dim_feedforward=2048,
+            dropout=0.1,
+            pos_enc_at_attn=False,
+            pre_norm=True,
+            pos_enc_at_cross_attn_queries=False,
+            pos_enc_at_cross_attn_keys=True,
+        ),
+        use_act_ckpt=True,
+        add_cls=True,
+        add_post_encode_proj=True,
+    )
+
+    # Create the SAM3SemanticModel model
+    model = SAM3SemanticModel(
+        backbone=backbone,
+        transformer=transformer,
+        input_geometry_encoder=input_geometry_encoder,
+        segmentation_head=segmentation_head,
+        num_feature_levels=1,
+        o2m_mask_predict=True,
+        dot_prod_scoring=dot_prod_scoring,
+        use_instance_query=False,
+        multimask_output=True,
+    )
+
+    # Load checkpoint
+    model = _load_checkpoint(model, checkpoint_path)
+    model.eval()
+    return model
+
+
+def build_interactive_sam3(checkpoint_path: str, compile=None, with_backbone=True) -> SAM3Model:
+    """Build the SAM3 Tracker module for video tracking.
+
+    Returns:
+        Sam3TrackerPredictor: Wrapped SAM3 Tracker module
+    """
+    # Create model components
+    memory_encoder = MemoryEncoder(out_dim=64, interpol_size=[1152, 1152])
+    memory_attention = MemoryAttention(
+        batch_first=True,
+        d_model=256,
+        pos_enc_at_input=True,
+        layer=MemoryAttentionLayer(
+            dim_feedforward=2048,
+            dropout=0.1,
+            pos_enc_at_attn=False,
+            pos_enc_at_cross_attn_keys=True,
+            pos_enc_at_cross_attn_queries=False,
+            self_attn=RoPEAttention(
+                embedding_dim=256,
+                num_heads=1,
+                downsample_rate=1,
+                rope_theta=10000.0,
+                feat_sizes=[72, 72],
+            ),
+            d_model=256,
+            cross_attn=RoPEAttention(
+                embedding_dim=256,
+                num_heads=1,
+                downsample_rate=1,
+                kv_in_dim=64,
+                rope_theta=10000.0,
+                feat_sizes=[72, 72],
+                rope_k_repeat=True,
+            ),
+        ),
+        num_layers=4,
+    )
+
+    backbone = (
+        SAM3VLBackbone(scalp=1, visual=_create_vision_backbone(compile_mode=compile), text=None)
+        if with_backbone
+        else None
+    )
+    model = SAM3Model(
+        image_size=1008,
+        image_encoder=backbone,
+        memory_attention=memory_attention,
+        memory_encoder=memory_encoder,
+        backbone_stride=14,
+        num_maskmem=7,
+        sigmoid_scale_for_mem_enc=20.0,
+        sigmoid_bias_for_mem_enc=-10.0,
+        use_mask_input_as_output_without_sam=True,
+        directly_add_no_mem_embed=True,
+        use_high_res_features_in_sam=True,
+        multimask_output_in_sam=True,
+        iou_prediction_use_sigmoid=True,
+        use_obj_ptrs_in_encoder=True,
+        add_tpos_enc_to_obj_ptrs=True,
+        only_obj_ptrs_in_the_past_for_eval=True,
+        pred_obj_scores=True,
+        pred_obj_scores_mlp=True,
+        fixed_no_obj_ptr=True,
+        multimask_output_for_tracking=True,
+        use_multimask_token_for_obj_ptr=True,
+        multimask_min_pt_num=0,
+        multimask_max_pt_num=1,
+        use_mlp_for_obj_ptr_proj=True,
+        compile_image_encoder=False,
+        no_obj_embed_spatial=True,
+        proj_tpos_enc_in_obj_ptrs=True,
+        use_signed_tpos_enc_to_obj_ptrs=True,
+        sam_mask_decoder_extra_args=dict(
+            dynamic_multimask_via_stability=True,
+            dynamic_multimask_stability_delta=0.05,
+            dynamic_multimask_stability_thresh=0.98,
+        ),
+    )
+
+    # Load checkpoint if provided
+    model = _load_checkpoint(model, checkpoint_path, interactive=True)
+
+    # Setup device and mode
+    model.eval()
+    return model
+
+
+def _load_checkpoint(model, checkpoint, interactive=False):
+    """Load SAM3 model checkpoint from file."""
+    with open(checkpoint, "rb") as f:
+        ckpt = torch_load(f)
+    if "model" in ckpt and isinstance(ckpt["model"], dict):
+        ckpt = ckpt["model"]
+    sam3_image_ckpt = {k.replace("detector.", ""): v for k, v in ckpt.items() if "detector" in k}
+    if interactive:
+        sam3_image_ckpt.update(
+            {
+                k.replace("backbone.vision_backbone", "image_encoder.vision_backbone"): v
+                for k, v in sam3_image_ckpt.items()
+                if "backbone.vision_backbone" in k
+            }
+        )
+        sam3_image_ckpt.update(
+            {
+                k.replace("tracker.transformer.encoder", "memory_attention"): v
+                for k, v in ckpt.items()
+                if "tracker.transformer" in k
+            }
+        )
+        sam3_image_ckpt.update(
+            {
+                k.replace("tracker.maskmem_backbone", "memory_encoder"): v
+                for k, v in ckpt.items()
+                if "tracker.maskmem_backbone" in k
+            }
+        )
+        sam3_image_ckpt.update({k.replace("tracker.", ""): v for k, v in ckpt.items() if "tracker." in k})
+    model.load_state_dict(sam3_image_ckpt, strict=False)
+    return model

{dgenerate_ultralytics_headless-8.3.236 → dgenerate_ultralytics_headless-8.3.237}/ultralytics/models/sam/model.py

@@ -21,7 +21,7 @@ from pathlib import Path
 from ultralytics.engine.model import Model
 from ultralytics.utils.torch_utils import model_info
 
-from .predict import Predictor, SAM2Predictor
+from .predict import Predictor, SAM2Predictor, SAM3Predictor
 
 
 class SAM(Model):
@@ -59,6 +59,7 @@ class SAM(Model):
         if model and Path(model).suffix not in {".pt", ".pth"}:
             raise NotImplementedError("SAM prediction requires pre-trained *.pt or *.pth model.")
         self.is_sam2 = "sam2" in Path(model).stem
+        self.is_sam3 = "sam3" in Path(model).stem
         super().__init__(model=model, task="segment")
 
     def _load(self, weights: str, task=None):
@@ -72,9 +73,14 @@ class SAM(Model):
             >>> sam = SAM("sam_b.pt")
             >>> sam._load("path/to/custom_weights.pt")
         """
-        from .build import build_sam  # slow import
+        if self.is_sam3:
+            from .build_sam3 import build_interactive_sam3
 
-        self.model = build_sam(weights)
+            self.model = build_interactive_sam3(weights)
+        else:
+            from .build import build_sam  # slow import
+
+            self.model = build_sam(weights)
 
     def predict(self, source, stream: bool = False, bboxes=None, points=None, labels=None, **kwargs):
         """Perform segmentation prediction on the given image or video source.
@@ -158,4 +164,6 @@ class SAM(Model):
             >>> print(task_map)
             {'segment': {'predictor': <class 'ultralytics.models.sam.predict.Predictor'>}}
         """
-        return {"segment": {"predictor": SAM2Predictor if self.is_sam2 else Predictor}}
+        return {
+            "segment": {"predictor": SAM2Predictor if self.is_sam2 else SAM3Predictor if self.is_sam3 else Predictor}
+        }