garmentiq.segmentation.model_definition.sam

251def load_sam_config(model_type: str = "sam-vit-b", model_dir: Optional[str] = None):
252    """
253    Reads and loads the configuration for a specified Segment Anything Model variant.
254
255    This function facilitates strictly offline initialization by dynamically locating the
256    `config.json` file associated with the chosen SAM variant within the local package
257    directory. It reads the JSON file securely and converts it into the Hugging Face
258    configuration object matching the variant's family, entirely bypassing external
259    network requests.
260
261    Configurations for SAM 1 (`sam-vit-*`), SAM 2.1 (`sam2.1-hiera-*`) and SAM 3 (`sam3`)
262    all ship with the package, so no Hugging Face access is needed. SAM 3 weights are still
263    gated, so download the checkpoint once and point `model_path` at it when loading the
264    model; `model_dir` here only overrides the bundled configuration.
265
266    Released SAM 3 configurations describe the video model and nest the image model under
267    `detector_config`. GarmentIQ segments still images, so that nested section is unwrapped
268    automatically.
269
270    Args:
271        model_type (str, optional): The identifier for the desired SAM variant. Must be one
272                                    of `["sam-vit-b", "sam-vit-l", "sam-vit-h",
273                                    "sam2.1-hiera-tiny", "sam2.1-hiera-small",
274                                    "sam2.1-hiera-base-plus", "sam2.1-hiera-large", "sam3"]`.
275                                    Default is `"sam-vit-b"`.
276        model_dir (str, optional): Path to a local directory containing the variant's
277                                   `config.json`, used to override the bundled configuration.
278                                   Default is None.
279
280    Raises:
281        ValueError: If the provided `model_type` is not within the supported variants list.
282        FileNotFoundError: If the corresponding offline `config.json` file cannot be located.
283
284    Returns:
285        Union[SamConfig, Sam2Config, Sam3Config]: The loaded configuration object ready to be
286                                                  passed into the matching SAM model class.
287    """
288    family = sam_family(model_type)
289    config_path = os.path.join(_variant_dir(model_type, model_dir), "config.json")
290
291    if not os.path.exists(config_path):
292        raise FileNotFoundError(f"Offline config missing at {config_path}.")
293
294    with open(config_path, "r") as f:
295        config_dict = json.load(f)
296
297    config_dict = _unwrap_detector_config(config_dict)
298
299    return _resolve(SAM_FAMILIES[family]["config_class"]).from_dict(config_dict)

319def load_sam_processor(
320    model_type: str = "sam-vit-b",
321    use_fast: Optional[bool] = None,
322    model_dir: Optional[str] = None,
323    backend: Optional[str] = None,
324):
325    """
326    Loads the offline processor for a specified Segment Anything Model variant.
327
328    This function instantiates the processor matching the variant's family by reading
329    bundled preprocessor (and, for SAM 3, tokenizer) configuration files from the local
330    variant directory. Loading from a local path ensures the package remains completely
331    air-gapped. The image resizing backend can be selected to balance speed against
332    strict backward compatibility.
333
334    Args:
335        model_type (str, optional): The identifier for the desired SAM variant.
336                                    Default is `"sam-vit-b"`.
337        use_fast (bool, optional): Deprecated alias for `backend`, kept for backwards
338                                   compatibility. `True` maps to `"torchvision"` and `False`
339                                   maps to `"pil"`. Only applies to SAM 1. Default is None.
340        model_dir (str, optional): Path to a local directory containing the variant's
341                                   processor files, used to override the bundled ones.
342                                   Default is None.
343        backend (str, optional): Image resizing backend, either `"pil"` or `"torchvision"`.
344                                 `"torchvision"` is faster, while `"pil"` matches the
345                                 historical GarmentIQ default. Only applies to SAM 1.
346                                 Default is None, which resolves to `"pil"`.
347
348    Raises:
349        ValueError: If the provided `model_type` is not within the supported variants list,
350                    or if `backend` is neither `"pil"` nor `"torchvision"`.
351        FileNotFoundError: If the corresponding offline processor configuration cannot be
352                           located.
353
354    Returns:
355        Union[SamProcessor, Sam2Processor, Sam3Processor]: The instantiated processor ready
356                                                           for image and prompt transformations.
357    """
358    family = sam_family(model_type)
359    variant_dir = _variant_dir(model_type, model_dir)
360
361    if backend is None and use_fast is not None:
362        backend = "torchvision" if use_fast else "pil"
363    if backend is None:
364        # Historical GarmentIQ default. Selected explicitly so results stay identical
365        # across transformers releases, which have shifted their own default over time.
366        backend = "pil"
367    if backend not in ("pil", "torchvision"):
368        raise ValueError(
369            f"Invalid backend {backend!r}. Choose either 'pil' or 'torchvision'."
370        )
371
372    if family == "sam1":
373        processor_path = os.path.join(variant_dir, "preprocessor_config.json")
374        if not os.path.exists(processor_path):
375            raise FileNotFoundError(
376                f"Offline processor config missing at {processor_path}."
377            )
378        # transformers 5.4 replaced `use_fast` with `backend`; passing the deprecated
379        # flag to a newer release emits a warning, so select the kwarg by version.
380        backend_kwargs = (
381            {"backend": backend}
382            if _supports_backend_kwarg()
383            else {"use_fast": backend == "torchvision"}
384        )
385        # Loading from a local directory disables Hugging Face network calls
386        return _resolve("SamProcessor").from_pretrained(
387            processor_path, **backend_kwargs
388        )
389
390    if family == "sam2":
391        processor_path = os.path.join(variant_dir, "preprocessor_config.json")
392        if not os.path.exists(processor_path):
393            raise FileNotFoundError(
394                f"Offline processor config missing at {processor_path}."
395            )
396        with open(processor_path, "r") as f:
397            image_processor_dict = json.load(f)
398        # The bundled config names the deprecated `*Fast` class; build the image
399        # processor explicitly so loading stays offline and warning-free.
400        image_processor_dict.pop("image_processor_type", None)
401        image_processor = _resolve("Sam2ImageProcessor")(**image_processor_dict)
402        return _resolve("Sam2Processor")(image_processor=image_processor)
403
404    # SAM 3 bundles a tokenizer alongside the image processor, so the whole
405    # directory is handed to the processor loader.
406    return _resolve("Sam3Processor").from_pretrained(variant_dir)

