File size: 4,560 Bytes
07302a4 | 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 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 | import onnxruntime as ort
import numpy as np
from PIL import Image
import pandas as pd
import requests
import io
import sys
# CONFIG SECTION
ONNX_MODEL_PATH = "./lsnet_xl_artist-dynamo-opset18_merged.onnx"
CSV_PATH = "./class_mapping.csv"
IMAGE_URL = "https://cdn.donmai.us/sample/9f/bb/__vampire_s_sister_original_drawn_by_gogalking__sample-9fbb30aa76bdc8242a1c122d3d6b41d9.jpg"
IMAGE_SIZE = (224, 224)
TOP_K = 5
PREDICTION_THRESHOLD = 0.0
# ------------------------------------------------------------
def preprocess_image_from_url(image_url, size=(224, 224)):
"""
Downloads an image from a URL, preprocesses it, and prepares it for the model.
"""
try:
response = requests.get(image_url)
response.raise_for_status() # Raise an exception for bad status codes
image_bytes = io.BytesIO(response.content)
image = Image.open(image_bytes).convert("RGB")
except requests.exceptions.RequestException as e:
print(f"Error: Failed to download image from URL '{image_url}'.\nDetails: {e}")
sys.exit(1)
except Exception as e:
print(f"Error: Could not process the downloaded image. It may not be a valid image file.\nDetails: {e}")
sys.exit(1)
# Resize the image
image = image.resize(size, Image.Resampling.LANCZOS)
# Convert image to numpy array and scale to [0, 1]
image_np = np.array(image, dtype=np.float32) / 255.0
# Define ImageNet mean and standard deviation for normalization
mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
std = np.array([0.229, 0.224, 0.225], dtype=np.float32)
# Normalize the image
normalized_image = (image_np - mean) / std
# Transpose the dimensions from (H, W, C) to (C, H, W)
transposed_image = normalized_image.transpose((2, 0, 1))
# Add a batch dimension to create a shape of (1, C, H, W)
batched_image = np.expand_dims(transposed_image, axis=0)
return batched_image
def load_labels(csv_path):
"""
Loads the class labels from the provided CSV file into a dictionary,
handling the header row and stripping quotes from names.
"""
try:
df = pd.read_csv(csv_path)
if 'class_id' not in df.columns or 'class_name' not in df.columns:
print(f"Error: CSV file must have 'class_id' and 'class_name' columns.")
sys.exit(1)
df['class_name'] = df['class_name'].str.strip("'")
return dict(zip(df['class_id'], df['class_name']))
except FileNotFoundError:
print(f"Error: CSV file not found at '{csv_path}'")
sys.exit(1)
except Exception as e:
print(f"Error reading CSV file: {e}")
sys.exit(1)
def softmax(x):
"""Compute softmax values for a set of scores."""
e_x = np.exp(x - np.max(x))
return e_x / e_x.sum(axis=0)
def main():
"""
Main function to run the ONNX model inference.
"""
print("1. Loading class labels...")
labels = load_labels(CSV_PATH)
print(f" Loaded {len(labels)} labels.")
print("\n2. Downloading and preprocessing image from URL...")
input_tensor = preprocess_image_from_url(IMAGE_URL, IMAGE_SIZE)
print(f" Image shape: {input_tensor.shape}, Data type: {input_tensor.dtype}")
print("\n3. Initializing ONNX runtime session...")
try:
session = ort.InferenceSession(ONNX_MODEL_PATH, providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name
print(" ONNX session created successfully.")
except Exception as e:
print(f"Error loading ONNX model: {e}")
sys.exit(1)
print("\n4. Running inference...")
results = session.run([output_name], {input_name: input_tensor})
logits = results[0][0]
print(" Inference complete.")
print("\n5. Processing results...")
probabilities = softmax(logits)
top_k_indices = np.argsort(probabilities)[-TOP_K:][::-1]
print(f"\n--- Predictions for image URL (Top K: {TOP_K}, Threshold: {PREDICTION_THRESHOLD:.1%}) ---")
predictions_found = 0
for i, index in enumerate(top_k_indices):
score = probabilities[index]
if score >= PREDICTION_THRESHOLD:
class_name = labels.get(index, f"Unknown Class #{index}")
print(f"Rank {i+1}: {class_name} (Score: {score:.2%})")
predictions_found += 1
if predictions_found == 0:
print("No predictions met the specified threshold.")
if __name__ == "__main__":
main() |