onnx_runtime_squeezenet_image_classification_with_preprocessing.py

1import numpy as np
2import onnxruntime as ort
3from PIL import Image
4
5# Load the model and create an InferenceSession
6# Note: You must have 'squeezenet1.1-7.onnx' in your directory
7# or download it from the ONNX Model Zoo
8session = ort.InferenceSession("squeezenet1.1-7.onnx", providers=["CPUExecutionProvider"])
9
10# Preprocess the image
11def preprocess(image_path):
12    img = Image.open(image_path)
13    img = img.resize((224, 224))
14    img_data = np.array(img).transpose(2, 0, 1)
15    img_data = img_data.astype('float32')
16    # Standard ImageNet normalization
17    mean_vec = np.array([0.485, 0.456, 0.406])
18    stddev_vec = np.array([0.229, 0.224, 0.225])
19    norm_img_data = np.zeros(img_data.shape).astype('float32')
20    for i in range(img_data.shape[0]):
21        norm_img_data[i,:,:] = (img_data[i,:,:] / 255 - mean_vec[i]) / stddev_vec[i]
22    norm_img_data = norm_img_data[np.newaxis, ...]
23    return norm_img_data
24
25# Run inference
26input_name = session.get_inputs()[0].name
27output_name = session.get_outputs()[0].name
28input_data = preprocess("test_image.jpg")
29
30raw_result = session.run([output_name], {input_name: input_data})
31res = np.array(raw_result)
32
33# Postprocess: Get the top predicted class
34predicted_class = np.argmax(res)
35print(f"Predicted class index: {predicted_class}")