467def load_sam3_tracker(
468    model_path: str,
469    model_dir: Optional[str] = None,
470    device: Union[str, torch.device] = "cpu",
471):
472    """
473    Loads SAM 3's point-promptable tracker from a released SAM 3 checkpoint.
474
475    SAM 3 ships as a single checkpoint holding both the detector and the tracker, and the
476    two share one vision encoder that is stored only once under the detector's prefix. This
477    function reassembles the tracker by pairing `tracker_model.*` weights with the shared
478    `detector_model.vision_encoder.*` backbone, then verifies every expected tensor was
479    found so a mismatched checkpoint fails loudly instead of yielding a random model.
480
481    Use this when you want to prompt SAM 3 with points on a still image. For text or box
482    prompts use the detector via `Sam3Model` and `load_sam_config("sam3")` instead.
483
484    Args:
485        model_path (str): Local path to the SAM 3 `model.safetensors` checkpoint.
486        model_dir (str, optional): Directory holding `tracker_config.json`, used to
487                                   override the bundled configuration. Default is None.
488        device (Union[str, torch.device], optional): The device to load onto, e.g. `"cpu"`,
489                                                     `"cuda"`, or `"mps"`. Default is `"cpu"`.
490
491    Raises:
492        FileNotFoundError: If `model_path` does not exist.
493        ValueError: If the checkpoint does not contain the tensors the tracker expects, or
494                    if the requested `device` is unavailable.
495        ImportError: If the installed transformers release does not provide SAM 3.
496
497    Returns:
498        Sam3TrackerModel: The assembled tracker in evaluation mode, placed on `device`.
499    """
500    from safetensors.torch import load_file
501
502    from garmentiq.utils.device import resolve_device
503
504    if not os.path.exists(model_path):
505        raise FileNotFoundError(f"SAM 3 checkpoint not found at {model_path}.")
506
507    device = resolve_device(device)
508
509    model = _resolve("Sam3TrackerModel")(load_sam3_tracker_config(model_dir))
510    expected = model.state_dict()
511
512    if model_path.endswith(".safetensors"):
513        state_dict = load_file(model_path, device="cpu")
514    else:
515        state_dict = torch.load(model_path, map_location="cpu", weights_only=True)
516
517    shared_prefix = SAM3_DETECTOR_PREFIX + "vision_encoder."
518    neck_prefix = shared_prefix + "neck."
519    assembled = {}
520    for key, value in state_dict.items():
521        if key.startswith(neck_prefix):
522            # The detector has its own neck with the same layer names as the tracker's.
523            # Skip it: the tracker's neck is stored separately under `tracker_neck.`,
524            # and using the detector's would silently produce a garbage mask.
525            continue
526        if key.startswith(shared_prefix):
527            # The vision encoder backbone is shared, and stored once under the detector.
528            assembled[key[len(SAM3_DETECTOR_PREFIX) :]] = value
529        elif key.startswith(SAM3_TRACKER_NECK_PREFIX):
530            assembled[
531                "vision_encoder.neck." + key[len(SAM3_TRACKER_NECK_PREFIX) :]
532            ] = value
533        elif key.startswith(SAM3_TRACKER_PREFIX):
534            stripped = key[len(SAM3_TRACKER_PREFIX) :]
535            if stripped in expected:
536                assembled[stripped] = value
537
538    missing = set(expected) - set(assembled)
539    if missing:
540        raise ValueError(
541            f"SAM 3 tracker weights are incomplete: {len(missing)} of {len(expected)} "
542            f"tensors were not found in {model_path!r} "
543            f"(first missing: {sorted(missing)[:3]}). Check that the checkpoint is a "
544            f"full SAM 3 release."
545        )
546
547    model.load_state_dict(assembled, strict=True)
548    model = model.to(device)
549    model.eval()
550    return model

