# Converts Mask2Former (facebook/mask2former-swin-base-coco-panoptic) to a Core ML model for
# modules/Segment on macOS, then checks it against the ONNX model on a folder of images, segment by
# segment. Run once, on a Mac, after tools/export_mask2former.py; mcap only uses the .mlpackage it writes.
#
#   python3 -m venv ~/mcap-coreml
#   ~/mcap-coreml/bin/pip install torch==2.7.0 transformers coremltools==9.0 onnxruntime pillow scipy
#   ~/mcap-coreml/bin/pip install --no-deps timm
#   ~/mcap-coreml/bin/python tools/export_mask2former_coreml.py --check FOLDER_OF_IMAGES
#
# The model as tools/export_mask2former.py exports it (pixel_values [1, 3, 384, 384] to
# class_queries_logits [1, 100, 134] and masks_queries_logits [1, 100, 96, 96]), as an ML Program in
# float32, one image a run (a batch of 4 did not convert). Three things coremltools 9.0 cannot take
# are worked around for the conversion only, none changing what the model computes:
#   - it casts a one-element tensor with int(array), which NumPy 2 refuses;
#   - its meshgrid refuses the inputs Mask2Former's traced linspaces come out as;
#   - Core ML takes tensors of at most 5 dimensions, and the multi-scale deformable attention makes 6
#     (batch, queries, heads, levels, points, xy): it is run a level at a time instead, which gives
#     the same output in PyTorch (checked below) with at most 5.
# Measured on an M5 Pro: 42 ms an image on the GPU against 166 ms for onnxruntime on the CPU, with
# the same segments (IoU 1.0) on every test image.

import argparse, glob, os, sys, time
import numpy as np
import torch
from torch import nn
from PIL import Image
import coremltools as ct
from coremltools.converters.mil import Builder as mb
from coremltools.converters.mil.frontend.torch import ops as ct_ops
from coremltools.converters.mil.frontend.torch.torch_op_registry import register_torch_op
from transformers import AutoImageProcessor, Mask2FormerForUniversalSegmentation
from transformers.utils import logging
import transformers.models.mask2former.modeling_mask2former as modeling

logging.set_verbosity_error()
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from export_mask2former import Logits

CHECKPOINT = "facebook/mask2former-swin-base-coco-panoptic"
MODELS = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "modules", "Segment", "models")
ONNX = os.path.join(MODELS, "mask2former-swin-base-coco-panoptic.onnx")
OUTPUT = os.path.join(MODELS, "mask2former-swin-base-coco-panoptic.mlpackage")
SIZE = 384


_original_cast = ct_ops._cast
def _cast_one_element(context, node, dtype, dtype_name):
    x = ct_ops._get_inputs(context, node, expected=1)[0]
    if x.can_be_folded_to_const() and np.ndim(x.val) > 0 and np.size(x.val) == 1:
        context.add(ct_ops.mb.const(val=dtype(np.asarray(x.val).item()), name=node.name), node.name)
        return
    _original_cast(context, node, dtype, dtype_name)
ct_ops._cast = _cast_one_element


@register_torch_op(override=True)
def meshgrid(context, node):
    """torch.meshgrid, from its inputs flattened to 1-d first."""
    inputs = ct_ops._get_inputs(context, node, expected=[1, 2])
    indexing = inputs[1].val if len(inputs) > 1 else "ij"
    indexing = ct_ops._get_kwinputs(context, node, "indexing", default=[indexing])[0]
    if hasattr(indexing, "val"):
        indexing = indexing.val
    tensors = [mb.reshape(x=t, shape=[-1], name=f"{node.name}_flat_{i}") for i, t in enumerate(inputs[0])]
    sizes = [t.shape[0] for t in tensors]
    grids = []
    for i, t in enumerate(tensors):
        view = [1] * len(tensors); view[i] = -1
        reps = list(sizes); reps[i] = 1
        grid = mb.tile(x=mb.reshape(x=t, shape=view, name=f"{node.name}_view_{i}"), reps=reps, name=f"{node.name}_grid_{i}")
        if indexing == "xy":
            grid = mb.transpose(x=grid, perm=[1, 0] + list(range(2, len(tensors))), name=f"{node.name}_xy_{i}")
        grids.append(grid)
    context.add(tuple(grids), node.name)


