garmentiq.matting

12def generate_trimap(
13    mask: np.ndarray,
14    erode_size: int = 10,
15    dilate_size: int = 10,
16    threshold: int = 127,
17):
18    """
19    Builds a trimap from a binary segmentation mask.
20
21    Trimap-based matting models such as ViTMatte need to be told which pixels are
22    definitely foreground, definitely background, and uncertain. This function derives
23    that from a segmentation mask by eroding it to obtain confident foreground, dilating
24    it to obtain confident background, and marking the band in between as unknown, which
25    is where the matting model is free to predict soft alpha values.
26
27    The size of the unknown band matters: too narrow and hair or fabric edges fall outside
28    it and can never become soft, too wide and the model has to guess over large areas.
29
30    Args:
31        mask (numpy.ndarray): Segmentation mask as a 2D array. Values may be binary
32                              (`0`/`1`) or 8-bit (`0`-`255`).
33        erode_size (int, optional): Erosion kernel size controlling how far the confident
34                                    foreground is pulled inward. Default is 10.
35        dilate_size (int, optional): Dilation kernel size controlling how far the unknown
36                                     band extends outward. Default is 10.
37        threshold (int, optional): Cutoff used to binarise an 8-bit mask. Default is 127.
38
39    Raises:
40        ValueError: If `mask` is not 2D, or if `erode_size` or `dilate_size` is negative.
41
42    Returns:
43        numpy.ndarray: Trimap as `uint8` with values `0` (background), `128` (unknown),
44                       and `255` (foreground), matching `mask` in shape.
45    """
46    if mask is None or not hasattr(mask, "ndim"):
47        raise ValueError("mask must be a numpy array.")
48
49    mask = np.asarray(mask)
50    if mask.ndim == 3 and mask.shape[-1] == 1:
51        mask = mask[..., 0]
52    if mask.ndim != 2:
53        raise ValueError(
54            f"mask must be a 2D array, got shape {mask.shape}. Pass a single-channel "
55            f"segmentation mask."
56        )
57    if erode_size < 0 or dilate_size < 0:
58        raise ValueError(
59            f"erode_size and dilate_size must be non-negative, got "
60            f"erode_size={erode_size}, dilate_size={dilate_size}."
61        )
62
63    if mask.dtype == bool:
64        binary = mask.astype(np.uint8)
65    elif mask.max() <= 1:
66        binary = (mask > 0).astype(np.uint8)
67    else:
68        binary = (mask > threshold).astype(np.uint8)
69
70    foreground = binary
71    if erode_size > 0:
72        erode_kernel = cv2.getStructuringElement(
73            cv2.MORPH_ELLIPSE, (erode_size, erode_size)
74        )
75        foreground = cv2.erode(binary, erode_kernel, iterations=1)
76
77    background = binary
78    if dilate_size > 0:
79        dilate_kernel = cv2.getStructuringElement(
80            cv2.MORPH_ELLIPSE, (dilate_size, dilate_size)
81        )
82        background = cv2.dilate(binary, dilate_kernel, iterations=1)
83
84    trimap = np.full(binary.shape, 128, dtype=np.uint8)
85    trimap[background == 0] = 0
86    trimap[foreground == 1] = 255
87    return trimap

 38def matte(
 39    model,
 40    image_path: Union[str, np.ndarray, Image.Image],
 41    processor=None,
 42    mask: Optional[np.ndarray] = None,
 43    trimap: Optional[np.ndarray] = None,
 44    prompt: Optional[dict] = None,
 45    trimap_args: Optional[dict] = None,
 46    device: Union[str, torch.device] = "cpu",
 47):
 48    """
 49    Extracts a soft alpha matte for an image, refining a hard segmentation into soft edges.
 50
 51    Segmentation answers "which pixels belong to the garment" with a yes/no decision, which
 52    leaves stair-stepped borders and loses semi-transparent detail such as loose fabric,
 53    lace, and stray fibres. Matting instead predicts a continuous alpha value per pixel, so
 54    compositing onto a new background looks natural.
 55
 56    The two supported approaches need different guidance:
 57        - **ViTMatte** requires a trimap marking definite foreground, definite background,
 58          and the uncertain band between them. Pass `trimap` directly, or pass `mask` and a
 59          trimap will be derived from it via `generate_trimap`.
 60        - **Matting Anything (MAM)** needs no trimap. It prompts a frozen SAM with `prompt`
 61          and refines the resulting coarse mask into an alpha matte.
 62
 63    Args:
 64        model: A loaded matting model, either `VitMatteForImageMatting` or `MattingAnything`.
 65        image_path (Union[str, numpy.ndarray, PIL.Image.Image]): The image to matte, given as
 66                                                                 a file path, RGB array, or
 67                                                                 PIL image.
 68        processor (VitMatteImageProcessor, optional): Required for ViTMatte, which uses it to
 69                                                      stack the image and trimap into a
 70                                                      four-channel input. Default is None.
 71        mask (numpy.ndarray, optional): A binary segmentation mask used to derive a trimap
 72                                        when `trimap` is not supplied. Default is None.
 73        trimap (numpy.ndarray, optional): An explicit trimap with values `0`, `128`, and
 74                                          `255`. Takes precedence over `mask`. Default is None.
 75        prompt (dict, optional): Prompt for MAM's internal SAM, accepting `"points"`,
 76                                 `"labels"`, and/or `"boxes"` exactly as
 77                                 `garmentiq.segmentation.extract` does. Default is None.
 78        trimap_args (dict, optional): Options forwarded to `generate_trimap`, such as
 79                                      `"erode_size"` and `"dilate_size"`. Default is None.
 80        device (Union[str, torch.device], optional): The device to run inference on, e.g.
 81                                                     `"cpu"`, `"cuda"`, or `"mps"`. Hardware
 82                                                     acceleration is opt-in. Default is `"cpu"`.
 83
 84    Raises:
 85        ValueError: If the model is not a recognised matting model, if ViTMatte is used
 86                    without a processor or without either `trimap` or `mask`, if MAM is used
 87                    without a prompt, or if the requested `device` is unavailable.
 88        FileNotFoundError: If `image_path` points to a file that does not exist.
 89
 90    Returns:
 91        tuple (numpy.ndarray, numpy.ndarray): The original image as an RGB array, and the
 92                                              alpha matte as a `uint8` array in `[0, 255]`
 93                                              matching the image's height and width.
 94    """
 95    device = resolve_device(device)
 96    family = _detect_matting_family(model)
 97    if family is None:
 98        raise ValueError(
 99            "Unrecognised matting model. Expected a ViTMatte model "
100            "(VitMatteForImageMatting) or a Matting Anything model (MattingAnything)."
101        )
102
103    if isinstance(image_path, str):
104        image = Image.open(image_path).convert("RGB")
105    elif isinstance(image_path, np.ndarray):
106        image = Image.fromarray(image_path.astype(np.uint8)).convert("RGB")
107    else:
108        image = image_path.convert("RGB")
109
110    image_np = np.array(image)
111    model = model.to(device)
112
113    if (
114        device.type == "mps"
115        and image_np.shape[0] * image_np.shape[1] > _MPS_PIXEL_WARN_THRESHOLD
116    ):
117        warnings.warn(
118            f"Matting a {image_np.shape[1]}x{image_np.shape[0]} image on the MPS backend. "
119            f"Above roughly {_MPS_PIXEL_WARN_THRESHOLD / 1e6:.1f} megapixels MPS has been "
120            f"observed to return a degenerate alpha matte that differs substantially from "
121            f"CPU. Consider device='cpu' for images this large, or downscale first.",
122            RuntimeWarning,
123            stacklevel=2,
124        )
125
126    if family == "vitmatte":
127        if processor is None:
128            raise ValueError(
129                "ViTMatte requires a processor. Pass processor=load_vitmatte_processor(...)."
130            )
131        if trimap is None:
132            if mask is None:
133                raise ValueError(
134                    "ViTMatte is a trimap-based model, so it needs either an explicit "
135                    "trimap= or a segmentation mask= to derive one from. Run segmentation "
136                    "first, then pass its mask here."
137                )
138            trimap = generate_trimap(mask, **(trimap_args or {}))
139
140        trimap_arr = np.asarray(trimap)
141        if trimap_arr.shape[:2] != image_np.shape[:2]:
142            raise ValueError(
143                f"trimap shape {trimap_arr.shape[:2]} does not match image shape "
144                f"{image_np.shape[:2]}."
145            )
146
147        # The processor rescales trimaps by its own rescale_factor (1/255), so hand it
148        # the raw 0-255 trimap. Pre-normalising here would rescale twice and collapse
149        # the trimap to near-zero, making the model see everything as background.
150        inputs = inputs_to_device(
151            processor(
152                images=image, trimaps=trimap_arr.astype(np.uint8), return_tensors="pt"
153            ),
154            device,
155        )
156
157        with torch.no_grad():
158            outputs = model(**inputs)
159
160        alpha = outputs.alphas
161        # ViTMatte pads to a multiple of its patch size; crop back to the true size.
162        alpha = alpha[..., : image_np.shape[0], : image_np.shape[1]]
163        alpha_np = alpha[0, 0].detach().cpu().numpy()
164        del inputs, outputs
165
166    else:
167        from garmentiq.matting.model_definition.mam.mam import fuse_alpha
168
169        if not prompt:
170            raise ValueError(
171                "Matting Anything prompts a frozen SAM, so it needs a prompt. Pass "
172                "prompt={'points': [[[x, y]]]} or prompt={'boxes': [[[x0, y0, x1, y1]]]}."
173            )
174
175        sam_processor = model.sam_processor
176        processor_kwargs = {}
177        if prompt.get("points") is not None:
178            processor_kwargs["input_points"] = prompt["points"]
179        if prompt.get("labels") is not None:
180            processor_kwargs["input_labels"] = prompt["labels"]
181        if prompt.get("boxes") is not None:
182            processor_kwargs["input_boxes"] = prompt["boxes"]
183
184        inputs = inputs_to_device(
185            sam_processor(images=image, return_tensors="pt", **processor_kwargs),
186            device,
187        )
188
189        pixel_values = inputs.pop("pixel_values")
190        original_sizes = inputs.pop("original_sizes")
191        reshaped_sizes = inputs.pop("reshaped_input_sizes")
192
193        with torch.no_grad():
194            pred, post_mask = model(
195                pixel_values=pixel_values,
196                reshaped_size=reshaped_sizes[0].tolist(),
197                original_size=original_sizes[0].tolist(),
198                **inputs,
199            )
200            alpha = fuse_alpha(
201                pred,
202                reshaped_size=reshaped_sizes[0].tolist(),
203                original_size=original_sizes[0].tolist(),
204                post_mask=post_mask,
205            )
206
207        alpha_np = alpha[0, 0].detach().cpu().numpy()
208        del inputs, pred
209
210    alpha_np = np.clip(alpha_np, 0.0, 1.0)
211    alpha_u8 = (alpha_np * 255).astype(np.uint8)
212
213    empty_cache(device)
214    return image_np, alpha_u8