409def load_sam3_tracker_config(model_dir: Optional[str] = None):
410    """
411    Reads and loads the bundled configuration for SAM 3's point-promptable tracker.
412
413    SAM 3's image model is an open-vocabulary detector prompted by text or boxes, and does
414    not accept points. The release also contains a tracker, whose prompt encoder is the
415    SAM-style one that understands points, boxes and masks. This function loads that
416    tracker's configuration so still images can be segmented from a click.
417
418    Args:
419        model_dir (str, optional): Path to a local directory containing
420                                   `tracker_config.json`, used to override the bundled
421                                   configuration. Default is None.
422
423    Raises:
424        FileNotFoundError: If the offline `tracker_config.json` cannot be located.
425        ImportError: If the installed transformers release does not provide SAM 3.
426
427    Returns:
428        Sam3TrackerConfig: The configuration for `Sam3TrackerModel`.
429    """
430    variant_dir = _variant_dir("sam3", model_dir)
431    config_path = os.path.join(variant_dir, "tracker_config.json")
432
433    if not os.path.exists(config_path):
434        # A directly downloaded SAM 3 release nests it inside the video config.
435        fallback = os.path.join(variant_dir, "config.json")
436        if os.path.exists(fallback):
437            with open(fallback, "r") as f:
438                nested = json.load(f).get("tracker_config")
439            if isinstance(nested, dict):
440                return _resolve("Sam3TrackerConfig").from_dict(nested)
441        raise FileNotFoundError(f"Offline tracker config missing at {config_path}.")
442
443    with open(config_path, "r") as f:
444        config_dict = json.load(f)
445
446    return _resolve("Sam3TrackerConfig").from_dict(config_dict)

449def load_sam3_tracker_processor():
450    """
451    Builds the processor for SAM 3's point-promptable tracker.
452
453    The tracker takes only geometric prompts, so unlike `load_sam_processor("sam3")` no
454    tokenizer is involved. The image processor is constructed from its defaults, which
455    keeps loading fully offline.
456
457    Raises:
458        ImportError: If the installed transformers release does not provide SAM 3.
459
460    Returns:
461        Sam3TrackerProcessor: The processor for `Sam3TrackerModel`.
462    """
463    image_processor = _resolve("Sam3ImageProcessor")()
464    return _resolve("Sam3TrackerProcessor")(image_processor=image_processor)

142def sam_family(model_type: str) -> str:
143    """
144    Returns the Segment Anything family a variant belongs to.
145
146    Args:
147        model_type (str): A SAM variant identifier, e.g. `"sam-vit-b"`,
148                          `"sam2.1-hiera-tiny"`, or `"sam3"`.
149
150    Raises:
151        ValueError: If `model_type` is not a recognised SAM variant.
152
153    Returns:
154        str: One of `"sam1"`, `"sam2"`, or `"sam3"`.
155    """
156    for family, spec in SAM_FAMILIES.items():
157        if model_type in spec["models"]:
158            return family
159    raise ValueError(
160        f"Invalid model_type '{model_type}'. Choose from: {ALL_SAM_MODELS}"
161    )

164def sam_capabilities(model_type: str) -> dict:
165    """
166    Reports which prompt types a SAM variant accepts natively.
167
168    Args:
169        model_type (str): A SAM variant identifier, e.g. `"sam-vit-b"` or `"sam3"`.
170
171    Raises:
172        ValueError: If `model_type` is not a recognised SAM variant.
173
174    Returns:
175        dict: Mapping of `"points"`, `"labels"`, `"boxes"`, and `"text"` to booleans.
176              SAM 1 and SAM 2 report `text=False` because they contain no text encoder,
177              and SAM 3 reports `points=False` because it accepts only text and box prompts.
178    """
179    return dict(SAM_CAPABILITIES[sam_family(model_type)])

182def sam_model_class(model_type: str):
183    """
184    Returns the Hugging Face model class matching a SAM variant.
185
186    Args:
187        model_type (str): A SAM variant identifier, e.g. `"sam2.1-hiera-tiny"`.
188
189    Raises:
190        ValueError: If `model_type` is not a recognised SAM variant.
191
192    Returns:
193        type: `SamModel`, `Sam2Model`, or `Sam3Model`.
194    """
195    return _resolve(SAM_FAMILIES[sam_family(model_type)]["model_class"])