Hugging Face's logo

Panorama-grounding
/
PANORAMA-2B

Image-Text-to-Text
Transformers
Safetensors
panorama_chat
feature-extraction
panoptic-grounded-captioning
grounded-captioning
segmentation
multimodal
conversational
custom_code
Model card Files
xet
 Community
PANORAMA-2B / visualize.py
HuggingSara's picture
HuggingSara
Initial commit
b9d0394
Raw History Blame Contribute Delete
10.5 kB
#!/usr/bin/env python3
"""Visualise PANORAMA's masks: one panel per phrase, plus a combined overlay.
python visualize.py --image test.png --pred ./out # after predict_panocaps.py
python visualize.py --image ./my_images/ --pred ./out # a whole folder
For each image it writes a folder <stem>_viz/ containing:
00_<phrase>.png ... the image with ONLY that phrase's mask(s) lit; everything else dimmed
overlay.png all masks tinted per phrase + a colour->phrase legend
caption.txt the caption verbatim, with the <p>...</p>[SEG] markers
caption_clean.txt the same caption with the markers stripped
Per-phrase panels dim the background instead of tinting the mask, so you judge the mask against the
object's real colours (a 50% tint turns a red mug green). The overlay is for coverage; the panels
are for correctness.
Pillow + numpy only — no GPU, no extra libraries.
"""
import argparse
import json
import os
import re
import numpy as np
from PIL import Image, ImageDraw, ImageFont
# Colour is per PHRASE, not per mask: one phrase can own several masks ("three chairs"), and they
# belong together. masks_per_phrase in the json gives the split.
# Okabe-Ito (colour-blind friendly), extended; cycles if a caption names more regions than colours.
# Names are the canonical Okabe-Ito ones for the first 8; the rest are our tints/extensions.
# Index == the phrase number printed in the per-phrase panels, so "phrase 03" is always yellow.
PALETTE_NAMES = [
"orange", # 0 (230, 159, 0) Okabe-Ito
"sky-blue", # 1 ( 86, 180, 233) Okabe-Ito
"bluish-green", # 2 ( 0, 158, 115) Okabe-Ito
"yellow", # 3 (240, 228, 66) Okabe-Ito
"blue", # 4 ( 0, 114, 178) Okabe-Ito
"vermillion", # 5 (213, 94, 0) Okabe-Ito
"reddish-purple", # 6 (204, 121, 167) Okabe-Ito
"grey", # 7 (153, 153, 153) Okabe-Ito (black slot, lightened)
"salmon", # 8 (255, 128, 128) extension
"mint", # 9 (128, 255, 128) extension
"periwinkle", # 10 (128, 128, 255) extension
"apricot", # 11 (255, 200, 100) extension
"amethyst", # 12 (170, 110, 200) extension
"teal", # 13 (100, 200, 190) extension
"raspberry", # 14 (200, 90, 140) extension
]
PALETTE = [
(230, 159, 0), (86, 180, 233), (0, 158, 115), (240, 228, 66), (0, 114, 178),
(213, 94, 0), (204, 121, 167), (153, 153, 153), (255, 128, 128), (128, 255, 128),
(128, 128, 255), (255, 200, 100), (170, 110, 200), (100, 200, 190), (200, 90, 140),
]
assert len(PALETTE_NAMES) == len(PALETTE), "palette names and colours out of sync"
def colour_name(i):
"""Name of the colour phrase `i` is drawn in (wraps like PALETTE does)."""
return PALETTE_NAMES[i % len(PALETTE_NAMES)]
OVERLAY_ALPHA = 0.45 # combined view: strong enough to read coverage
DIM = 0.28 # per-phrase panels: how much of the image survives outside the mask
def phrase_of_mask(counts, n_masks):
"""mask index -> phrase index, from masks_per_phrase (falls back to 1:1)."""
if not counts:
return list(range(n_masks))
out = []
for pi, c in enumerate(counts):
out.extend([pi] * int(c))
if len(out) < n_masks: # trust the masks, not the counts
out.extend([len(counts) - 1] * (n_masks - len(out)))
return out[:n_masks]
def outline(mask):
"""1px boundary of a boolean mask."""
e = np.zeros_like(mask)
e[:-1, :] |= mask[:-1, :] ^ mask[1:, :]
e[:, :-1] |= mask[:, :-1] ^ mask[:, 1:]
return e
def _font(size=15):
for p in ('DejaVuSans.ttf', '/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf'):
try:
return ImageFont.truetype(p, size)
except Exception:
pass
return ImageFont.load_default()
def panel(img_arr, masks, colour, label):
"""One phrase: its mask(s) at full brightness, everything else dimmed."""
m = np.zeros(img_arr.shape[:2], bool)
for k in masks:
m |= k
out = (img_arr * DIM).astype(np.float32) # dim everything
out[m] = img_arr[m] # restore the mask region
out[outline(m)] = colour # trace it
im = Image.fromarray(out.astype(np.uint8))
d = ImageDraw.Draw(im)
f = _font(16)
px = int(m.sum())
txt = f'{label} — {px:,} px' + (f' ({len(masks)} masks)' if len(masks) > 1 else '')
x0, y0, x1, y1 = d.textbbox((0, 0), txt, font=f)
d.rectangle([6, 6, 14 + (x1 - x0), 14 + (y1 - y0)], fill=(0, 0, 0))
d.rectangle([6, 6, 12, 14 + (y1 - y0)], fill=colour) # colour tab
d.text((16, 10), txt, fill=(255, 255, 255), font=f)
return im
def _wrap(draw, text, font, width):
"""Greedy word-wrap to `width` px."""
lines, cur = [], ''
for w in text.split():
t = f'{cur} {w}'.strip()
if draw.textlength(t, font=font) <= width or not cur:
cur = t
else:
lines.append(cur)
cur = w
if cur:
lines.append(cur)
return lines
def overlay(img_arr, masks, phrases, m2p, caption=''):
arr = img_arr.astype(np.float32)
for i, m in enumerate(masks):
if not m.any():
continue
c = np.array(PALETTE[m2p[i] % len(PALETTE)], np.float32)
arr[m] = (1 - OVERLAY_ALPHA) * arr[m] + OVERLAY_ALPHA * c
for i, m in enumerate(masks):
if m.any():
arr[outline(m)] = PALETTE[m2p[i] % len(PALETTE)]
im = Image.fromarray(arr.astype(np.uint8))
f, fc = _font(14), _font(15)
probe = ImageDraw.Draw(im)
# caption block: the [SEG] markers are noise on a picture that already shows the masks
cap = re.sub(r'\s*\[SEG\]', '', caption).replace('<p>', '').replace('</p>', '')
cap = re.sub(r'\s+', ' ', cap).strip()
cap_lines = _wrap(probe, cap, fc, im.width - 24) if cap else []
cap_h = (len(cap_lines) * 20 + 16) if cap_lines else 0
lh = 20
strip_h = cap_h + max(lh * len(phrases) + 12, 24)
strip = Image.new('RGB', (im.width, strip_h), (255, 255, 255))
sd = ImageDraw.Draw(strip)
y = 8
for ln in cap_lines:
sd.text((12, y), ln, fill=(0, 0, 0), font=fc)
y += 20
if cap_lines:
y += 8
sd.line([12, y - 5, im.width - 12, y - 5], fill=(210, 210, 210))
for pi, ph in enumerate(phrases):
sd.rectangle([6, y + 3, 20, y + 15], fill=PALETTE[pi % len(PALETTE)], outline=(0, 0, 0))
n = m2p.count(pi)
sd.text((28, y), f'{pi:02d} {ph.strip()}' + (f' ({n} masks)' if n > 1 else ''),
fill=(0, 0, 0), font=f)
y += lh
out = Image.new('RGB', (im.width, im.height + strip_h), (255, 255, 255))
out.paste(im, (0, 0))
out.paste(strip, (0, im.height))
return out
def slug(s, n=40):
s = re.sub(r'[^a-z0-9]+', '_', s.strip().lower()).strip('_')
return (s[:n] or 'phrase')
def run(image_path, pred_dir):
stem = os.path.splitext(os.path.basename(image_path))[0]
jf, mf = os.path.join(pred_dir, f'{stem}.json'), os.path.join(pred_dir, f'{stem}_masks.npz')
if not (os.path.isfile(jf) and os.path.isfile(mf)):
miss = [os.path.basename(f) for f in (jf, mf) if not os.path.isfile(f)]
print(f' skip {stem}: missing {", ".join(miss)} in {pred_dir} — run predict_panocaps.py '
f'with --out {pred_dir} first')
return
d = json.load(open(jf))
masks = np.load(mf)['masks'].astype(bool)
phrases = d['phrases']
img = np.asarray(Image.open(image_path).convert('RGB'))
if masks.size and masks.shape[1:] != img.shape[:2]:
print(f' !! {stem}: masks are {masks.shape[1:]} but the image is {img.shape[:2]}')
return
m2p = phrase_of_mask(d.get('masks_per_phrase'), len(masks))
outdir = os.path.join(pred_dir, f'{stem}_viz')
os.makedirs(outdir, exist_ok=True)
# The caption, next to the panels, in both forms:
# caption.txt — verbatim, with the <p>...</p>[SEG] markers. This is the alignment
# record: the i-th <p>...</p> is the i-th phrase, hence panel 00, 01, ...
# caption_clean.txt — the markers stripped, i.e. the prose on its own, for reading or for
# feeding somewhere that does not want the markup.
raw = d['caption']
clean = re.sub(r'\s*\[SEG\]', '', raw).replace('<p>', '').replace('</p>', '')
clean = re.sub(r'\s+', ' ', clean)
# the model writes "<p> red mug </p> [SEG]." — the spaces live INSIDE the tags, so removing them
# strands a space before the punctuation ("a red mug ."). Close it up.
clean = re.sub(r'\s+([,.;:!?%)\]])', r'\1', clean)
clean = re.sub(r'([(\[])\s+', r'\1', clean).strip()
open(os.path.join(outdir, 'caption.txt'), 'w').write(raw.strip() + '\n')
open(os.path.join(outdir, 'caption_clean.txt'), 'w').write(clean + '\n')
for pi, ph in enumerate(phrases):
mine = [masks[i] for i in range(len(masks)) if m2p[i] == pi]
if not mine:
continue
panel(img, mine, PALETTE[pi % len(PALETTE)], f'{pi:02d} {ph.strip()}').save(
os.path.join(outdir, f'{pi:02d}_{slug(ph)}.png'))
overlay(img, masks, phrases, m2p).save(os.path.join(outdir, 'overlay.png'))
print(f'{stem}: {len(masks)} mask(s), {len(phrases)} phrase(s) -> {outdir}/ '
f'({len(phrases)} panels + overlay.png + caption.txt/caption_clean.txt)')
def main():
ap = argparse.ArgumentParser()
ap.add_argument('--image', required=True, help='the image file, or the directory you predicted on')
ap.add_argument('--pred', default='./out',
help='the --out DIRECTORY predict_panocaps.py wrote (not a file)')
a = ap.parse_args()
if os.path.isfile(a.pred): # easy slip: passing <stem>_masks.npz
a.pred = os.path.dirname(os.path.abspath(a.pred))
print(f'--pred is a file; using its directory: {a.pred}')
elif not os.path.isdir(a.pred):
raise SystemExit(f'--pred: no such directory: {a.pred}')
imgs = ([os.path.join(a.image, f) for f in sorted(os.listdir(a.image))
if f.lower().endswith(('.jpg', '.jpeg', '.png', '.webp', '.bmp'))]
if os.path.isdir(a.image) else [a.image])
for p in imgs:
run(p, a.pred)
if __name__ == '__main__':
main()