def deformable_attention_rank5(self, hidden_states, attention_mask=None, encoder_hidden_states=None, encoder_attention_mask=None,
                               position_embeddings=None, reference_points=None, spatial_shapes_list=None, level_start_index=None,
                               output_attentions=False):
    """Mask2FormerPixelDecoderEncoderMultiscaleDeformableAttention.forward, a level at a time."""
    if position_embeddings is not None:
        hidden_states = self.with_pos_embed(hidden_states, position_embeddings)
    if reference_points.shape[-1] != 2:
        raise ValueError("only 2-d reference points are rewritten")
    B, Q, _ = hidden_states.shape
    _, S, _ = encoder_hidden_states.shape
    H, L, P, D = self.n_heads, self.n_levels, self.n_points, self.d_model // self.n_heads
    value = self.value_proj(encoder_hidden_states)
    if attention_mask is not None:
        value = value.masked_fill(attention_mask[..., None], float(0))
    value = value.view(B, S, H, D)
    offsets = self.sampling_offsets(hidden_states).view(B, Q, H, L, P * 2)
    weights = nn.functional.softmax(self.attention_weights(hidden_states).view(B, Q, H, L * P), -1)
    weights = weights.transpose(1, 2).reshape(B * H, 1, Q, L * P)
    values = value.split([h * w for h, w in spatial_shapes_list], dim=1)
    sampled = []
    for level, (h, w) in enumerate(spatial_shapes_list):
        value_l = values[level].flatten(2).transpose(1, 2).reshape(B * H, D, h, w)
        offset_l = offsets[:, :, :, level].reshape(B, Q, H, P, 2)
        location = reference_points[:, :, None, level, None, :] + offset_l / torch.tensor([w, h], dtype=offset_l.dtype)
        grid = (2 * location - 1).transpose(1, 2).flatten(0, 1)
        sampled.append(nn.functional.grid_sample(value_l, grid, mode="bilinear", padding_mode="zeros", align_corners=False))
    output = (torch.stack(sampled, dim=-2).flatten(-2) * weights).sum(-1).view(B, H * D, Q).transpose(1, 2).contiguous()
    return self.output_proj(output), None


def segments_of(processor, class_logits, mask_logits):
    holder = type("Outputs", (), {"class_queries_logits": torch.from_numpy(class_logits), "masks_queries_logits": torch.from_numpy(mask_logits)})()
    return processor.post_process_panoptic_segmentation(holder, target_sizes=[(SIZE, SIZE)])[0]


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--check", metavar="FOLDER", help="images to check the Core ML model against the ONNX one on")
    args = parser.parse_args()

    processor = AutoImageProcessor.from_pretrained(CHECKPOINT)
    model = Mask2FormerForUniversalSegmentation.from_pretrained(CHECKPOINT).eval()
    example = torch.zeros(1, 3, SIZE, SIZE)
    with torch.no_grad():
        original = Logits(model)(torch.randn(1, 3, SIZE, SIZE, generator=torch.Generator().manual_seed(0)))
        modeling.Mask2FormerPixelDecoderEncoderMultiscaleDeformableAttention.forward = deformable_attention_rank5
        rewritten = Logits(model)(torch.randn(1, 3, SIZE, SIZE, generator=torch.Generator().manual_seed(0)))
        difference = max((a - b).abs().max().item() for a, b in zip(original, rewritten))
        print(f"The attention a level at a time, against the original in PyTorch: max difference {difference:.2e}")
        if difference > 1e-4:
            sys.exit("The rewritten attention does not give the same output.")
        traced = torch.jit.trace(Logits(model).eval(), example)
    start = time.perf_counter()
    ml = ct.convert(traced, inputs=[ct.TensorType(name="pixel_values", shape=(1, 3, SIZE, SIZE))],
                    outputs=[ct.TensorType(name="class_queries_logits"), ct.TensorType(name="masks_queries_logits")],
                    convert_to="mlprogram", compute_precision=ct.precision.FLOAT32, minimum_deployment_target=ct.target.macOS14)
    ml.short_description = f"{CHECKPOINT}, class and mask logits, float32, for mcap"
    ml.save(OUTPUT)
    print(f"Wrote {OUTPUT} in {time.perf_counter() - start:.0f} s.")

    if not args.check:
        return
    import onnxruntime as ort
    paths = sorted(p for p in glob.glob(os.path.join(args.check, "**", "*"), recursive=True)
                   if p.lower().endswith((".png", ".jpg", ".jpeg", ".webp")))
    if not paths:
        sys.exit(f"No images in {args.check}.")
    session = ort.InferenceSession(ONNX, providers=["CPUExecutionProvider"])
    coreml = ct.models.MLModel(OUTPUT, compute_units=ct.ComputeUnit.CPU_AND_GPU)
    total, worst_iou, unmatched, times = 0, 1.0, 0, []
    for path in paths:
        pixels = processor(images=Image.open(path).convert("RGB"), return_tensors="np")["pixel_values"].astype(np.float32)
        ref = session.run(["class_queries_logits", "masks_queries_logits"], {"pixel_values": pixels})
        start = time.perf_counter()
        got = coreml.predict({"pixel_values": pixels})
        times.append(time.perf_counter() - start)
        a = segments_of(processor, ref[0], ref[1])
        b = segments_of(processor, got["class_queries_logits"].astype(np.float32), got["masks_queries_logits"].astype(np.float32))
        for sa in a["segments_info"]:
            total += 1
            best = 0.0
            for sb in b["segments_info"]:
                if sb["label_id"] == sa["label_id"]:
                    ma, mbk = (a["segmentation"] == sa["id"]).numpy(), (b["segmentation"] == sb["id"]).numpy()
                    best = max(best, (ma & mbk).sum() / max((ma | mbk).sum(), 1))
            if best == 0:
                unmatched += 1
            else:
                worst_iou = min(worst_iou, best)
    print(f"{len(paths)} images, {total} segments: worst IoU against the ONNX model {worst_iou:.4f}, unmatched {unmatched}; "
          f"{np.median(times[1:] or times) * 1000:.0f} ms an image on the GPU")
    if unmatched or worst_iou < 0.98:
        sys.exit("The Core ML model segments differently from the ONNX one.")


if __name__ == "__main__":
    main()