Mask2Former and universal segmentation

pip install numpy scipy

import numpy as np
from scipy.optimize import linear_sum_assignment

H, W, D, Q = 8, 14, 4, 6                 # image, embedding size, number of queries
CLASSES = ["sky", "road", "car", "no object"]
NO_OBJECT = 3

# ---- the pixel decoder's output: one embedding vector per pixel ----
region = np.zeros((H, W), int)           # 0 sky, 1 road, 2 car A, 3 car B
region[4:] = 1
region[4:7, 1:5] = 2
region[4:7, 8:12] = 3
pixel_emb = np.eye(D)[region].transpose(2, 0, 1)      # (D, H, W)

# ---- the transformer decoder's output: one row per query ----
# each query holds a mask embedding and a distribution over classes
mask_emb = np.array([
    [1, 0, 0, 0],    # q0 -> sky
    [0, 1, 0, 0],    # q1 -> road
    [0, 0, 1, 0],    # q2 -> car A
    [0, 0, 0, 1],    # q3 -> car B
    [0, 0, 1, 0],    # q4 -> car A again, a duplicate
    [0, 0, 0, 0],    # q5 -> nothing in particular
], float)
class_prob = np.array([
    [0.97, 0.01, 0.01, 0.01],
    [0.01, 0.96, 0.02, 0.01],
    [0.02, 0.02, 0.94, 0.02],
    [0.02, 0.03, 0.88, 0.07],
    [0.03, 0.05, 0.52, 0.40],
    [0.05, 0.05, 0.05, 0.85],
])

# ---- masks come from a dot product, nothing more ----
mask_logits = np.einsum("qd,dhw->qhw", mask_emb, pixel_emb) * 8 - 4
mask_prob = 1 / (1 + np.exp(-mask_logits))
print("one dot product turns", Q, "query vectors into", Q, "masks:", mask_logits.shape)

GLYPH = " .:-=+*#@"
def strip(m, row=5):
    return "".join(GLYPH[int(v * 8)] for v in m[row])
print("\nrows 1 and 5 of each query mask (@ = confident foreground):")
for q in range(Q):
    print(f"   q{q} |{strip(mask_prob[q], 1)}| |{strip(mask_prob[q], 5)}|  "
          f"best class: {CLASSES[class_prob[q].argmax()]}")

# ---- ONE set of outputs, THREE different tasks ----
def semantic(mask_prob, class_prob):
    thing = class_prob[:, :NO_OBJECT]                       # drop the no-object column
    return np.einsum("qc,qhw->chw", thing, mask_prob).argmax(0)

THINGS = {2}                                                 # only cars are countable here

def instances(mask_prob, class_prob, thr=0.6):
    out = []
    for q in range(len(class_prob)):
        c = class_prob[q].argmax()
        if c in THINGS and class_prob[q, c] > thr:
            out.append((q, CLASSES[c], class_prob[q, c], (mask_prob[q] > 0.5).sum()))
    return out

def panoptic(mask_prob, class_prob):
    score = class_prob[:, :NO_OBJECT].max(1)[:, None, None] * mask_prob
    winner = score.argmax(0)                                 # one query wins each pixel
    keep = score.max(0) > 0.5
    return np.where(keep, winner, -1)

sem = semantic(mask_prob, class_prob)
print("\nsemantic map from those queries (s sky, r road, c car):")
for row in sem:
    print("   " + "".join("src"[v] for v in row))

print("\ninstances: countable classes only, score above 0.6:")
for q, name, s, area in instances(mask_prob, class_prob):
    print(f"   q{q}  {name:5s} score {s:.2f}  {area:>3} pixels")

pan = panoptic(mask_prob, class_prob)
print("\npanoptic map, the winning query id per pixel (. = nothing above 0.5):")
for row in pan:
    print("   " + "".join("." if v < 0 else str(v) for v in row))
print("   note q4 wins nothing: q2 predicts the same mask with a higher class score")

# ---- training: which query is responsible for which ground-truth mask? ----
gt_masks = np.stack([region == 0, region == 1, region == 2, region == 3]).astype(float)
gt_class = [0, 1, 2, 2]

def dice_cost(p, g):
    return 1 - (2 * (p * g).sum() + 1) / (p.sum() + g.sum() + 1)

cost = np.zeros((Q, len(gt_masks)))
for q in range(Q):
    for g in range(len(gt_masks)):
        cost[q, g] = dice_cost(mask_prob[q], gt_masks[g]) - class_prob[q, gt_class[g]]

print("\nmatching cost, queries down, ground-truth masks across:")
print("        " + "".join(f"gt{g}     " for g in range(len(gt_masks))))
for q in range(Q):
    print(f"   q{q}  " + "".join(f"{cost[q, g]:+7.3f} " for g in range(len(gt_masks))))

rows, cols = linear_sum_assignment(cost)
print("\nHungarian assignment (each ground-truth mask gets exactly one query):")
for q, g in zip(rows, cols):
    print(f"   q{q} -> gt{g} ({CLASSES[gt_class[g]]}), cost {cost[q, g]:+.3f}")
print("   unmatched queries are trained towards 'no object':",
      [f"q{q}" for q in range(Q) if q not in rows])

from transformers import AutoImageProcessor, Mask2FormerForUniversalSegmentation
from PIL import Image
import torch

name = "facebook/mask2former-swin-tiny-coco-instance"     # about 190 MB of weights
processor = AutoImageProcessor.from_pretrained(name)
model = Mask2FormerForUniversalSegmentation.from_pretrained(name)

image = Image.open("your_photo.jpg").convert("RGB")
inputs = processor(image, return_tensors="pt")
with torch.no_grad():
    outputs = model(**inputs)

# the two tensors the toy script above imitates
print(outputs.class_queries_logits.shape)   # (batch, num_queries, num_classes + 1)
print(outputs.masks_queries_logits.shape)   # (batch, num_queries, h, w)

result = processor.post_process_instance_segmentation(
    outputs, target_sizes=[(image.height, image.width)]
)[0]
print(result["segmentation"].shape)