Previous tutorial: How to run a custom OCR model
Next tutorial: Convert PyTorch classification model and launch with OpenCV C++
| Original author | Anastasia Murzova |
|---|---|
| Compatibility | OpenCV >= 4.5 |
Goals
In this tutorial you will learn how to
- Convert PyTorch classification model to ONNX format
- Run the converted PyTorch model using the OpenCV Python API
- Evaluate PyTorch and OpenCV DNN models.
We will discuss the above points using the ResNet-50 architecture as an example.
Introduction
Let’s briefly review the key concepts involved in converting PyTorch models to the OpenCV API. The first step in converting a PyTorch model to cv.dnn.Net is to convert the model to ONNX format. ONNX aims to enable the interchange of neural networks between different frameworks. There is a built-in function in PyTorch for ONNX conversion: torch.onnx.export. Additionally, the obtained .onnx model is passed into cv.dnn.readNetFromONNX.
Requirements
To experiment with the code below, you need to install a set of libraries. For this we will use a virtual environment with python3.7+:
virtualenv -p /usr/bin/python3.7 <env_dir_path> source <env_dir_path>/bin/activate
To build OpenCV-Python from source, follow the instructions in Introduction to OpenCV.
Before you start installing the library, you can customize requirements.txt to exclude or include (for example, opencv-python) certain dependencies. The following line installs the startup requirements into the previously activated virtual environment:
pip install -r requirements.txt
Practice
In this part we will cover the following points:
- Create a classification model transformation pipeline and provide inferences
- Evaluate and test classification models
If you only want to run an evaluation or test model pipeline, you can skip the “Model Transformation Pipeline” section.
Model conversion pipeline
The code for this subchapter is located in the dnn_model_runner module and can be executed through the following command line:
python -m dnn_model_runner.dnn_conversion.pytorch.classification.py_to_py_resnet50
The following code contains instructions for the following steps:
- Instantiate a PyTorch model
- Convert PyTorch model to
.onnx - Reading transmitted network using OpenCV API
- Prepare to enter data
- provide reasoning
# Initialize PyTorch ResNet-50 model
original_model = models.resnet50(pretrained=True)
# Get the path converted to ONNX PyTorch model
full_model_path = get_pytorch_onnx_model(original_model)
# Use OpenCV API to read the converted .onnx model
opencv_net = cv2.dnn.readNetFromONNX(full_model_path)
print("OpenCV model read successfully. Layer ID: \\
", opencv_net.getLayerNames())
# Get preprocessed image
input_img = get_preprocessed_img("../data/squirrel_cls.jpg")
# Get image network tags
imagenet_labels = get_imagenet_labels("../data/dnn/classification_classes_ILSVRC2012.txt")
# Get OpenCV DNN prediction results
get_opencv_dnn_prediction(opencv_net, input_img, imagenet_labels)
# Get the original PyTorch ResNet50 prediction results
get_pytorch_dnn_prediction(original_model, input_img, imagenet_labels)
To provide model inference, we will use the following photo of a squirrel corresponding to ImageNet class ID 335 (licensed under CC0 license):
Fox squirrel, Eastern fox squirrel, Sciurus niger
Classification model input image
In order to label decode the predictions, we also need the imagenet_classes.txt file, which contains the complete list of ImageNet classes.
Let’s take a pre-trained PyTorch ResNet-50 as an example and dive into each step:
- Instantiate the PyTorch ResNet-50 model:
# Initialize PyTorch ResNet-50 model original_model = models.resnet50(pretrained=True)
- Convert PyTorch model to ONNX:
# Define the saving directory for further converted models
onnx_model_path = "models" # Define names for further converted models
# Define names for further converted models
onnx_model_name = "resnet50.onnx" # Create a directory for further converted models
# Create directories for further converted models
os.makedirs(onnx_model_path, exist_ok=True)
# Get the full path of the converted model
full_model_path = os.path.join(onnx_model_path, onnx_model_name)
# Generate model input
generated_input = Variable(
torch.randn(1, 3, 224, 224)
)
# Export the model to ONNX format
torch.onnx.export(
original_model,
generated_input,
full_model_path,
verbose=True,
input_names=["input"]、
output_names=["output"]、
opset_version=11
)
After successfully executing the above code, we will get models/resnet50.onnx.
- Use cv.dnn.readNetFromONNX to read the transferred network and pass the ONNX model obtained in the previous step into it:
# Read the converted .onnx model using OpenCV API opencv_net = cv2.dnn.readNetFromONNX(full_model_path)
- Prepare to enter data:
# Read image
input_img = cv2.imread(img_path, cv2.IMREAD_COLOR)
input_img = input_img.astype(np.float32)
input_img = cv2.resize(input_img, (256, 256))
# Define preprocessing parameters
mean = np.array([0.485, 0.456, 0.406]) * 255.0
scale=1/255.0
std = [0.229, 0.224, 0.225] # Prepare input blob to fit the model input: 0.229, 0.224, 0.225
# Prepare the input blob to fit the model input:
# Subtract the average
# 2. Scale pixel value from 0 to 1
input_blob = cv2.dnn.blobFromImage(
image=input_img、
scalefactor=scale、
size=(224, 224), # Image target size
mean=mean、
swapRB=True, # BGR -> RGB
crop=True # Center crop
)
# 3. Divide by standard
input_blob[0] /= np.asarray(std, dtype=np.float32).reshape(3, 1, 1)
In this step, we read the image and prepare the model input using the cv.dnn.blobFromImage function, which returns a 4-dimensional blob. Note that cv.dnn.blobFromImage first subtracts the average value and only then multiplies the pixel value by the scale. Therefore, the average is multiplied by 255.0 to reproduce the preprocessing sequence of the original image:
img /= 255.0 img -= [0.485, 0.456, 0.406] (0.485, 0.456, 0.406) img /= [0.229, 0.224, 0.225]
- OpenCV cv.dnn.Net inference:
# Set up OpenCV DNN input
opencv_net.setInput(preproc_img)
# Set up OpenCV DNN input
out = opencv_net.forward()
print("OpenCV DNN prediction: \\
")
print("* shape: ", out.shape)
# Get the predicted class ID
imagenet_class_id = np.argmax(out)
# Get confidence
confidence = out[0][imagenet_class_id].
print("* class ID: {}, label: {}".format(imagenet_class_id, imagenet_labels[imagenet_class_id]))
print("* confidence: {:.4f}".format(confidence))
After executing the above code, we will get the following output:
OpenCV DNN prediction: * Shape: (1, 1000) * Class ID: 335, label: fox squirrel, eastern fox squirrel, Sciurus niger * Confidence 14.8308
- PyTorch ResNet-50 model inference:
original_net.eval()
preproc_img = torch.FloatTensor(preproc_img)
# Reasoning
out = original_net(preproc_img)
print("\\
PyTorch model prediction:\\
")
print("* shape: ", out.shape)
# Get the predicted class ID
imagenet_class_id = torch.argmax(out, axis=1).item()
print("* class ID: {}, label: {}".format(imagenet_class_id, imagenet_labels[imagenet_class_id]))
# Get confidence
confidence = out[0][imagenet_class_id].
print("* confidence: {:.4f}".format(confidence.item())
After executing the above code, we will get the following output:
PyTorch model prediction: * shape: torch.Size([1, 1000]) * Class ID: 335, label: fox squirrel, eastern fox squirrel, Sciurus niger * Confidence 14.8308
The original ResNet-50 model and cv.dnn.Net have the same inference results. For extended evaluation of models, we can use py_to_py_cls from the dnn_model_runner module. This module part will be introduced in the next chapter.
Model Evaluation
The dnn_model_runner module proposed in samples/dnn allows running the complete evaluation pipeline on the ImageNet dataset and test execution of the following PyTorch classification models:
- alexnet
- vgg11
- vgg13
- vgg16
- vgg19
- resnet18
- resnet34
- resnet50
- resnet101
- resnet152
- squeezenet1_0
- squeezenet1_1
- resnext50_32x4d
- resnext101_32x8d
- wide_resnet50_2
- wide_resnet101_2
This list can also be extended with further appropriate evaluation pipeline configurations.
Evaluation Mode
The following line indicates running the module in evaluation mode:
python -m dnn_model_runner.dnn_conversion.pytorch.classification.py_to_py_cls--model_name<pytorch_cls_model_name>
The classification model selected from the list will be read into an OpenCV cv.dnn.Net object. The evaluation results (accuracy, inference time, L1) of PyTorch and OpenCV models are written to log files. Inference time values are also displayed graphically to summarize the model information obtained.
The necessary evaluation configuration is defined in test_config.py, which can be modified based on the actual path to the data location:
@dataclass
class TestClsConfig:
batch_size: int = 50
frame_size: int = 224
img_root_dir: str = "./ILSVRC2012_img_val" # Position of image class matching
# Position of image class matching
img_cls_file: str = "./val.txt" # Position of image class matching
bgr_too_rgb: bool = True
To start the evaluation of PyTorch ResNet-50, run the following line:
python -m dnn_model_runner.dnn_conversion.pytorch.classification.py_to_py_cls --model_name resnet50
After the script is started, a log file containing the evaluation data will be generated in dnn_model_runner/dnn_conversion/logs:
The model PyTorch resnet50 has been successfully obtained and converted to OpenCV DNN resnet50
===== The model is being evaluated using the following parameters:
* Val data location: ./ILSVRC2012_img_val
* Log file location: dnn_model_runner/dnn_conversion/logs/PyTorch_resnet50_log.txt
Test Mode
The following line represents running the module in test mode, i.e. providing the steps for model inference:
python -m dnn_model_runner.dnn_conversion.pytorch.classification.py_to_py_cls --model_name <pytorch_cls_model_name> --test True --default_img_preprocess <True/False> --evaluate False
The default_img_preprocess keyword here defines whether you want to parameterize the model testing process with some specific values, or use default values such as scale, mean or std.
The test configuration is represented in the test_config.py TestClsModuleConfig class:
@dataclass
class TestClsModuleConfig:
cls_test_data_dir: str = "../data"
test_module_name: str = "classification"
test_module_path: str = "classification.py"
input_img: str = os.path.join(cls_test_data_dir, "squirrel_cls.jpg")
model: str = ""
frame_height: str = str(TestClsConfig.frame_size)
frame_width: str = str(TestClsConfig.frame_size)
scale: str = "1.0"
mean: List[str] = field(default_factory=lambda: ["0.0", "0.0", "0.0"])
std: List[str] = field(default_factory=list)
crop: str = "False"
rgb: str = "True"
rsz_height: str = ""
rsz_width: str = ""
classes: str = os.path.join(cls_test_data_dir, "dnn", "classification_classes_ILSVRC2012.txt")
Default image preprocessing options are defined in default_preprocess_config.py. For example
BASE_IMG_SCALE_FACTOR = 1 / 255.0
PYTORCH_RSZ_HEIGHT = 256
PYTORCH_RSZ_WIDTH = 256
pytorch_resize_input_blob = {<!-- -->
"mean": ["123.675", "116.28", "103.53"],
"scale": str(BASE_IMG_SCALE_FACTOR),
"std": ["0.229", "0.224", "0.225"],
"crop": "True",
"rgb": "True",
"rsz_height": str(PYTORCH_RSZ_HEIGHT),
"rsz_width": str(PYTORCH_RSZ_WIDTH)
}
The basis for model testing is provided in samples/dnn/classification.py. Classification.py can be executed automatically using the conversion model provided in --input and the fill parameters of cv.dnn.blobFromImage.
To reproduce the OpenCV steps described in “Model Conversion Pipeline” using dnn_model_runner from scratch, execute the following line:
python -m dnn_model_runner.dnn_conversion.pytorch.classification.py_to_py_cls --model_name resnet50 --test True --default_img_preprocess True --evaluate False
The network prediction results are displayed in the upper left corner of the output window:
ResNet50 OpenCV inference output