#!/usr/bin/env python3
"""PANORAMA — panoptic captioning + segmentation on your own images.

The model writes a caption in which every object and background region is named inside
<p>...</p> and immediately followed by a [SEG] token; each [SEG] returns a binary mask.

    python predict_panocaps.py --model . --image photo.jpg --out ./out

"""
import argparse
import json
import os
import re

import numpy as np
import torch
from PIL import Image
from transformers import AutoModel, AutoProcessor, AutoTokenizer

# The EXACT prompt this model is evaluated with. Do not paraphrase it: near-paraphrases
# ("describe the image in detail", ...) were trained toward a SPARSER output style, so
# rewording quietly thins the coverage.
PANOCAPS_PROMPT = ('Write a caption covering all objects and background regions in the image, '
                   'inserting the matching segmentation mask after each phrase.')

PHRASE_RE = re.compile(r'<p>(.*?)</p>', re.S)


def load(model_path):
    model = AutoModel.from_pretrained(
        model_path,
        dtype=torch.bfloat16,         # REQUIRED: the model is trained and evaluated in bf16
        low_cpu_mem_usage=True,
        trust_remote_code=True,       # the architecture ships with the weights
    ).eval().cuda()
    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
    processor = AutoProcessor.from_pretrained(model_path, trust_remote_code=True)
    return model, tokenizer, processor


@torch.no_grad()
def predict(model, tokenizer, processor, image):
    out = model.predict_forward(
        image=image,
        text='<image>\n' + PANOCAPS_PROMPT,
        tokenizer=tokenizer,
        processor=processor,
    )
    caption = out['prediction'].replace('<|im_end|>', '').strip()

    # Prefer the per-INSTANCE masks: one phrase can name several things ("three dogs"), and the
    # model selects a set per phrase. prediction_masks gives one union mask per [SEG] instead.
    inst = out.get('prediction_instance_masks')
    counts = out.get('prediction_phrase_seg_counts')
    src = inst if (inst is not None and counts is not None) else out.get('prediction_masks')
    masks = (np.stack([np.asarray(m) for m in src])[:, 0] if src is not None and len(src)
             else np.zeros((0, image.size[1], image.size[0]), dtype=bool))
    return caption, masks, counts


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--model', required=True, help='path to (or HF id of) the model')
    ap.add_argument('--image', required=True, help='an image file, or a directory of images')
    ap.add_argument('--out', default='./out')
    a = ap.parse_args()

    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])
    assert imgs, (f'no images found directly in {a.image} '
                  f'(jpg/jpeg/png/webp/bmp; not recursive — subfolders are ignored)')
    # outputs are keyed by STEM, so cat.jpg and cat.png would both write cat.json/cat_masks.npz and
    # the second would clobber the first. Fail loudly instead of silently losing a result.
    seen = {}
    for p in imgs:
        seen.setdefault(os.path.splitext(os.path.basename(p))[0], []).append(os.path.basename(p))
    clash = {k: v for k, v in seen.items() if len(v) > 1}
    assert not clash, (f'filename clash — these share a stem and would overwrite each other: {clash}. '
                       f'Rename so each image has a unique name (without extension).')
    os.makedirs(a.out, exist_ok=True)

    model, tokenizer, processor = load(a.model)
    print(f'loaded. {len(imgs)} image(s)\n')

    for p in imgs:
        image = Image.open(p).convert('RGB')
        caption, masks, counts = predict(model, tokenizer, processor, image)
        stem = os.path.splitext(os.path.basename(p))[0]

        np.savez_compressed(os.path.join(a.out, f'{stem}_masks.npz'), masks=masks)
        json.dump({'image': p, 'caption': caption,
                   'phrases': PHRASE_RE.findall(caption),
                   'n_masks': int(len(masks)),
                   'masks_per_phrase': counts},
                  open(os.path.join(a.out, f'{stem}.json'), 'w'), indent=2)

        print(f'--- {os.path.basename(p)}')
        print(f'    {caption}')
        print(f'    {len(masks)} mask(s) -> {stem}_masks.npz  (bool, N x H x W, image resolution)\n')


if __name__ == '__main__':
    main()