# Exports a Mask2Former checkpoint from Hugging Face to ONNX for modules/Segment, then checks the ONNX
# model against PyTorch on a folder of images: the raw logits, and the panoptic segmentation made from
# them, segment by segment. Run once, on a developer machine; mcap only ships the .onnx it writes.
#
#   conda create -n mcap-seg python=3.11
#   conda run -n mcap-seg pip install torch --index-url https://download.pytorch.org/whl/cpu
#   conda run -n mcap-seg pip install transformers onnx onnxruntime pillow scipy timm
#   conda run -n mcap-seg python tools/export_mask2former.py --check FOLDER_OF_IMAGES [--reference OUT]
#
# --reference also writes PyTorch's own segmentation of each image, which bin/segment-probe -check
# compares modules/Segment's against.
#
# The model takes pixel_values [batch, 3, 384, 384] (RGB, resized to 384 x 384 without keeping the
# aspect ratio, as the checkpoint's processor does; scaled to 0..1, then ImageNet mean and std) and
# gives class_queries_logits [batch, 100, 134] (133 classes and "no object") and masks_queries_logits
# [batch, 100, 96, 96]. Making segments of those (the panoptic merge) is mcap's, in Jai.

import argparse, json, os, sys, time
import numpy as np
import torch
from PIL import Image
from transformers import AutoImageProcessor, Mask2FormerForUniversalSegmentation
from transformers.utils import logging

logging.set_verbosity_error()


class Logits(torch.nn.Module):
    """The model, from pixel values to the two logit tensors, as the ONNX graph has it."""

    def __init__(self, model):
        super().__init__()
        self.model = model

    def forward(self, pixel_values):
        out = self.model(pixel_values=pixel_values)
        return out.class_queries_logits, out.masks_queries_logits


def export(model, path, size, opset):
    wrapper = Logits(model).eval()
    dummy = torch.randn(1, 3, size, size)
    start = time.time()
    torch.onnx.export(
        wrapper, (dummy,), path,
        input_names=["pixel_values"],
        output_names=["class_queries_logits", "masks_queries_logits"],
        dynamic_axes={"pixel_values": {0: "batch"}, "class_queries_logits": {0: "batch"}, "masks_queries_logits": {0: "batch"}},
        opset_version=opset,
        do_constant_folding=True,
        dynamo=False,
    )
    print(f"exported {path} ({os.path.getsize(path) / 1e6:.0f} MB) in {time.time() - start:.0f} s")


class Outputs:
    """What post_process_panoptic_segmentation reads, from ONNX's arrays."""

    def __init__(self, class_logits, mask_logits):
        self.class_queries_logits = torch.from_numpy(class_logits)
        self.masks_queries_logits = torch.from_numpy(mask_logits)


def segments(processor, outputs, size):
    """[(label, mask)] of the panoptic segmentation, at the image's size."""
    result = processor.post_process_panoptic_segmentation(outputs, target_sizes=[size])[0]
    seg = result["segmentation"]
    return [(s["label_id"], (seg == s["id"]).numpy(), s["score"]) for s in result["segments_info"]]


def iou(a, b):
    union = np.logical_or(a, b).sum()
    return 1.0 if union == 0 else np.logical_and(a, b).sum() / union