217def composite(
218    image_np: np.ndarray,
219    alpha_np: np.ndarray,
220    background_color=(255, 255, 255),
221):
222    """
223    Composites a matted foreground onto a solid background colour.
224
225    Unlike hard mask replacement, this blends each pixel by its alpha value, so soft edges
226    stay soft instead of showing a cut-out outline.
227
228    Args:
229        image_np (numpy.ndarray): The original RGB image, shape `(H, W, 3)`.
230        alpha_np (numpy.ndarray): The alpha matte, shape `(H, W)`, `uint8` in `[0, 255]`
231                                  or float in `[0, 1]`.
232        background_color (tuple[int, int, int], optional): RGB colour to composite onto.
233                                                           Default is white.
234
235    Raises:
236        ValueError: If the image and alpha shapes do not match.
237
238    Returns:
239        numpy.ndarray: The composited RGB image as `uint8`, shape `(H, W, 3)`.
240    """
241    image_np = np.asarray(image_np)
242    alpha_np = np.asarray(alpha_np)
243
244    if image_np.shape[:2] != alpha_np.shape[:2]:
245        raise ValueError(
246            f"image shape {image_np.shape[:2]} does not match alpha shape "
247            f"{alpha_np.shape[:2]}."
248        )
249
250    alpha = alpha_np.astype(np.float32)
251    if alpha.max() > 1.0:
252        alpha = alpha / 255.0
253    alpha = alpha[..., None]
254
255    background = np.array(background_color, dtype=np.float32).reshape(1, 1, 3)
256    out = alpha * image_np.astype(np.float32) + (1.0 - alpha) * background
257    return np.clip(out, 0, 255).astype(np.uint8)

