| """Standalone inference example for the wclip-base ONNX encoders. |
| |
| Only needs onnxruntime + transformers (the copied tokenizer). No original |
| WeCLIP code or weights required. |
| |
| Usage: |
| python wclip-base-inference-example.py [image_path] [text1] [text2] ... |
| """ |
| import sys, os |
| import numpy as np |
| import onnxruntime as ort |
| import cv2 |
| from transformers import AutoTokenizer |
|
|
| BASE = os.path.dirname(os.path.abspath(__file__)) |
| IMG_ONNX = os.path.join(BASE, 'wclip-base-image-encoder.onnx') |
| TXT_ONNX = os.path.join(BASE, 'wclip-base-text-encoder.onnx') |
| TOK_DIR = os.path.join(BASE, 'wclip-base-tokenizer') |
| TEMPERATURE = 0.01 |
|
|
|
|
| def preprocess_image(img_path): |
| data = open(img_path, 'rb').read() |
| arr = np.frombuffer(data, np.uint8) |
| img = cv2.imdecode(arr, cv2.IMREAD_COLOR)[..., ::-1] |
| img = cv2.resize(img, (224, 224)).astype(np.float32) / 255.0 |
| mean = np.array([0.48145466, 0.4578275, 0.40821073], np.float32) |
| std = np.array([0.26862954, 0.26130258, 0.27577711], np.float32) |
| img = (img - mean) * (1.0 / std) |
| return img.transpose(2, 0, 1)[None].astype(np.float32) |
|
|
|
|
| def main(): |
| args = sys.argv[1:] |
| img_path = args[0] if args else os.path.join(BASE, '..', 'weclip', 'assets', 'demo.jpeg') |
| texts = args[1:] if len(args) > 1 else ['狗', '猫', '兔子', '熊猫'] |
|
|
| img_t = preprocess_image(img_path) |
|
|
| tok = AutoTokenizer.from_pretrained(TOK_DIR, padding_side='right') |
| enc = tok(texts, padding='max_length', truncation=True, max_length=72, |
| return_tensors='np', return_attention_mask=True) |
| ids = enc['input_ids'].astype(np.int64) |
| mask = enc['attention_mask'].astype(np.int64) |
|
|
| img_sess = ort.InferenceSession(IMG_ONNX, providers=['CPUExecutionProvider']) |
| txt_sess = ort.InferenceSession(TXT_ONNX, providers=['CPUExecutionProvider']) |
|
|
| imf = img_sess.run(['image_embedding_1024'], {'image_tensor': img_t})[0] |
| outputs = txt_sess.run(['text_embedding_1024'], |
| {'text_token_id_tensor': ids, 'text_mask_tensor': mask})[0] |
|
|
| imf = imf / np.linalg.norm(imf, axis=-1, keepdims=True) |
| outputs = outputs / np.linalg.norm(outputs, axis=-1, keepdims=True) |
| logits = 100.0 * (imf @ outputs.T).squeeze() |
| logits = logits - logits.max() |
| probs = np.exp(logits) / np.exp(logits).sum() |
|
|
| print(f'image : {os.path.basename(img_path)}') |
| print(f'texts : {texts}') |
| print(f'probs : {[round(float(p), 4) for p in probs]}') |
| print(f'pred : {texts[int(np.argmax(probs))]}') |
| return texts[int(np.argmax(probs))] |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
