garmentiq.matting.matte

  1"""Producing an alpha matte and compositing with it.
  2
  3Dispatches on the model family: ViTMatte is guided by a trimap, which can be supplied
  4directly or derived from a segmentation mask, while Matting Anything is guided by a SAM
  5prompt and needs no trimap at all.
  6"""
  7from PIL import Image
  8import numpy as np
  9import torch
 10import warnings
 11from typing import Optional, Union
 12
 13from garmentiq.utils.device import resolve_device, empty_cache, inputs_to_device
 14from garmentiq.matting.trimap import generate_trimap
 15
 16# Above roughly this many pixels the MPS backend has been observed to diverge from CPU
 17# for ViTMatte, silently producing a degenerate matte. Verified identical up to
 18# 2048x1536 (3.15 MP) and divergent at 2400x1800 (4.32 MP) on Apple Silicon.
 19_MPS_PIXEL_WARN_THRESHOLD = 3_500_000
 20
 21
 22def _detect_matting_family(model):
 23    """
 24    Identifies which matting approach a loaded model implements.
 25
 26    Detection walks the class hierarchy by name rather than importing the transformers
 27    classes, so a model is recognised without requiring every backend to be installed.
 28    """
 29    names = {cls.__name__ for cls in type(model).__mro__}
 30    if "MattingAnything" in names:
 31        return "mam"
 32    if {"VitMatteForImageMatting", "VitMattePreTrainedModel"} & names:
 33        return "vitmatte"
 34    return None
 35
 36
 37def matte(
 38    model,
 39    image_path: Union[str, np.ndarray, Image.Image],
 40    processor=None,
 41    mask: Optional[np.ndarray] = None,
 42    trimap: Optional[np.ndarray] = None,
 43    prompt: Optional[dict] = None,
 44    trimap_args: Optional[dict] = None,
 45    device: Union[str, torch.device] = "cpu",
 46):
 47    """
 48    Extracts a soft alpha matte for an image, refining a hard segmentation into soft edges.
 49
 50    Segmentation answers "which pixels belong to the garment" with a yes/no decision, which
 51    leaves stair-stepped borders and loses semi-transparent detail such as loose fabric,
 52    lace, and stray fibres. Matting instead predicts a continuous alpha value per pixel, so
 53    compositing onto a new background looks natural.
 54
 55    The two supported approaches need different guidance:
 56        - **ViTMatte** requires a trimap marking definite foreground, definite background,
 57          and the uncertain band between them. Pass `trimap` directly, or pass `mask` and a
 58          trimap will be derived from it via `generate_trimap`.
 59        - **Matting Anything (MAM)** needs no trimap. It prompts a frozen SAM with `prompt`
 60          and refines the resulting coarse mask into an alpha matte.
 61
 62    Args:
 63        model: A loaded matting model, either `VitMatteForImageMatting` or `MattingAnything`.
 64        image_path (Union[str, numpy.ndarray, PIL.Image.Image]): The image to matte, given as
 65                                                                 a file path, RGB array, or
 66                                                                 PIL image.
 67        processor (VitMatteImageProcessor, optional): Required for ViTMatte, which uses it to
 68                                                      stack the image and trimap into a
 69                                                      four-channel input. Default is None.
 70        mask (numpy.ndarray, optional): A binary segmentation mask used to derive a trimap
 71                                        when `trimap` is not supplied. Default is None.
 72        trimap (numpy.ndarray, optional): An explicit trimap with values `0`, `128`, and
 73                                          `255`. Takes precedence over `mask`. Default is None.
 74        prompt (dict, optional): Prompt for MAM's internal SAM, accepting `"points"`,
 75                                 `"labels"`, and/or `"boxes"` exactly as
 76                                 `garmentiq.segmentation.extract` does. Default is None.
 77        trimap_args (dict, optional): Options forwarded to `generate_trimap`, such as
 78                                      `"erode_size"` and `"dilate_size"`. Default is None.
 79        device (Union[str, torch.device], optional): The device to run inference on, e.g.
 80                                                     `"cpu"`, `"cuda"`, or `"mps"`. Hardware
 81                                                     acceleration is opt-in. Default is `"cpu"`.
 82
 83    Raises:
 84        ValueError: If the model is not a recognised matting model, if ViTMatte is used
 85                    without a processor or without either `trimap` or `mask`, if MAM is used
 86                    without a prompt, or if the requested `device` is unavailable.
 87        FileNotFoundError: If `image_path` points to a file that does not exist.
 88
 89    Returns:
 90        tuple (numpy.ndarray, numpy.ndarray): The original image as an RGB array, and the
 91                                              alpha matte as a `uint8` array in `[0, 255]`
 92                                              matching the image's height and width.
 93    """
 94    device = resolve_device(device)
 95    family = _detect_matting_family(model)
 96    if family is None:
 97        raise ValueError(
 98            "Unrecognised matting model. Expected a ViTMatte model "
 99            "(VitMatteForImageMatting) or a Matting Anything model (MattingAnything)."
100        )
101
102    if isinstance(image_path, str):
103        image = Image.open(image_path).convert("RGB")
104    elif isinstance(image_path, np.ndarray):
105        image = Image.fromarray(image_path.astype(np.uint8)).convert("RGB")
106    else:
107        image = image_path.convert("RGB")
108
109    image_np = np.array(image)
110    model = model.to(device)
111
112    if (
113        device.type == "mps"
114        and image_np.shape[0] * image_np.shape[1] > _MPS_PIXEL_WARN_THRESHOLD
115    ):
116        warnings.warn(
117            f"Matting a {image_np.shape[1]}x{image_np.shape[0]} image on the MPS backend. "
118            f"Above roughly {_MPS_PIXEL_WARN_THRESHOLD / 1e6:.1f} megapixels MPS has been "
119            f"observed to return a degenerate alpha matte that differs substantially from "
120            f"CPU. Consider device='cpu' for images this large, or downscale first.",
121            RuntimeWarning,
122            stacklevel=2,
123        )
124
125    if family == "vitmatte":
126        if processor is None:
127            raise ValueError(
128                "ViTMatte requires a processor. Pass processor=load_vitmatte_processor(...)."
129            )
130        if trimap is None:
131            if mask is None:
132                raise ValueError(
133                    "ViTMatte is a trimap-based model, so it needs either an explicit "
134                    "trimap= or a segmentation mask= to derive one from. Run segmentation "
135                    "first, then pass its mask here."
136                )
137            trimap = generate_trimap(mask, **(trimap_args or {}))
138
139        trimap_arr = np.asarray(trimap)
140        if trimap_arr.shape[:2] != image_np.shape[:2]:
141            raise ValueError(
142                f"trimap shape {trimap_arr.shape[:2]} does not match image shape "
143                f"{image_np.shape[:2]}."
144            )
145
146        # The processor rescales trimaps by its own rescale_factor (1/255), so hand it
147        # the raw 0-255 trimap. Pre-normalising here would rescale twice and collapse
148        # the trimap to near-zero, making the model see everything as background.
149        inputs = inputs_to_device(
150            processor(
151                images=image, trimaps=trimap_arr.astype(np.uint8), return_tensors="pt"
152            ),
153            device,
154        )
155
156        with torch.no_grad():
157            outputs = model(**inputs)
158
159        alpha = outputs.alphas
160        # ViTMatte pads to a multiple of its patch size; crop back to the true size.
161        alpha = alpha[..., : image_np.shape[0], : image_np.shape[1]]
162        alpha_np = alpha[0, 0].detach().cpu().numpy()
163        del inputs, outputs
164
165    else:
166        from garmentiq.matting.model_definition.mam.mam import fuse_alpha
167
168        if not prompt:
169            raise ValueError(
170                "Matting Anything prompts a frozen SAM, so it needs a prompt. Pass "
171                "prompt={'points': [[[x, y]]]} or prompt={'boxes': [[[x0, y0, x1, y1]]]}."
172            )
173
174        sam_processor = model.sam_processor
175        processor_kwargs = {}
176        if prompt.get("points") is not None:
177            processor_kwargs["input_points"] = prompt["points"]
178        if prompt.get("labels") is not None:
179            processor_kwargs["input_labels"] = prompt["labels"]
180        if prompt.get("boxes") is not None:
181            processor_kwargs["input_boxes"] = prompt["boxes"]
182
183        inputs = inputs_to_device(
184            sam_processor(images=image, return_tensors="pt", **processor_kwargs),
185            device,
186        )
187
188        pixel_values = inputs.pop("pixel_values")
189        original_sizes = inputs.pop("original_sizes")
190        reshaped_sizes = inputs.pop("reshaped_input_sizes")
191
192        with torch.no_grad():
193            pred, post_mask = model(
194                pixel_values=pixel_values,
195                reshaped_size=reshaped_sizes[0].tolist(),
196                original_size=original_sizes[0].tolist(),
197                **inputs,
198            )
199            alpha = fuse_alpha(
200                pred,
201                reshaped_size=reshaped_sizes[0].tolist(),
202                original_size=original_sizes[0].tolist(),
203                post_mask=post_mask,
204            )
205
206        alpha_np = alpha[0, 0].detach().cpu().numpy()
207        del inputs, pred
208
209    alpha_np = np.clip(alpha_np, 0.0, 1.0)
210    alpha_u8 = (alpha_np * 255).astype(np.uint8)
211
212    empty_cache(device)
213    return image_np, alpha_u8
214
215
216def composite(
217    image_np: np.ndarray,
218    alpha_np: np.ndarray,
219    background_color=(255, 255, 255),
220):
221    """
222    Composites a matted foreground onto a solid background colour.
223
224    Unlike hard mask replacement, this blends each pixel by its alpha value, so soft edges
225    stay soft instead of showing a cut-out outline.
226
227    Args:
228        image_np (numpy.ndarray): The original RGB image, shape `(H, W, 3)`.
229        alpha_np (numpy.ndarray): The alpha matte, shape `(H, W)`, `uint8` in `[0, 255]`
230                                  or float in `[0, 1]`.
231        background_color (tuple[int, int, int], optional): RGB colour to composite onto.
232                                                           Default is white.
233
234    Raises:
235        ValueError: If the image and alpha shapes do not match.
236
237    Returns:
238        numpy.ndarray: The composited RGB image as `uint8`, shape `(H, W, 3)`.
239    """
240    image_np = np.asarray(image_np)
241    alpha_np = np.asarray(alpha_np)
242
243    if image_np.shape[:2] != alpha_np.shape[:2]:
244        raise ValueError(
245            f"image shape {image_np.shape[:2]} does not match alpha shape "
246            f"{alpha_np.shape[:2]}."
247        )
248
249    alpha = alpha_np.astype(np.float32)
250    if alpha.max() > 1.0:
251        alpha = alpha / 255.0
252    alpha = alpha[..., None]
253
254    background = np.array(background_color, dtype=np.float32).reshape(1, 1, 3)
255    out = alpha * image_np.astype(np.float32) + (1.0 - alpha) * background
256    return np.clip(out, 0, 255).astype(np.uint8)
257
258
259__all__ = ["matte", "composite"]