12def load_model(
13    model_class: Type[nn.Module],
14    model_path: str,
15    model_args: dict = None,
16    device: Union[str, torch.device] = "cpu",
17    **kwargs,
18):
19    """
20    Loads a matting model from a local checkpoint and prepares it for inference.
21
22    This mirrors `garmentiq.segmentation.load_model`: the model class is instantiated with
23    safely filtered configuration arguments, weights are read from a local `.pth` or
24    `.safetensors` file, the model is moved to the requested device and set to evaluation
25    mode, and common weight prefixes are stripped for compatibility.
26
27    Matting Anything checkpoints are not loaded here because they pair a decoder with a
28    separate SAM model; use `garmentiq.matting.model_definition.mam.load_mam` for those.
29
30    Args:
31        model_class (Type[nn.Module]): The uninstantiated model class, typically
32                                       `VitMatteForImageMatting`.
33        model_path (str): Local path to the checkpoint weights, ending in `.pth` or
34                          `.safetensors`.
35        model_args (dict, optional): Configuration arguments for initializing the model,
36                                     e.g. `{"config": load_vitmatte_config(...)}`.
37                                     Incompatible arguments are safely ignored.
38                                     Default is None.
39        device (Union[str, torch.device], optional): The device to load the model onto, e.g.
40                                                     `"cpu"`, `"cuda"`, or `"mps"`. Hardware
41                                                     acceleration is opt-in. Default is `"cpu"`.
42        **kwargs: Additional arbitrary keyword arguments.
43
44    Raises:
45        ValueError: If the requested `device` is invalid or unavailable on this machine.
46        Exception: If the weights cannot be loaded or the file format is unsupported.
47
48    Returns:
49        nn.Module: The loaded model in evaluation mode, placed on `device`.
50    """
51    model_args = model_args or {}
52
53    if not isinstance(model_args, dict):
54        if hasattr(model_args, "to_dict"):
55            model_args = model_args.to_dict()
56        elif hasattr(model_args, "__dict__"):
57            model_args = vars(model_args)
58        else:
59            raise TypeError(
60                "model_args must be a dictionary or a configuration object."
61            )
62
63    device = resolve_device(device)
64
65    sig = inspect.signature(model_class.__init__)
66    valid_params = set(sig.parameters.keys())
67    has_kwargs = any(
68        p.kind == inspect.Parameter.VAR_KEYWORD for p in sig.parameters.values()
69    )
70
71    if not has_kwargs:
72        filtered_args = {k: v for k, v in model_args.items() if k in valid_params}
73    else:
74        filtered_args = model_args
75
76    model = model_class(**filtered_args).to(device)
77
78    if model_path.endswith(".safetensors"):
79        state_dict = load_file(model_path, device=str(device))
80    else:
81        state_dict = torch.load(model_path, map_location=device, weights_only=True)
82
83    new_state_dict = {
84        k.removeprefix("module.").removeprefix("model."): v
85        for k, v in state_dict.items()
86    }
87
88    load_state_dict_checked(model, new_state_dict, model_path)
89    model.eval()
90
91    return model