Skip to main content

4.2 Python Inference Example

This section demonstrates how to run model inference using Python, using image classification as the example task. Image classification identifies the main object category in an image (e.g., cats, dogs, cars).

This example uses MobileNetV2 to run high-performance inference on the SpaceMIT platform with onnxruntime and the spacemit-ort acceleration library.

Clone the Code​

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

Install Dependencies and Set Up the Environment​

We recommend creating a Python virtual environment:

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

Install the required dependencies from the SpaceMIT private package 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
pip install pillow --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 the 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 an ONNX Runtime session
session_options.intra_op_num_threads = 4  # Set the number of threads (recommended: 4)

# Load the model and set up the inference session
session = onnxruntime.InferenceSession(
    args.model,
    sess_options=session_options,
    providers=["SpaceMITExecutionProvider"]
)

Model Inference Execution​

input_name = session.get_inputs()[0].name # Get the input node name
output_name = session.get_outputs()[0].name # Get the output node name

# Execute inference and store the results in result
result = session.run([output_name], {input_name: img})[0]

Post-processing of Inference Results​

# Get the top-5 class indices, sorted by probability from high to low
top_k = result.argsort()[-5:][::-1]