def check(model, processor, path, folder, names, providers):
    import onnxruntime as ort
    session = ort.InferenceSession(path, providers=providers)
    print(f"onnxruntime {ort.__version__} on {session.get_providers()[0]}")
    files = sorted(f for f in os.listdir(folder) if f.lower().endswith((".jpg", ".jpeg", ".png")))
    worst_logit, ious, unmatched, torch_s, onnx_s = 0.0, [], 0, 0.0, 0.0
    for name in files:
        image = Image.open(os.path.join(folder, name)).convert("RGB")
        inputs = processor(images=image, return_tensors="pt")
        start = time.time()
        with torch.no_grad():
            reference = model(pixel_values=inputs["pixel_values"])
        torch_s += time.time() - start
        start = time.time()
        class_logits, mask_logits = session.run(None, {"pixel_values": inputs["pixel_values"].numpy()})
        onnx_s += time.time() - start

        diff = max(np.abs(class_logits - reference.class_queries_logits.numpy()).max(),
                   np.abs(mask_logits - reference.masks_queries_logits.numpy()).max())
        worst_logit = max(worst_logit, diff)

        size = image.size[::-1]
        ref = segments(processor, reference, size)
        got = segments(processor, Outputs(class_logits, mask_logits), size)
        # Each reference segment against the best ONNX segment of the same class.
        line = []
        for label, mask, score in ref:
            best = max((iou(mask, m) for l, m, _ in got if l == label), default=0.0)
            ious.append(best)
            if best < 0.5:
                unmatched += 1
            line.append(f"{names[label]} {best:.3f}")
        extra = len(got) - len(ref)
        print(f"  {name}: logits max diff {diff:.1e}; {len(ref)} segments{f' ({extra:+d} in ONNX)' if extra else ''}: " + ", ".join(line[:6]) + (" ..." if len(line) > 6 else ""))

    ious = np.array(ious)
    print(f"\n{len(files)} images, {len(ious)} segments")
    print(f"  logits: worst max abs diff {worst_logit:.2e}")
    print(f"  segment IoU, ONNX vs PyTorch: min {ious.min():.4f}, mean {ious.mean():.4f}, below 0.98: {(ious < 0.98).sum()}, unmatched: {unmatched}")
    print(f"  time per image: PyTorch {torch_s / len(files) * 1000:.0f} ms, ONNX {onnx_s / len(files) * 1000:.0f} ms (CPU, batch 1)")
    return ious.min() >= 0.98 and unmatched == 0


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--model", default="facebook/mask2former-swin-base-coco-panoptic")
    parser.add_argument("--out", default="modules/Segment/models")
    parser.add_argument("--opset", type=int, default=17)
    parser.add_argument("--check", help="a folder of images to check the export on")
    parser.add_argument("--skip-export", action="store_true", help="check an .onnx already written")
    parser.add_argument("--reference", help="write PyTorch's segmentation of the --check images here, for segment-probe -check")
    args = parser.parse_args()

    processor = AutoImageProcessor.from_pretrained(args.model)
    model = Mask2FormerForUniversalSegmentation.from_pretrained(args.model).eval()
    size = processor.size["height"]
    assert processor.size["width"] == size, "the export assumes a square input"
    names = model.config.id2label

    os.makedirs(args.out, exist_ok=True)
    base = args.model.split("/")[-1]
    path = os.path.join(args.out, base + ".onnx")
    if not args.skip_export:
        export(model, path, size, args.opset)

    # The class list, for models.config: index, whether it is a countable thing, the model's name, and
    # the name people see and search (has:): COCO's "wall-other-merged" is "wall".
    with open(os.path.join(args.out, base + ".classes.txt"), "w", encoding="utf-8") as f:
        f.write("@version 1\n")
        for i in range(len(names)):
            thing = "thing" if i < 80 else "stuff"   # COCO panoptic: 80 things, then 53 stuff
            f.write(f'class: {i} {thing} "{names[i]}" "{display_name(names[i])}"\n')
    print(f"wrote {base}.classes.txt ({len(names)} classes)")

    if args.reference:
        if not args.check:
            sys.exit("--reference needs --check FOLDER, the images to make it of")
        write_reference(model, processor, args.check, args.reference)

    if args.check:
        ok = check(model, processor, path, args.check, names, ["CPUExecutionProvider"])
        print("PASS" if ok else "FAIL")
        sys.exit(0 if ok else 1)


def display_name(name):
    """COCO panoptic's class names for people: no -merged, -other or -stuff, spaces for dashes."""
    for suffix in ("-merged", "-other", "-stuff"):
        name = name.replace(suffix, "")
    return name.replace("-", " ").strip()


MASK_MAX_SIDE = 1024   # mcap's mask map: the image scaled so its longest side is at most this


def map_size(width, height):
    """The mask map's size for an image, as modules/Segment works it out (rounding half up)."""
    scale = min(1.0, MASK_MAX_SIDE / max(width, height))
    return max(1, int(width * scale + 0.5)), max(1, int(height * scale + 0.5))


def write_reference(model, processor, folder, out):
    """PyTorch's panoptic segmentation of each image, at the mask map's size, for segment-probe -check:
    NAME.ref (b"M2FR", u16 width, u16 height, then a u8 segment id per pixel, row by row) and
    NAME.ref.txt (a line per segment: id, class index, score)."""
    import struct
    from PIL import ImageOps
    os.makedirs(out, exist_ok=True)
    files = sorted(f for f in os.listdir(folder) if f.lower().endswith((".jpg", ".jpeg", ".png")))
    for name in files:
        image = ImageOps.exif_transpose(Image.open(os.path.join(folder, name))).convert("RGB")   # turned, as mcap shows it
        w, h = map_size(*image.size)
        inputs = processor(images=image, return_tensors="pt")
        with torch.no_grad():
            outputs = model(pixel_values=inputs["pixel_values"])
        result = processor.post_process_panoptic_segmentation(outputs, target_sizes=[(h, w)])[0]
        seg = result["segmentation"].numpy().astype(np.int64)
        seg[seg < 0] = 0   # no segment
        assert seg.max() < 256
        with open(os.path.join(out, name + ".ref"), "wb") as f:
            f.write(b"M2FR" + struct.pack("<HH", w, h) + seg.astype(np.uint8).tobytes())
        with open(os.path.join(out, name + ".ref.txt"), "w", encoding="utf-8") as f:
            for s in result["segments_info"]:
                f.write(f"{s['id']} {s['label_id']} {s['score']:.6f}\n")
    print(f"wrote references for {len(files)} images to {out}")


if __name__ == "__main__":
    main()