Panoptic segmentation

pip install numpy

import numpy as np

H, W = 10, 16
SKY, ROAD, TREE, CAR, PERSON = 0, 1, 2, 3, 4
NAME = {SKY: "sky", ROAD: "road", TREE: "tree", CAR: "car", PERSON: "person"}
STUFF = (SKY, ROAD, TREE)

# ---------------- ground truth panoptic map ----------------
gt_seg = np.zeros((H, W), int)                 # one segment id per pixel
gt_cls = {1: SKY, 2: ROAD, 3: CAR, 4: CAR, 5: PERSON}
gt_seg[:4] = 1
gt_seg[4:] = 2
gt_seg[4:8, 1:6] = 3
gt_seg[4:8, 8:12] = 4
gt_seg[5:9, 13:16] = 5

# ---------------- what the two heads predicted ----------------
sem = np.zeros((5, H, W))                      # semantic head score per class per pixel
sem[SKY, :4] = 3.0
sem[ROAD, 4:] = 3.0
sem[ROAD, 3, :] = 3.5                          # the horizon is one row too low
sem[TREE, 9, 0:3] = 3.5                        # three stray tree pixels

def box(r0, r1, c0, c1):
    m = np.zeros((H, W), bool); m[r0:r1, c0:c1] = True; return m

instances = [                                   # (score, class, mask) from the instance head
    (0.95, CAR,    box(4, 8, 1, 6)),            # car 1, exactly right
    (0.88, CAR,    box(4, 8, 8, 13)),           # car 2, one column too wide
    (0.71, CAR,    box(4, 8, 4, 10)),           # a third car nobody asked for
    (0.55, PERSON, box(5, 9, 12, 16)),          # person, one column too wide
]

# ---------------- merge them into one map ----------------
def merge(instances, sem, overlap_keep=0.5, min_stuff_area=8):
    pan = np.zeros((H, W), int)                 # 0 means still unclaimed
    cls_of, next_id = {}, 1
    for score, cls, mask in sorted(instances, key=lambda t: -t[0]):
        free = mask & (pan == 0)                # pixels a higher-scoring mask has not taken
        kept = free.sum() / mask.sum()
        if kept < overlap_keep:
            print(f"   dropped {NAME[cls]:6s} score {score:.2f}: only {kept:.0%} of it was still free")
            continue
        pan[free] = next_id; cls_of[next_id] = cls; next_id += 1
        print(f"   kept    {NAME[cls]:6s} score {score:.2f}: {free.sum():>3} pixels ({kept:.0%} of its mask)")
    for c in STUFF:                             # all pixels of one stuff class = ONE segment
        region = (sem.argmax(0) == c) & (pan == 0)
        if region.sum() >= min_stuff_area:
            pan[region] = next_id; cls_of[next_id] = c; next_id += 1
        elif region.sum():
            print(f"   dropped {NAME[c]:6s}         : {region.sum()} pixels is under the area floor")
    return pan, cls_of

print("merge order is score, highest first:")
pred_seg, pred_cls = merge(instances, sem)

GLYPH = {SKY: ".", ROAD: "-", TREE: "T", CAR: "C", PERSON: "P"}
def show(seg, cls_of, title):
    print("\n" + title)
    for row in seg:
        print("   " + " ".join(GLYPH[cls_of[v]] if v else "?" for v in row))

show(gt_seg, gt_cls, "ground truth")
show(pred_seg, pred_cls, "merged prediction  (? = void, claimed by nothing)")

# ---------------- panoptic quality ----------------
def pq(gt_seg, gt_cls, pr_seg, pr_cls):
    out = {}
    for c in sorted(set(gt_cls.values()) | set(pr_cls.values())):
        g = [i for i, k in gt_cls.items() if k == c]
        p = [i for i, k in pr_cls.items() if k == c]
        hit_g, hit_p, ious = set(), set(), []
        for gi in g:
            for pi in p:
                inter = np.sum((gt_seg == gi) & (pr_seg == pi))
                if not inter:
                    continue
                iou = inter / np.sum((gt_seg == gi) | (pr_seg == pi))
                if iou > 0.5:                   # above one half the match can only be unique
                    hit_g.add(gi); hit_p.add(pi); ious.append(iou)
        tp, fp, fn = len(ious), len(p) - len(hit_p), len(g) - len(hit_g)
        sq = float(np.mean(ious)) if tp else 0.0
        rq = tp / (tp + 0.5 * fp + 0.5 * fn)
        out[c] = (tp, fp, fn, sq, rq, sq * rq)
    return out

print(f"\n{'class':8s} {'TP':>3} {'FP':>3} {'FN':>3} {'SQ':>6} {'RQ':>6} {'PQ':>6}")
rows = pq(gt_seg, gt_cls, pred_seg, pred_cls)
for c, (tp, fp, fn, sq, rq, p) in rows.items():
    print(f"{NAME[c]:8s} {tp:>3} {fp:>3} {fn:>3} {sq:>6.3f} {rq:>6.3f} {p:>6.3f}")
print(f"{'PQ':8s} {'':>3} {'':>3} {'':>3} {'':>6} {'':>6} "
      f"{np.mean([r[5] for r in rows.values()]):>6.3f}")

# what a missed car does, with everything else identical
second_car = pred_seg[4, 9]                      # the segment id sitting on the second car
worse_seg = np.where(pred_seg == second_car, 0, pred_seg)
worse_cls = {i: c for i, c in pred_cls.items() if i != second_car}
rows2 = pq(gt_seg, gt_cls, worse_seg, worse_cls)
print(f"\ndelete the second car and only recognition quality moves:")
print(f"   car  SQ {rows[CAR][3]:.3f} -> {rows2[CAR][3]:.3f}   "
      f"RQ {rows[CAR][4]:.3f} -> {rows2[CAR][4]:.3f}   PQ {rows[CAR][5]:.3f} -> {rows2[CAR][5]:.3f}")