File size: 3,389 Bytes
82b04f2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 | import numpy as np
import onnxruntime as ort
from PIL import Image
from transformers import AutoTokenizer
def predict_thumbnail_onnx(onnx_path: str, image_path: str, image_size: int = 299, threshold: float = 0.5):
# 1. Load the ONNX model graph
session = ort.InferenceSession(onnx_path)
input_name = session.get_inputs()[0].name
# 2. Preprocess the image (replicates T.Compose without PyTorch)
img = Image.open(image_path).convert('RGB').resize((image_size, image_size))
img_np = np.array(img, dtype=np.float32) / 255.0 # ToTensor part 1: Scale 0-1
img_np = np.transpose(img_np, (2, 0, 1)) # ToTensor part 2: HWC -> CHW
# Normalize
mean = np.array([0.485, 0.456, 0.406]).reshape(3, 1, 1)
std = np.array([0.229, 0.224, 0.225]).reshape(3, 1, 1)
img_np = (img_np - mean) / std
# Add batch dimension (1, C, H, W)
input_tensor = np.expand_dims(img_np, axis=0).astype(np.float32)
# 3. Run inference through the ONNX graph
outputs = session.run(None, {input_name: input_tensor})
logit = outputs[0].item()
# 4. Apply Sigmoid and threshold
probability = 1 / (1 + np.exp(-logit))
prediction = 1 if probability >= threshold else 0
return probability, prediction
# --- Usage ---
#prob, pred = predict_thumbnail_onnx("./thumbnail_predictor.onnx", "/content/TBxS0XhdfmU-HD.jpg", image_size=299)
#print(f"Class: {pred}, Probability: {prob:.4f}")
def predict_title_onnx(onnx_path: str, title: str, tokenizer_name: str = "sentence-transformers/all-MiniLM-L6-v2", threshold: float = 0.5):
# 1. Initialize the tokenizer
# (In a production environment, load this once outside the function to save time)
tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
# 2. Tokenize the input text and return NumPy arrays ("np")
inputs = tokenizer(
title,
return_tensors="np",
padding="max_length",
truncation=True,
max_length=100 # 128 is plenty for YouTube titles
)
# Convert to int64, which is what ONNX exported transformer models expect
input_ids = inputs["input_ids"].astype(np.int64)
attention_mask = inputs["attention_mask"].astype(np.int64)
# 3. Load the ONNX model graph
session = ort.InferenceSession(onnx_path)
# Prepare the input dictionary
ort_inputs = {
"input_ids": input_ids,
"attention_mask": attention_mask
}
# Check if the exported ONNX graph also expects 'token_type_ids'
# (MiniLM usually doesn't need them, but exporting to ONNX sometimes forces them in)
expected_inputs = [inp.name for inp in session.get_inputs()]
if "token_type_ids" in expected_inputs and "token_type_ids" in inputs:
ort_inputs["token_type_ids"] = inputs["token_type_ids"].astype(np.int64)
# 4. Run inference
outputs = session.run(None, ort_inputs)
# Outputs are shape (Batch, 1). Extract the raw logit.
logit = outputs[0].item()
# 5. Apply Sigmoid and threshold (because of BCEWithLogitsLoss)
probability = 1 / (1 + np.exp(-logit))
prediction = 1 if probability >= threshold else 0
return probability, prediction
# --- Usage ---
#prob, pred = predict_title_onnx(onnx_path: str, title: str, tokenizer_name: str = "sentence-transformers/all-MiniLM-L6-v2", threshold: float = 0.5)
#print(f"Class: {pred}, Probability: {prob:.4f}") |