"""
[Yclip-B] ONNX 独立推理示例（不依赖原 open_clip 包）。
仅需要：onnxruntime + transformers(tokenizer) + PIL + numpy。
用法：python inference_onnx_example.py [图片路径] [文本1] [文本2] ...
"""
import sys, os
import numpy as np
from PIL import Image
from transformers import AutoTokenizer

BASE = os.path.dirname(os.path.abspath(__file__))
IMG_ONNX = os.path.join(BASE, "Yclip-B-image-encoder.onnx")
TXT_ONNX = os.path.join(BASE, "Yclip-B-text-encoder.onnx")
TOK_DIR = os.path.join(BASE, "Yclip-B-tokenizer")

IMAGE_SIZE = 224
MEAN = np.array([0.5, 0.5, 0.5], dtype=np.float32)
STD = np.array([0.5, 0.5, 0.5], dtype=np.float32)
CONTEXT_LENGTH = 77


def preprocess_image(path):
    img = Image.open(path).convert("RGB").resize((IMAGE_SIZE, IMAGE_SIZE), Image.BICUBIC)
    arr = np.asarray(img, dtype=np.float32) / 255.0
    arr = (arr - MEAN) / STD
    arr = arr.transpose(2, 0, 1)[None]  # [1,3,224,224]
    return arr.astype(np.float32)


def tokenize(texts):
    tok = AutoTokenizer.from_pretrained(TOK_DIR, padding_side="right")
    enc = tok(texts, padding="max_length", truncation=True,
              return_tensors="np", return_attention_mask=False, max_length=CONTEXT_LENGTH)
    return enc["input_ids"].astype(np.int64)


def softmax(x):
    e = np.exp(x - x.max()); return e / e.sum()


def main():
    args = sys.argv[1:]
    # 默认示例图：请将一张测试图放到 export 目录下命名为 demo.png，或用第 1 个参数指定路径
    img_path = args[0] if args else os.path.join(BASE, "demo.png")
    texts = args[1:] if len(args) > 1 else ["一只猪", "一个狗", "一只猫"]
    if not os.path.exists(img_path):
        print(f"[提示] 未找到示例图 {img_path}，请传入图片路径：python inference_onnx_example.py <图片> <文本1> <文本2> ...")
        return

    import onnxruntime as ort
    img_sess = ort.InferenceSession(IMG_ONNX, providers=["CPUExecutionProvider"])
    txt_sess = ort.InferenceSession(TXT_ONNX, providers=["CPUExecutionProvider"])

    imf = img_sess.run(["image_features"], {"pixel_values": preprocess_image(img_path)})[0]
    txt_ids = tokenize(texts)
    rows = [txt_sess.run(["text_features"], {"input_ids": txt_ids[i:i+1]})[0] for i in range(len(texts))]
    txtf = np.concatenate(rows, axis=0)

    imf = imf / np.linalg.norm(imf, axis=-1, keepdims=True)
    txtf = txtf / np.linalg.norm(txtf, axis=-1, keepdims=True)
    logits = 100.0 * (imf @ txtf.T).squeeze()
    probs = softmax(logits)

    print(f"图片: {img_path}")
    for t, p in zip(texts, probs):
        print(f"  {t}: {p:.4f}")
    print(f"预测: {texts[int(np.argmax(probs))]} (置信度 {probs.max():.4f})")


if __name__ == "__main__":
    main()