Text-guided removal: auto-generate the mask with SAM 3

README.md

@@ -30,7 +30,7 @@ object you point at and leaves the rest of the photo as it was. Paint a rough ma
 say what it is, and it is replaced with what would be behind it, its shadow or reflection included.
 
 - **`/` — Studio** (`gr.Server` API + custom UI, Pruna design, dark and light themes): paint the mask
-  with a brush and eraser, write the instruction, pick a variant; the card's 18 examples load photo,
+  with a brush and eraser, or let SAM 3 find it from the instruction, write the instruction, pick a variant; the card's 18 examples load photo,
   mask and instruction in one click; before/after split view and run history.
 - **`/workflow` — Workflow canvas** (`gr.Workflow`): photo + mask + instruction into a removal node,
   chained into a second removal. To change the starter graph, edit `workflow.json`.
@@ -45,8 +45,17 @@ The mask is drawn as the red outline the model expects (`draw_outline.py`, from
 the matching sentence is added to the instruction. With no mask, the instruction alone says what to
 remove. No CFG.
 
+**Text-guided masks.** Instead of painting, write the instruction and press *Auto-generate mask from instruction*:
+[SAM 3](https://huggingface.co/facebook/sam3) masks every object the instruction names
+(`text_mask.py`; "the phone on the right" picks one), and the mask lands on the canvas to touch up.
+A run with nothing painted does this by itself, in the same GPU call as the removal, and falls back
+to the instruction alone if SAM 3 finds nothing. `facebook/sam3` is gated: the Space needs an `HF_TOKEN` secret with
+access, and without one it runs with painted masks only.
+
 The `remove` API endpoint takes `source` and `mask` (uploaded file paths; mask white = remove, may
-be empty), `instruction`, `variant`, `seed` and `randomize`.
+be empty), `instruction`, `variant`, `seed`, `randomize` and `auto_mask` (with no mask, find it from
+the instruction with SAM 3; the mask found comes back as `mask_data`). The `find_mask` endpoint takes
+`source` and `text` and returns the SAM 3 mask as `mask_data`.
 
 Example photos are from [Unsplash](https://unsplash.com) (credits on the model card); the masks are
 the card's own, except `c9`'s, rebuilt from its outlined input.

app.py

@@ -3,6 +3,9 @@
   - "/"          Studio: custom UI (studio.html) over the `remove` API endpoint
   - "/workflow"  Workflow canvas: node graph from workflow.json (pq_remove)
 
+Text-guided masks: with no painted mask, SAM 3 can find the object the instruction names
+(text_mask.py) and the removal runs on its outline, in the same GPU call.
+
 The removal LoRA runs as its model card says: the few-step variants on top of PrunaAI's distillation
 LoRA (at 1.0) at strength 0.7, the base variant alone at 40 steps; no CFG in all cases.
 """
@@ -29,6 +32,7 @@ from diffusers import FlowMatchEulerDiscreteScheduler, QwenImage21Pipeline
 
 from draw_outline import draw_outline, outline_prompt
 from examples import EXAMPLES
+from text_mask import TextMasker
 
 APP_DIR = os.path.dirname(os.path.abspath(__file__))
 EXAMPLES_DIR = os.path.join(APP_DIR, "examples")
@@ -118,6 +122,24 @@ def _activate(name: str) -> dict:
 
 _activate(DEFAULT_VARIANT)
 pipe = pipe.to("cuda")
+# Decoding a ~2 MP result in one piece runs out of memory on a ZeroGPU slot; tiles of 1024 px leave
+# results up to 1024 px a side decoded exactly as before.
+pipe.vae.enable_tiling(
+    tile_sample_min_height=1024, tile_sample_min_width=1024,
+    tile_sample_stride_height=768, tile_sample_stride_width=768,
+)
+
+# SAM 3 for text-guided masks. facebook/sam3 is gated: without access (HF_TOKEN) the Space still
+# runs, with painted masks only. It stays in CPU memory and visits the GPU only to find a mask, so
+# the removal keeps the GPU memory it had without it.
+try:
+    masker = TextMasker(device="cpu")
+except Exception as error:  # noqa: BLE001
+    print(f"Text-guided masks disabled, SAM 3 did not load: {error}")
+    masker = None
+
+# SAM 3 works at 1008 px; segmenting a larger copy only costs time.
+FIND_MAX_SIDE = 1536
 
 
 def _resolve_variant(variant) -> str:
@@ -132,13 +154,36 @@ def _resolve_variant(variant) -> str:
 # ------------------------------------------------------------------ rendering
 
 
-def _gpu_seconds(prompt, variant, seed, image, resolution):
-    return VARIANTS[variant]["gpu_seconds"]
+def _find(image: Image.Image, text: str):
+    """SAM 3 mask (L, white = the object) of what `text` names in `image`, at its size, or None."""
+    if masker is None or not text:
+        return None
+    small = image
+    if max(image.size) > FIND_MAX_SIDE:
+        scale = FIND_MAX_SIDE / max(image.size)
+        small = image.resize((round(image.width * scale), round(image.height * scale)), Image.LANCZOS)
+    try:
+        masker.model.to("cuda")
+        mask, _ = masker(small, text)
+    finally:
+        masker.model.to("cpu")
+        torch.cuda.empty_cache()
+    return None if mask is None else mask.resize(image.size, Image.NEAREST)
+
+
+def _gpu_seconds(prompt, variant, seed, image, resolution, find=""):
+    return VARIANTS[variant]["gpu_seconds"] + (10 if find else 0)
 
 
 @spaces.GPU(duration=_gpu_seconds)
-def _render(prompt: str, variant: str, seed: int, image: Image.Image, resolution: int):
+def _render(prompt: str, variant: str, seed: int, image: Image.Image, resolution: int, find: str = ""):
+    """The removal. With `find`, SAM 3 first looks for the object it names; when found, the photo
+    gets its red outline and the prompt the outline sentence. Returns (image, seconds, mask or None)."""
     spec = _activate(variant)
+    mask = _find(image, find) if find else None
+    if mask is not None:
+        image = draw_outline(image, mask)
+        prompt = outline_prompt(prompt)
     torch.cuda.synchronize()
     start = time.perf_counter()
     result = pipe(
@@ -151,7 +196,12 @@ def _render(prompt: str, variant: str, seed: int, image: Image.Image, resolution
         generator=torch.Generator("cuda").manual_seed(seed),
     ).images[0]
     torch.cuda.synchronize()
-    return result, time.perf_counter() - start
+    return result, time.perf_counter() - start, mask
+
+
+@spaces.GPU(duration=20)
+def _find_gpu(image: Image.Image, text: str):
+    return _find(image, text)
 
 
 def _canvas(ratio: float, side: int) -> tuple[int, int]:
@@ -160,33 +210,39 @@ def _canvas(ratio: float, side: int) -> tuple[int, int]:
     return round(width / 32) * 32, round(width / ratio / 32) * 32
 
 
-def _remove(photo: Image.Image, mask, instruction: str, variant, seed: int):
+def _remove(photo: Image.Image, mask, instruction: str, variant, seed: int, auto_mask: bool = False):
     """Object removal as the model card does it: work at the photo's area (sides rounded to 32,
     capped per variant), draw the red outline at that size with the matching sentence, and resize
     the result back to the photo's aspect ratio. Without a mask (or an empty one) it runs on the
-    instruction alone. Returns (image, gpu seconds, variant, outlined)."""
+    instruction alone, or, with `auto_mask`, on the outline of what SAM 3 finds from the
+    instruction. Returns (image, gpu seconds, variant, mode, found mask at the photo's size or None)."""
     name = _resolve_variant(variant)
     instruction = (instruction or "").strip()
     area = photo.width * photo.height
     side = min(round(area**0.5), VARIANTS[name]["max_side"])
     work_size = _canvas(photo.width / photo.height, side)
     work = photo.resize(work_size, Image.LANCZOS)
-    outlined = False
+    mode = "prompt only"
     if mask is not None:
         mask = mask.resize(work_size, Image.NEAREST)
         if (np.array(mask) > 127).any():
             work = draw_outline(work, mask)
             instruction = outline_prompt(instruction or DEFAULT_INSTRUCTION)
-            outlined = True
+            mode = "outline"
     if not instruction:
         raise gr.Error("Paint over the object to remove, or describe it.")
+    find = instruction if auto_mask and mode == "prompt only" and masker is not None else ""
     try:
-        result, seconds = _render(instruction, name, seed, work, side)
+        result, seconds, found = _render(instruction, name, seed, work, side, find)
     except (RuntimeError, ValueError) as error:
         raise gr.Error(str(error).split("\n")[0])
+    if find:
+        mode = "text mask (SAM 3)" if found is not None else "prompt only, SAM 3 found nothing"
+    if found is not None:
+        found = found.resize(photo.size, Image.NEAREST)
     scale = side / area**0.5
     out_size = (round(photo.width * scale), round(photo.height * scale)) if scale < 1 else photo.size
-    return result.resize(out_size, Image.LANCZOS), seconds, name, outlined
+    return result.resize(out_size, Image.LANCZOS), seconds, name, mode, found
 
 
 # ------------------------------------------------------------------ helpers
@@ -307,7 +363,7 @@ def pq_remove(image, prompt="", variant: str = DEFAULT_VARIANT, seed: str = "42"
     """Remove an object with a mask outline, a text instruction, or both (no CFG)."""
     if _empty_media(image):
         raise gr.Error("Connect a source image first.")
-    result, _, _, _ = _remove(
+    result, _, _, _, _ = _remove(
         _open_rgb(image), _open_mask(mask), _as_instruction(prompt), variant, _seed(seed)
     )
     return _as_media(_publish(result))
@@ -328,6 +384,7 @@ def _studio_page():
         "examples": [{"id": eid, "label": label, "prompt": prompt} for eid, label, prompt in EXAMPLES],
         "variants": VARIANT_CHOICES,
         "defaultVariant": DEFAULT_VARIANT,
+        "textMask": masker is not None,
     }
     return page.replace("__STUDIO_CONFIG__", json.dumps(config).replace("</", "<\\/"))
 
@@ -352,15 +409,36 @@ def api_remove(
     variant: str = DEFAULT_VARIANT,
     seed: float = 42,
     randomize: bool = True,
+    auto_mask: bool = False,
 ) -> dict:
-    """Remove the masked object (white = remove) with P-Q-Image-2.1-Remove-Object."""
+    """Remove the masked object (white = remove) with P-Q-Image-2.1-Remove-Object. With no mask and
+    `auto_mask`, SAM 3 finds the object the instruction names and its mask is returned as `mask_data`."""
     if not source:
         raise gr.Error("Upload a photo first.")
     seed_value = _resolve_seed(seed, randomize)
-    image, seconds, name, outlined = _remove(_open_rgb(source), _open_mask(mask), instruction, variant, seed_value)
-    mode = "outline" if outlined else "prompt only"
+    image, seconds, name, mode, found = _remove(
+        _open_rgb(source), _open_mask(mask), instruction, variant, seed_value, auto_mask
+    )
     info = f"{image.width}×{image.height}, seed {seed_value} · {name} · {mode} · {seconds:.1f} s on the GPU"
-    return {"image_data": _as_data_uri(image), "seed": seed_value, "info": info}
+    payload = {"image_data": _as_data_uri(image), "seed": seed_value, "info": info}
+    if found is not None:
+        payload["mask_data"] = _as_data_uri(found)
+    return payload
+
+
+@server.api(name="find_mask", concurrency_limit=2)
+def api_find_mask(source: str, text: str = "") -> dict:
+    """Mask (white = the object) of what `text` names in the photo, found with SAM 3."""
+    if masker is None:
+        raise gr.Error("Text-guided masks are not available on this Space.")
+    if not source:
+        raise gr.Error("Upload a photo first.")
+    if not (text or "").strip():
+        raise gr.Error("Write what to remove first.")
+    mask = _find_gpu(_open_rgb(source), text.strip())
+    if mask is None:
+        raise gr.Error("SAM 3 found nothing matching the instruction. Try naming the object more simply.")
+    return {"mask_data": _as_data_uri(mask)}
 
 
 # The @server.api endpoints are deferred until Server.launch(), which also fires

studio.html

@@ -239,6 +239,11 @@
   .tool-toggle button.active { background: var(--accent-soft); color: var(--lavender); }
   .mask-clear { border: 1px solid var(--border); background: var(--raised); margin-top: 8px; width: 100%; }
   .mask-clear:hover { color: var(--text); border-color: var(--accent); }
+  .panel label.instruction-label { margin-top: 14px; padding-top: 12px; border-top: 1px solid var(--hairline); }
+  .mask-find { display: inline-flex; align-items: center; justify-content: center; gap: 5px; }
+  .mask-find { min-width: 0; }
+  .mask-find svg { width: 13px; height: 13px; flex-shrink: 0; }
+  .mask-find:disabled { opacity: 0.5; cursor: wait; }
   input[type="range"] { width: 100%; accent-color: var(--accent); }
   @media (max-width: 860px) {
     html, body { height: auto !important; overflow: visible !important; }
@@ -355,12 +360,11 @@
         <label for="brush-size">Brush size</label>
         <input type="range" id="brush-size" min="4" max="120" value="28" />
         <button class="mask-clear" id="mask-clear">Clear mask</button>
-        <div class="hint">Paint roughly over the object; include its shadow or reflection if they should go too, and keep the mask off anything that should stay. It is drawn as a red outline for the model. With no mask, the instruction alone says what to remove.</div>
-      </section>
-      <section class="panel">
-        <div class="panel-title">Instruction</div>
+        <div class="hint" id="mask-hint">Paint roughly over the object; include its shadow or reflection if they should go too, and keep the mask off anything that should stay. It is drawn as a red outline for the model. With no mask, the instruction alone says what to remove.</div>
+        <label for="prompt" class="instruction-label">Instruction</label>
         <textarea id="prompt" rows="4" placeholder="e.g. Remove the man on the right."></textarea>
         <div class="hint">Name the object as it appears, and say which one when there are several.</div>
+        <button class="mask-clear mask-find" id="mask-find" title="Mask what the instruction names, with SAM 3" hidden><svg viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round" aria-hidden="true"><path d="M12 3l1.9 5.1L19 10l-5.1 1.9L12 17l-1.9-5.1L5 10l5.1-1.9z"/><path d="M19 15l.8 2.2L22 18l-2.2.8L19 21l-.8-2.2L16 18l2.2-.8z"/></svg>Auto-generate mask from instruction</button>
       </section>
       <section class="panel">
         <div class="panel-title">Variant</div>
@@ -627,6 +631,32 @@
     return state.maskPath;
   }
 
+  // ---------- text-guided mask (SAM 3) ----------
+  if (CONFIG.textMask) {
+    $("mask-find").hidden = false;
+    $("mask-hint").textContent = "Paint roughly over the object, its shadow or reflection too if they should go, or let SAM 3 find it from the instruction below and touch that up; with nothing painted, Remove Object finds it by itself. The mask is drawn as a red outline for the model.";
+  }
+  $("mask-find").onclick = async () => {
+    const prompt = $("prompt").value.trim();
+    if (!state.sourcePath) { setStatus("Upload a photo or pick an example first."); return; }
+    if (!prompt) { setStatus("Write what to remove in the instruction first."); $("prompt").focus(); return; }
+    const sourceUrl = state.sourceUrl;
+    setStatus("Finding the object…", true);
+    $("mask-find").disabled = true;
+    try {
+      const app = await getApp();
+      const payload = unwrap(await app.predict("/find_mask", { source: state.sourcePath, text: prompt }));
+      if (typeof payload.mask_data !== "string") { setStatus("Nothing found (see console)."); return; }
+      if (state.sourceUrl !== sourceUrl) return;  // the photo changed while it ran
+      await loadMask(payload.mask_data);
+      setStatus("Mask found. Touch it up if needed, then hit Remove Object");
+    } catch (err) {
+      setStatus(err?.message || String(err));
+    } finally {
+      $("mask-find").disabled = false;
+    }
+  };
+
   // ---------- history ----------
   function showResultImage(uri) {
     $("result-img").src = uri;
@@ -757,7 +787,8 @@
     const run = {
       prompt, variant: $("variant").value, sourceUrl: state.sourceUrl, sourcePath: state.sourcePath,
     };
-    setStatus("Removing…", true);
+    const autoMask = CONFIG.textMask && !state.maskPainted;
+    setStatus(autoMask ? "Finding the object and removing it…" : "Removing…", true);
     try {
       const app = await getApp();
       const result = await app.predict("/remove", {
@@ -767,6 +798,7 @@
         variant: run.variant,
         seed: parseInt($("seed").value, 10) || 0,
         randomize: $("randomize").checked,
+        auto_mask: autoMask,
       });
       const payload = unwrap(result);
       if (typeof payload.image_data !== "string") { setStatus("Run finished but no image came back (see console)."); return; }
@@ -774,6 +806,8 @@
       if (state.sourceUrl === run.sourceUrl) {  // the photo may have changed while it ran
         $("result-meta").textContent = payload.info || "Done";
         showResultImage(payload.image_data);
+        // Show the mask SAM 3 found, so the next run can reuse or refine it.
+        if (typeof payload.mask_data === "string" && !state.maskPainted) await loadMask(payload.mask_data);
       }
       history.unshift({ ...run, uri: payload.image_data, seed: payload.seed, info: payload.info || "" });
       if (history.length > 8) history.pop();

text_mask.py

@@ -0,0 +1,152 @@
+"""Find the object to remove from a text description with SAM 3 (text-prompted segmentation).
+
+The instruction ("Remove the wooden ladder, the white stool and the potted plant.") is turned into
+short noun phrases ("wooden ladder", "white stool", "potted plant"), SAM 3 segments every instance
+of each, and the union is dilated a little so it reads like a rough painted mask: white = remove.
+"""
+
+from __future__ import annotations
+
+import re
+
+import cv2
+import numpy as np
+import torch
+from PIL import Image
+
+SAM3_REPO = "facebook/sam3"
+
+_VERBS = r"(?:please\s+)?(?:remove|erase|delete|get rid of|take out|take away|clear|eliminate|cut out)"
+_LEADING = re.compile(rf"^\s*{_VERBS}\s+", re.IGNORECASE)
+_ARTICLE = re.compile(r"^(?:the|a|an|all(?: of)?(?: the)?|both|with|plus)\s+", re.IGNORECASE)
+_SPLIT = re.compile(r"\s*(?:,|;|\band\b|\balong with\b|\bas well as\b)\s*", re.IGNORECASE)
+# Where the object is: stripped from the phrase SAM 3 sees, then used to pick among the instances.
+_WHERE = re.compile(
+    r"\s*\b(?:(?:on|to|at|in)\s+the\s+)?(?:far\s+)?(left|right|middle|center|centre|top|bottom|foreground|background)"
+    r"(?:\s+(?:side|edge|corner))?(?:\s+of\s+the\s+(?:image|photo|picture|frame))?\b",
+    re.IGNORECASE,
+)
+# Where a long description stops naming the object ("person | with the umbrella walking ...").
+_QUALIFIER = re.compile(
+    r"\s+(?:with|in|on|at|of|near|next to|behind|beside|holding|wearing|walking|standing|sitting|"
+    r"lying|running|riding|that|which|who)\b.*$",
+    re.IGNORECASE,
+)
+_COUNT = re.compile(r"^(?:two|three|four|five|six|\d+|several|many|some|all)\b", re.IGNORECASE)
+
+
+def phrases_from_instruction(text: str) -> list[str]:
+    """'Remove the boathouse and the red canoes on the shore.' -> ['boathouse', 'red canoes on the shore']"""
+    text = (text or "").strip()
+    # Only the first sentence names the object; later ones are extra guidance.
+    text = re.split(r"(?<=[.!?])\s+", text, maxsplit=1)[0].rstrip(" .!?")
+    text = _LEADING.sub("", text)
+    phrases = []
+    for part in _SPLIT.split(text):
+        part = part.strip()
+        while (stripped := _ARTICLE.sub("", part)) != part:
+            part = stripped
+        if part and part.lower() not in phrases:
+            phrases.append(part.lower())
+    return phrases
+
+
+def _split_where(phrase: str) -> tuple[str, list[str]]:
+    where = [w.lower() for w in _WHERE.findall(phrase)]
+    stripped = re.sub(r"\s+", " ", _WHERE.sub(" ", phrase)).strip()
+    return (stripped or phrase), where
+
+
+def _pick(masks: np.ndarray, scores: np.ndarray, where: list[str]) -> np.ndarray:
+    """The one instance the position words point at (instances: N x H x W bool): the instances on
+    that side of the photo, then the most prominent of them (people point at the obvious one)."""
+    h, w = masks.shape[1:]
+    sides = {
+        "left": lambda cx, cy: cx < 0.5, "right": lambda cx, cy: cx > 0.5,
+        "top": lambda cx, cy: cy < 0.5, "bottom": lambda cx, cy: cy > 0.5,
+        "middle": lambda cx, cy: abs(cx - 0.5) < 0.25, "center": lambda cx, cy: abs(cx - 0.5) < 0.25,
+        "centre": lambda cx, cy: abs(cx - 0.5) < 0.25,
+    }
+    stats = []
+    for m in masks:
+        ys, xs = np.nonzero(m)
+        stats.append((xs.mean() / w, ys.mean() / h, m.mean()))
+    keep = [i for i, (cx, cy, _) in enumerate(stats) if all(sides[x](cx, cy) for x in where if x in sides)]
+    keep = keep or list(range(len(masks)))
+    sign = -1 if "background" in where else 1  # background: the smallest, the furthest away
+    best = max(keep, key=lambda i: (sign * stats[i][2] * (0.5 + float(scores[i]))))
+    return masks[best : best + 1]
+
+
+def _plural(query: str) -> bool:
+    head = _QUALIFIER.sub("", query).split()
+    return bool(_COUNT.match(query)) or (bool(head) and head[-1].endswith("s") and not head[-1].endswith("ss"))
+
+
+def _fill(mask: np.ndarray) -> np.ndarray:
+    """Fill holes: the outline should go around the object, not through it (a ladder's rungs)."""
+    contours, _ = cv2.findContours(mask.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    filled = np.zeros(mask.shape, np.uint8)
+    cv2.drawContours(filled, contours, -1, 1, thickness=cv2.FILLED)
+    return filled > 0
+
+
+class TextMasker:
+    def __init__(self, device: str = "cuda", dtype=torch.bfloat16):
+        from transformers import Sam3Model, Sam3Processor
+
+        self.processor = Sam3Processor.from_pretrained(SAM3_REPO)
+        self.model = Sam3Model.from_pretrained(SAM3_REPO, torch_dtype=dtype).to(device).eval()
+        self.dtype = dtype
+
+    @torch.no_grad()
+    def _segment(self, image: Image.Image, phrases: list[str], threshold: float):
+        inputs = self.processor(images=[image] * len(phrases), text=phrases, return_tensors="pt").to(self.model.device)
+        inputs["pixel_values"] = inputs["pixel_values"].to(self.dtype)
+        outputs = self.model(**inputs)
+        results = self.processor.post_process_instance_segmentation(
+            outputs, threshold=threshold, mask_threshold=0.5, target_sizes=inputs["original_sizes"].tolist()
+        )
+        return [(r["masks"].bool().cpu().numpy(), r["scores"].float().cpu().numpy()) for r in results]
+
+    def __call__(
+        self, image: Image.Image, text: str, threshold: float = 0.4, dilate: float = 0.012
+    ) -> tuple[Image.Image | None, list[tuple[str, int]]]:
+        """Mask (L, white = remove) for every phrase in `text`, or None when nothing is found.
+        Also returns (phrase, instances used) per phrase."""
+        image = image.convert("RGB")
+        phrases = phrases_from_instruction(text)
+        if not phrases:
+            return None, []
+        queries = [_split_where(p) for p in phrases]
+        results = self._segment(image, [q for q, _ in queries], threshold)
+        # A long relational description ("studio light in the foreground on the left") can find
+        # nothing: retry those with just the head noun phrase, a little more permissively.
+        retry = [i for i, (masks, _) in enumerate(results) if not len(masks)]
+        # Then with just the noun ("wooden chapel" -> "chapel"). A retry is a looser guess, so it
+        # keeps only its best instance unless the phrase asks for several.
+        for shorten in (lambda q: _QUALIFIER.sub("", q), lambda q: _QUALIFIER.sub("", q).split()[-1]):
+            retry = [i for i, (masks, _) in enumerate(results) if not len(masks)]
+            if not retry:
+                break
+            again = self._segment(image, [shorten(queries[i][0]) for i in retry], threshold * 0.75)
+            for i, (masks, scores) in zip(retry, again):
+                if len(masks) > 1 and not _plural(queries[i][0]):
+                    keep = int(scores.argmax())
+                    masks, scores = masks[keep : keep + 1], scores[keep : keep + 1]
+                results[i] = (masks, scores)
+        union = np.zeros((image.height, image.width), dtype=bool)
+        found = []
+        for (query, where), phrase, (masks, scores) in zip(queries, phrases, results):
+            # One instance when the phrase is singular and says where it is ("the phone on the right").
+            if len(masks) > 1 and where and not _COUNT.match(query):
+                masks = _pick(masks, scores, where)
+            found.append((phrase, int(len(masks))))
+            if len(masks):
+                union |= masks.any(0)
+        if not union.any():
+            return None, found
+        radius = max(2, round(dilate * min(image.size)))
+        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2 * radius + 1, 2 * radius + 1))
+        union = _fill(cv2.dilate(union.astype(np.uint8), kernel) > 0)
+        return Image.fromarray((union * 255).astype(np.uint8), "L"), found