File size: 5,751 Bytes
4f08932 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 | #!/usr/bin/env python3
# coding=utf-8
"""ConCor-1 inference example: an image + its text in, correspondences out.
Given an image and a paired text, ConCor-1 predicts the full set of
correspondences between visually referential text spans and instance-level
image segments — the text spans are *not* given as queries; the model decides
which parts of the text are grounded.
Usage (the bundled `example.png` with its COCONut-PanCap caption):
python example_inference.py --image example.png \
--text "This image depicts a close-up of a brown bear in a natural outdoor setting. \
The background consists of lush green grass. In the foreground, a large brown bear is \
positioned centrally."
# a category list works just as well as a caption
python example_inference.py --image example.png --text "bear . grass . tree . person"
# write a mask overlay next to the printed correspondences
python example_inference.py --image example.png --text "..." --output overlay.png
"""
from __future__ import annotations
import argparse
from pathlib import Path
from typing import List
import numpy as np
import torch
from PIL import Image, ImageDraw
from transformers import AutoModel, AutoProcessor
# Visually distinct overlay colours.
COLORS = [
(239, 64, 64), (64, 204, 64), (64, 115, 255), (255, 204, 0), (255, 115, 0),
(217, 51, 217), (0, 204, 217), (153, 89, 13), (102, 255, 102), (140, 26, 255),
(255, 153, 204), (0, 140, 0), (179, 179, 0), (0, 26, 153), (204, 140, 51),
(128, 128, 128), (255, 0, 128), (0, 255, 140), (140, 0, 0), (0, 140, 140),
]
def overlay_masks(image: Image.Image, correspondences: List[dict], alpha: float = 0.5) -> Image.Image:
"""Blend each correspondence's mask over the image and label it with its phrases."""
canvas = np.array(image.convert("RGB"), dtype=np.float32)
for index, correspondence in enumerate(correspondences):
mask = correspondence.get("mask")
if mask is None or not mask.any():
continue
color = np.array(COLORS[index % len(COLORS)], dtype=np.float32)
canvas[mask] = (1.0 - alpha) * canvas[mask] + alpha * color
overlaid = Image.fromarray(canvas.astype(np.uint8))
draw = ImageDraw.Draw(overlaid)
for index, correspondence in enumerate(correspondences):
mask = correspondence.get("mask")
if mask is None or not mask.any():
continue
rows, columns = np.nonzero(mask)
label = " / ".join(correspondence["text_phrases"]) or "(no text span)"
label = f"{label} {correspondence['presence_score']:.2f}"
anchor = (int(columns.min()), max(int(rows.min()) - 12, 0))
draw.text(anchor, label, fill=COLORS[index % len(COLORS)])
return overlaid
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
parser.add_argument("--model", default=str(Path(__file__).parent), help="model repo id or local path")
parser.add_argument("--image", required=True, type=Path)
parser.add_argument("--text", required=True, help="the text to ground (caption, category list, referring expression)")
parser.add_argument("--output", type=Path, default=None, help="write a mask overlay here (PNG)")
parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
parser.add_argument(
"--attn_implementation",
default="flash_attention_2",
choices=["flash_attention_2", "sdpa", "eager"],
help="flash_attention_2 reproduces the paper's numbers exactly; sdpa needs no extra dependency",
)
parser.add_argument("--presence_threshold", type=float, default=None, help="default 0.1")
parser.add_argument("--text_threshold", type=float, default=None, help="default 0.45")
parser.add_argument("--image_threshold", type=float, default=None, help="default 0.45")
parser.add_argument("--nms_iou_threshold", type=float, default=None, help="default 0.5")
args = parser.parse_args()
processor = AutoProcessor.from_pretrained(args.model, trust_remote_code=True)
model = AutoModel.from_pretrained(
args.model,
trust_remote_code=True,
dtype=torch.bfloat16,
attn_implementation=args.attn_implementation,
).to(args.device).eval()
image = Image.open(args.image).convert("RGB")
inputs = processor(images=image, text=args.text, return_tensors="pt").to(args.device)
with torch.inference_mode(), torch.autocast(args.device, dtype=torch.bfloat16):
outputs = model(**inputs)
correspondences = processor.post_process_correspondences(
outputs,
text=args.text,
target_sizes=[(image.height, image.width)],
presence_threshold=args.presence_threshold,
text_threshold=args.text_threshold,
image_threshold=args.image_threshold,
nms_iou_threshold=args.nms_iou_threshold,
)[0]
print(f"\nimage: {args.image} ({image.width}x{image.height})")
print(f"text: {args.text}")
print(f"\n{len(correspondences)} correspondence(s):")
for correspondence in correspondences:
mask = correspondence["mask"]
phrases = " / ".join(correspondence["text_phrases"]) or "(no text span)"
area = 100.0 * mask.mean() if mask is not None else 0.0
print(
f" presence={correspondence['presence_score']:.3f} "
f"bridge={correspondence['bridge_index']:3d} "
f"mask={area:5.2f}% of image spans={correspondence['text_spans']} | {phrases}"
)
if args.output is not None:
overlay_masks(image, correspondences).save(args.output)
print(f"\nwrote {args.output}")
if __name__ == "__main__":
main()
|