Upload easy_inference.py

inference/easy_inference.py

@@ -0,0 +1,106 @@
+import argparse
+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):
+    session = ort.InferenceSession(onnx_path)
+    input_name = session.get_inputs()[0].name
+
+    img = Image.open(image_path).convert('RGB').resize((image_size, image_size))
+    img_np = np.array(img, dtype=np.float32) / 255.0         
+    img_np = np.transpose(img_np, (2, 0, 1))                 
+
+    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
+
+    input_tensor = np.expand_dims(img_np, axis=0).astype(np.float32)
+
+    outputs = session.run(None, {input_name: input_tensor})
+    logit = outputs[0].item()
+
+    probability = 1 / (1 + np.exp(-logit))
+    prediction = 1 if probability >= threshold else 0
+
+    return probability, prediction
+
+def predict_title_onnx(onnx_path: str, title: str, tokenizer_name: str = "sentence-transformers/all-MiniLM-L6-v2", threshold: float = 0.5):
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
+
+    inputs = tokenizer(
+        title,
+        return_tensors="np",
+        padding="max_length",
+        truncation=True,
+        max_length=100 
+    )
+
+    input_ids = inputs["input_ids"].astype(np.int64)
+    attention_mask = inputs["attention_mask"].astype(np.int64)
+
+    session = ort.InferenceSession(onnx_path)
+
+    ort_inputs = {
+        "input_ids": input_ids,
+        "attention_mask": attention_mask
+    }
+
+    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)
+
+    outputs = session.run(None, ort_inputs)
+    logit = outputs[0].item()
+
+    probability = 1 / (1 + np.exp(-logit))
+    prediction = 1 if probability >= threshold else 0
+
+    return probability, prediction
+
+def main():
+    parser = argparse.ArgumentParser(description="ONNX Inference CLI for Thumbnail and Title Predictors")
+    
+    parser.add_argument("--image_path", type=str, help="Path to the thumbnail image to classify.")
+    parser.add_argument("--title", type=str, help="Title text to classify.")
+    
+    # Models default to ./onnx/ directory
+    parser.add_argument("--thumbnail_model", type=str, default="./onnx/thumbnail_predictor.onnx", help="Path to ONNX thumbnail model.")
+    parser.add_argument("--title_model", type=str, default="./onnx/title_predictor.onnx", help="Path to ONNX title model.")
+    
+    # Follows the standard 299 resolution from the original training config[span_0](start_span)[span_0](end_span)
+    parser.add_argument("--image_size", type=int, default=299, help="Input resolution for the image model.")
+    parser.add_argument("--threshold", type=float, default=0.5, help="Classification threshold.")
+    
+    args = parser.parse_args()
+
+    if not args.image_path and not args.title:
+        print("Error: You must provide either --image_path or --title (or both).")
+        return
+
+    if args.image_path:
+        prob, pred = predict_thumbnail_onnx(
+            onnx_path=args.thumbnail_model, 
+            image_path=args.image_path, 
+            image_size=args.image_size, 
+            threshold=args.threshold
+        )
+        print("--- Thumbnail Inference ---")
+        print(f"File: {args.image_path}")
+        print(f"Probability: {prob:.4f}")
+        print(f"Prediction: {pred}\n")
+
+    if args.title:
+        prob, pred = predict_title_onnx(
+            onnx_path=args.title_model, 
+            title=args.title, 
+            threshold=args.threshold
+        )
+        print("--- Title Inference ---")
+        print(f"Text: '{args.title}'")
+        print(f"Probability: {prob:.4f}")
+        print(f"Prediction: {pred}\n")
+
+if __name__ == "__main__":
+    main()