Skip to main content

4.2 Python Inference Example

This article uses image classification as an example to demonstrate how to perform model inference in a Python environment. Image classification is a fundamental task in computer vision, aiming to identify the main object categories within an image (e.g., cats, dogs, cars). This example is based on the MobileNetV2 model and achieves high-performance inference on the SpaceMIT platform using onnxruntime and the spacemit-ort acceleration library.

Clone the Code​

git clone https://gitee.com/bianbu/spacemit-demo.git ~

Installation Dependencies​

It's recommended to use a Python virtual environment：

python -m venv .venv
source .venv/bin/activate

Install necessary dependencies using the SpaceMIT private source:

pip install opencv-python --index-url https://git.spacemit.com/api/v4/projects/33/packages/pypi/simple
pip install spacemit-ort --index-url https://git.spacemit.com/api/v4/projects/33/packages/pypi/simple

Run the Test​

Navigate to the project's example directory and execute the inference script:

cd ~/spacemit_demo/examples/CV/mobilenet_v2/python
python test_mobilenet_v2.py

example output:

Final result: class=n02123045 tabby, tabby cat;

Inference Process and Code Analysis​

Import the Acceleration Library​

import spacemit_ort

Image Loading and Preprocessing​

from PIL import Image
import numpy as np
import cv2

def preprocess(img):
    img = img / 255.0  # Normalize to [0, 1]
    img = cv2.resize(img, (256, 256))  # Resize
    y0, x0 = (256 - 224) // 2, (256 - 224) // 2
    img = img[y0:y0+224, x0:x0+224, :]  # Center Crop
    img = (img - [0.485, 0.456, 0.406]) / [0.229, 0.224, 0.225]  # Normalization
    img = np.transpose(img, (2, 0, 1))  # HWC → CHW
    img = np.expand_dims(img.astype(np.float32), axis=0)  # Add Batch Dimension
    return img

# Load and Preprocess Image
img = Image.open(args.image_path).convert('RGB')
img = preprocess(np.array(img))

Model Loading and Inference Session Setup​

session_options = onnxruntime.SessionOptions() # Create Onnxruntime Session
session_options.intra_op_num_threads = 4  # Set Number of Threads (Recommended: 4)

# Load Model and Set Up Inference Session
session = onnxruntime.InferenceSession(
    args.model,
    sess_options=session_options,
    providers=["SpaceMITExecutionProvider"]
)

Model Inference Execution​

input_name = session.get_inputs()[0].name # Get Input Node Name
output_name = session.get_outputs()[0].name # Get Output Node Name

# Execute Inference and Store Results in Result
result = session.run([output_name], {input_name: img})[0]

Post-processing of Inference Results​

# Get Top-5 Class Indices, Sorted by Probability (High to Low)
top_k = result.argsort()[-5:][::-1]