garmentiq.segmentation.extract

  1"""Running segmentation to obtain a garment mask.
  2
  3Dispatches on the model family, so BiRefNet and every SAM generation share one entry
  4point. SAM prompts arrive in a unified `prompt` dictionary accepting `points`,
  5`labels`, `boxes`, and `text`; nesting differences between SAM versions are normalised
  6here, and a text prompt for SAM 1 or SAM 2 is routed through a grounding model first.
  7"""
  8from PIL import Image
  9import torch
 10from torchvision import transforms
 11import numpy as np
 12from typing import Optional, Union
 13from garmentiq.utils.device import resolve_device, empty_cache, inputs_to_device
 14from garmentiq.segmentation.model_definition.sam import (
 15    SAM_CAPABILITIES,
 16    SAM_PROMPT_DEPTHS,
 17)
 18
 19PROMPT_KEYS = ("points", "labels", "boxes", "text")
 20
 21# Legacy keyword arguments accepted for backwards compatibility, mapped onto the
 22# unified prompt dictionary.
 23LEGACY_PROMPT_KWARGS = {
 24    "input_points": "points",
 25    "input_labels": "labels",
 26    "input_boxes": "boxes",
 27}
 28
 29
 30def _nesting_depth(value):
 31    """Returns how many list/tuple levels wrap the first scalar in `value`."""
 32    depth = 0
 33    current = value
 34    while isinstance(current, (list, tuple)):
 35        depth += 1
 36        if len(current) == 0:
 37            break
 38        current = current[0]
 39    return depth
 40
 41
 42def _to_depth(value, target: int, name: str):
 43    """
 44    Reshapes a nested prompt list to the nesting depth a processor expects.
 45
 46    SAM 1 and SAM 2 disagree on prompt nesting (SAM 2 requires exactly one extra
 47    level and raises otherwise), so prompts are normalised here to let callers use
 48    a single consistent format across every SAM family.
 49    """
 50    if isinstance(value, torch.Tensor) or isinstance(value, np.ndarray):
 51        return value
 52
 53    depth = _nesting_depth(value)
 54    while depth < target:
 55        value = [value]
 56        depth += 1
 57    while depth > target:
 58        if len(value) != 1:
 59            raise ValueError(
 60                f"Prompt '{name}' is nested {depth} levels deep but this SAM family "
 61                f"expects {target}, and it cannot be unwrapped because the outermost "
 62                f"level holds {len(value)} items. Provide '{name}' with {target} levels."
 63            )
 64        value = value[0]
 65        depth -= 1
 66    return value
 67
 68
 69def _detect_sam_family(model):
 70    """
 71    Identifies which SAM family a loaded model belongs to, if any.
 72
 73    Detection walks the class hierarchy by name rather than importing the transformers
 74    classes, so a model from any SAM generation is recognised without requiring every
 75    generation to exist in the installed transformers release.
 76    """
 77    names = {cls.__name__ for cls in type(model).__mro__}
 78    # Checked before Sam3Model: the tracker is a separate class that carries the SAM-style
 79    # prompt encoder, so it accepts points where the SAM 3 detector does not.
 80    if {"Sam3TrackerModel", "Sam3TrackerPreTrainedModel"} & names:
 81        return "sam3-tracker"
 82    if {"Sam3Model", "Sam3PreTrainedModel"} & names:
 83        return "sam3"
 84    if {"Sam2Model", "Sam2PreTrainedModel"} & names:
 85        return "sam2"
 86    if {"SamModel", "SamPreTrainedModel"} & names:
 87        return "sam1"
 88    return None
 89
 90
 91def _build_prompt(prompt: Optional[dict], kwargs: dict):
 92    """Merges the `prompt` dict with legacy `input_*` keyword arguments."""
 93    merged = {}
 94
 95    if prompt is not None:
 96        if not isinstance(prompt, dict):
 97            raise ValueError(
 98                f"`prompt` must be a dictionary with any of {list(PROMPT_KEYS)}, "
 99                f"got {type(prompt).__name__}."
100            )
101        unknown = set(prompt) - set(PROMPT_KEYS)
102        if unknown:
103            raise ValueError(
104                f"Unknown prompt key(s) {sorted(unknown)}. "
105                f"Valid keys are {list(PROMPT_KEYS)}."
106            )
107        merged.update({k: v for k, v in prompt.items() if v is not None})
108
109    for legacy_key, new_key in LEGACY_PROMPT_KWARGS.items():
110        if kwargs.get(legacy_key) is not None:
111            merged.setdefault(new_key, kwargs[legacy_key])
112
113    return merged
114
115
116def extract(
117    model: torch.nn.Module,
118    image_path: str,
119    processor=None,
120    prompt: Optional[dict] = None,
121    device: Union[str, torch.device] = "cpu",
122    grounding_model=None,
123    grounding_processor=None,
124    grounding_args: Optional[dict] = None,
125    **kwargs,
126):
127    """
128    Intelligently extracts an image segmentation mask from a given image using either a standard
129    PyTorch model or a Processor-based foundation model.
130
131    This function takes an image and processes it based on the model strategy. If a processor is supplied,
132    it delegates preprocessing (e.g., resizing, scaling, prompt-handling) to the processor. Otherwise,
133    it applies standard manual transformations based on provided kwargs. It then feeds the input into
134    the model to generate a segmentation mask. The original image and the mask are returned as numpy arrays.
135
136    For Segment Anything models the prompt is supplied through the unified `prompt` dictionary,
137    which may carry geometric prompts (`"points"`, `"labels"`, `"boxes"`) and/or a natural-language
138    prompt (`"text"`). At least one prompt is required, because an unprompted SAM silently returns
139    a meaningless mask rather than raising. Prompt nesting is normalised per family, so the same
140    `prompt` works across SAM 1, SAM 2, and SAM 3.
141
142    Text prompts are handled differently per family, because only SAM 3 has a text encoder:
143        - SAM 3 consumes the text directly.
144        - SAM 1 and SAM 2 have no text encoder, so `grounding_model` and `grounding_processor`
145          must be supplied. The phrase is grounded into boxes first, and those boxes are then
146          used as ordinary SAM prompts.
147
148    The model and its inputs are placed on `device`, so the same value should be passed here as was
149    used when loading the model. Any accelerator memory cached during inference is released before
150    returning.
151
152    Args:
153        model (torch.nn.Module): The pretrained PyTorch model to use for segmentation predictions.
154        image_path (str): The path to the image file on which to perform segmentation.
155        processor (Any, optional): The model-specific processor (e.g., from Hugging Face) used for preprocessing
156                                   inputs. If None, standard PyTorch manual transformations are applied.
157                                   Default is None.
158        prompt (dict, optional): The unified prompt for Segment Anything models. Recognised keys:
159                                 - `"points"` (list): Point coordinates, e.g. `[[[x, y]]]`.
160                                   Not supported by SAM 3.
161                                 - `"labels"` (list): Point labels, `1` for foreground and `0` for
162                                   background, e.g. `[[1]]`. Not supported by SAM 3.
163                                 - `"boxes"` (list): Boxes as `[[[x_min, y_min, x_max, y_max]]]`.
164                                 - `"text"` (str): Natural-language prompt, e.g. `"a shirt"`.
165                                 Default is None.
166        device (Union[str, torch.device], optional): The device to run inference on, e.g. `"cpu"`,
167                                                     `"cuda"`, `"cuda:0"`, or `"mps"`. Hardware
168                                                     acceleration is opt-in; pass it explicitly to
169                                                     use a GPU or Apple Silicon. Default is `"cpu"`.
170        grounding_model (GroundingDinoForObjectDetection, optional): Grounding model used to turn a
171                                                                     text prompt into boxes for SAM 1
172                                                                     and SAM 2. Default is None.
173        grounding_processor (GroundingDinoProcessor, optional): Processor matching `grounding_model`.
174                                                                Default is None.
175        grounding_args (dict, optional): Extra grounding options, accepting `"box_threshold"`,
176                                         `"text_threshold"`, and `"max_boxes"`. Default is None.
177        **kwargs: Additional arbitrary keyword arguments for model-specific configurations.
178                  For standard models (e.g., BiRefNet): `resize_dim`, `normalize_mean`, `normalize_std`.
179                  The legacy `input_points`, `input_labels`, and `input_boxes` arguments are still
180                  accepted and are folded into `prompt`.
181
182    Raises:
183        FileNotFoundError: If the image file at `image_path` does not exist.
184        ValueError: If the requested `device` is invalid or unavailable, if no prompt is supplied for
185                    a SAM model, if a prompt type is unsupported by the chosen SAM family, if a text
186                    prompt is given for SAM 1 or SAM 2 without a grounding model, or if the processor
187                    output format is unrecognized.
188
189    Returns:
190        tuple (numpy.ndarray, numpy.ndarray): The original image converted to a numpy array,
191                                              and the extracted segmentation mask as a numpy array.
192    """
193    device = resolve_device(device)
194    model = model.to(device)
195    image = Image.open(image_path).convert("RGB")
196
197    # Processor-Based Models (SAM 1 / SAM 2 / SAM 3)
198    if processor is not None:
199        family = _detect_sam_family(model)
200        prompt_dict = _build_prompt(prompt, kwargs)
201
202        if family is not None:
203            capabilities = SAM_CAPABILITIES[family]
204
205            if not prompt_dict:
206                options = []
207                if capabilities["points"]:
208                    options.append("prompt={'points': [[[x, y]]]}")
209                if capabilities["boxes"]:
210                    options.append("prompt={'boxes': [[[x0, y0, x1, y1]]]}")
211                if capabilities["text"]:
212                    options.append("prompt={'text': 'a shirt'}")
213                else:
214                    options.append(
215                        "prompt={'text': 'a shirt'} together with a grounding model"
216                    )
217                raise ValueError(
218                    f"A prompt is required for {family}. Pass one of: {', '.join(options)}. "
219                    f"Without a prompt SAM returns an arbitrary mask instead of failing."
220                )
221
222            # `text` is intentionally excluded here: it is always acceptable, either
223            # natively (SAM 3) or by grounding it into boxes first (SAM 1 / SAM 2).
224            unsupported = [
225                key
226                for key in prompt_dict
227                if key != "text" and key in capabilities and not capabilities[key]
228            ]
229            if unsupported:
230                detail = ""
231                if family == "sam3" and {"points", "labels"} & set(unsupported):
232                    # SAM 3's release does contain a point-promptable model, but it is
233                    # the tracker rather than the detector loaded here.
234                    detail = (
235                        " SAM 3's image model is an open-vocabulary detector, prompted by "
236                        "description rather than by clicking. For point prompts on a still "
237                        "image load SAM 3's tracker with "
238                        "garmentiq.segmentation.model_definition.sam.load_sam3_tracker(), "
239                        "or use SAM 1 / SAM 2."
240                    )
241                raise ValueError(
242                    f"Prompt key(s) {sorted(unsupported)} are not supported by {family}. "
243                    f"{family} supports: "
244                    f"{sorted(k for k, v in capabilities.items() if v)}.{detail}"
245                )
246
247            # Resolve a text prompt into boxes for the families without a text encoder.
248            text = prompt_dict.pop("text", None)
249            if text is not None and not capabilities["text"]:
250                if grounding_model is None or grounding_processor is None:
251                    raise ValueError(
252                        f"{family} has no text encoder, so a text prompt requires a grounding "
253                        f"model. Pass grounding_model and grounding_processor (see "
254                        f"garmentiq.grounding), or use SAM 3 which understands text natively."
255                    )
256                # Imported here so the grounding stack is only required when used.
257                from garmentiq.grounding import ground_text_to_boxes
258
259                grounded = ground_text_to_boxes(
260                    model=grounding_model,
261                    processor=grounding_processor,
262                    image=image,
263                    text=text,
264                    device=device,
265                    **(grounding_args or {}),
266                )
267                # Grounded boxes replace any caller-supplied boxes for this prompt.
268                prompt_dict["boxes"] = [grounded]
269                text = None
270
271            depths = SAM_PROMPT_DEPTHS[family]
272            processor_kwargs = {}
273            if "points" in prompt_dict:
274                processor_kwargs["input_points"] = _to_depth(
275                    prompt_dict["points"], depths["points"], "points"
276                )
277            if "labels" in prompt_dict:
278                processor_kwargs["input_labels"] = _to_depth(
279                    prompt_dict["labels"], depths["labels"], "labels"
280                )
281            if "boxes" in prompt_dict:
282                processor_kwargs["input_boxes"] = _to_depth(
283                    prompt_dict["boxes"], depths["boxes"], "boxes"
284                )
285            if text is not None:
286                processor_kwargs["text"] = text
287
288            inputs = inputs_to_device(
289                processor(images=image, return_tensors="pt", **processor_kwargs),
290                device,
291            )
292
293            with torch.no_grad():
294                outputs = model(**inputs)
295
296            if family == "sam3":
297                results = processor.post_process_instance_segmentation(
298                    outputs, target_sizes=[(image.height, image.width)]
299                )[0]
300                instance_masks = results["masks"]
301                if instance_masks.shape[0] == 0:
302                    raise ValueError(
303                        f"SAM 3 found no instance matching the prompt in {image_path!r}. "
304                        f"Try rephrasing the text prompt or lowering the threshold."
305                    )
306                # Keep the highest-scoring instance, mirroring the single-mask contract.
307                best = int(torch.argmax(results["scores"]))
308                best_mask = instance_masks[best].detach().cpu().numpy()
309            else:
310                post_kwargs = {}
311                if family == "sam1":
312                    post_kwargs["reshaped_input_sizes"] = inputs[
313                        "reshaped_input_sizes"
314                    ].cpu()
315                masks = processor.post_process_masks(
316                    outputs.pred_masks.cpu(),
317                    inputs["original_sizes"].cpu(),
318                    **post_kwargs,
319                )
320                if family == "sam3-tracker":
321                    # The tracker returns three candidate masks per prompt with an IoU
322                    # score for each, so pick the highest-scoring one rather than the
323                    # first, which is often the smallest sub-part.
324                    per_prompt_masks = masks[0]
325                    iou = outputs.iou_scores.cpu()
326                    selected = []
327                    for i in range(per_prompt_masks.shape[0]):
328                        scores = iou[0, i] if iou.ndim == 3 else iou[i]
329                        selected.append(per_prompt_masks[i, int(torch.argmax(scores))])
330                    per_prompt = torch.stack(selected, dim=0)
331                else:
332                    # masks[0] has shape (num_prompts, num_masks, H, W). Keep the first
333                    # mask of each prompt and union them, which is identical to the
334                    # previous single-prompt behaviour when only one prompt is given.
335                    per_prompt = masks[0][:, 0]
336                best_mask = (
337                    torch.any(per_prompt.bool(), dim=0).detach().cpu().numpy()
338                )
339
340            mask_np = (best_mask.astype(np.float32) * 255).astype(np.uint8)
341            del inputs, outputs
342
343        else:
344            # Unknown processor-based model: fall back to the generic SAM-like contract.
345            inputs = inputs_to_device(
346                processor(image, return_tensors="pt", **kwargs), device
347            )
348
349            with torch.no_grad():
350                outputs = model(**inputs)
351
352            if hasattr(outputs, "pred_masks"):
353                masks = processor.image_processor.post_process_masks(
354                    outputs.pred_masks.cpu(),
355                    inputs["original_sizes"].cpu(),
356                    inputs["reshaped_input_sizes"].cpu(),
357                )
358                best_mask = masks[0][0][0].numpy()
359                mask_np = (best_mask * 255).astype(np.uint8)
360            else:
361                raise ValueError("Unrecognized processor output format.")
362
363            del inputs, outputs, masks
364
365    # Standard Models (BiRefNet)
366    else:
367        # Extract BiRefNet-specific kwargs with safe defaults
368        resize_dim = kwargs.get("resize_dim", (1024, 1024))
369        normalize_mean = kwargs.get("normalize_mean", [0.485, 0.456, 0.406])
370        normalize_std = kwargs.get("normalize_std", [0.229, 0.224, 0.225])
371
372        transform = transforms.Compose(
373            [
374                transforms.Resize(resize_dim),
375                transforms.ToTensor(),
376                transforms.Normalize(normalize_mean, normalize_std),
377            ]
378        )
379
380        input_tensor = transform(image).unsqueeze(0).to(device)
381
382        with torch.no_grad():
383            preds = model(input_tensor)
384
385            # BiRefNet returns a tuple/list of tensors; we want the last one
386            if isinstance(preds, (list, tuple)):
387                preds = preds[-1]
388
389            preds = preds.sigmoid().cpu()
390
391        pred = preds[0].squeeze()
392        pred_pil = transforms.ToPILImage()(pred)
393        mask = pred_pil.resize(image.size)
394        mask_np = np.array(mask)
395
396        del input_tensor, preds
397
398    # Clean up and Return
399    image_np = np.array(image)
400    empty_cache(device)
401
402    return image_np, mask_np