init - 初始化项目

samples/dnn/dnn_model_runner/dnn_conversion/pytorch/classification/py_to_py_cls.py

@@ -0,0 +1,71 @@
+from torchvision import models
+
+from ..pytorch_model import (
+    PyTorchModelPreparer,
+    PyTorchModelProcessor,
+    PyTorchDnnModelProcessor
+)
+from ...common.evaluation.classification.cls_data_fetcher import PyTorchPreprocessedFetch
+from ...common.test.cls_model_test_pipeline import ClsModelTestPipeline
+from ...common.test.configs.default_preprocess_config import pytorch_resize_input_blob
+from ...common.test.configs.test_config import TestClsConfig
+from ...common.utils import set_pytorch_env, create_extended_parser
+
+model_dict = {
+    "alexnet": models.alexnet,
+
+    "vgg11": models.vgg11,
+    "vgg13": models.vgg13,
+    "vgg16": models.vgg16,
+    "vgg19": models.vgg19,
+
+    "resnet18": models.resnet18,
+    "resnet34": models.resnet34,
+    "resnet50": models.resnet50,
+    "resnet101": models.resnet101,
+    "resnet152": models.resnet152,
+
+    "squeezenet1_0": models.squeezenet1_0,
+    "squeezenet1_1": models.squeezenet1_1,
+
+    "resnext50_32x4d": models.resnext50_32x4d,
+    "resnext101_32x8d": models.resnext101_32x8d,
+
+    "wide_resnet50_2": models.wide_resnet50_2,
+    "wide_resnet101_2": models.wide_resnet101_2
+}
+
+
+class PyTorchClsModel(PyTorchModelPreparer):
+    def __init__(self, height, width, model_name, original_model):
+        super(PyTorchClsModel, self).__init__(height, width, model_name, original_model)
+
+
+def main():
+    set_pytorch_env()
+
+    parser = create_extended_parser(list(model_dict.keys()))
+    cmd_args = parser.parse_args()
+    model_name = cmd_args.model_name
+
+    cls_model = PyTorchClsModel(
+        height=TestClsConfig().frame_size,
+        width=TestClsConfig().frame_size,
+        model_name=model_name,
+        original_model=model_dict[model_name](pretrained=True)
+    )
+
+    pytorch_cls_pipeline = ClsModelTestPipeline(
+        network_model=cls_model,
+        model_processor=PyTorchModelProcessor,
+        dnn_model_processor=PyTorchDnnModelProcessor,
+        data_fetcher=PyTorchPreprocessedFetch,
+        cls_args_parser=parser,
+        default_input_blob_preproc=pytorch_resize_input_blob
+    )
+
+    pytorch_cls_pipeline.init_test_pipeline()
+
+
+if __name__ == "__main__":
+    main()

samples/dnn/dnn_model_runner/dnn_conversion/pytorch/classification/py_to_py_resnet50.py

@@ -0,0 +1,139 @@
+import os
+
+import cv2
+import numpy as np
+import torch
+import torch.onnx
+from torch.autograd import Variable
+from torchvision import models
+
+
+def get_pytorch_onnx_model(original_model):
+    # define the directory for further converted model save
+    onnx_model_path = "models"
+    # define the name of further converted model
+    onnx_model_name = "resnet50.onnx"
+
+    # create directory for further converted model
+    os.makedirs(onnx_model_path, exist_ok=True)
+
+    # get full path to 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)
+    )
+
+    # model export into ONNX format
+    torch.onnx.export(
+        original_model,
+        generated_input,
+        full_model_path,
+        verbose=True,
+        input_names=["input"],
+        output_names=["output"],
+        opset_version=11
+    )
+
+    return full_model_path
+
+
+def get_preprocessed_img(img_path):
+    # read the 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 preprocess 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:
+    # 1. subtract mean
+    # 2. scale to set pixel values from 0 to 1
+    input_blob = cv2.dnn.blobFromImage(
+        image=input_img,
+        scalefactor=scale,
+        size=(224, 224),  # img target size
+        mean=mean,
+        swapRB=True,  # BGR -> RGB
+        crop=True  # center crop
+    )
+    # 3. divide by std
+    input_blob[0] /= np.asarray(std, dtype=np.float32).reshape(3, 1, 1)
+    return input_blob
+
+
+def get_imagenet_labels(labels_path):
+    with open(labels_path) as f:
+        imagenet_labels = [line.strip() for line in f.readlines()]
+    return imagenet_labels
+
+
+def get_opencv_dnn_prediction(opencv_net, preproc_img, imagenet_labels):
+    # set OpenCV DNN input
+    opencv_net.setInput(preproc_img)
+
+    # OpenCV DNN inference
+    out = opencv_net.forward()
+    print("OpenCV DNN prediction: \n")
+    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))
+
+
+def get_pytorch_dnn_prediction(original_net, preproc_img, imagenet_labels):
+    original_net.eval()
+    preproc_img = torch.FloatTensor(preproc_img)
+
+    # inference
+    with torch.no_grad():
+        out = original_net(preproc_img)
+
+    print("\nPyTorch model prediction: \n")
+    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()))
+
+
+def main():
+    # initialize PyTorch ResNet-50 model
+    original_model = models.resnet50(pretrained=True)
+
+    # get the path to the converted into ONNX PyTorch model
+    full_model_path = get_pytorch_onnx_model(original_model)
+
+    # read converted .onnx model with OpenCV API
+    opencv_net = cv2.dnn.readNetFromONNX(full_model_path)
+    print("OpenCV model was successfully read. Layer IDs: \n", opencv_net.getLayerNames())
+
+    # get preprocessed image
+    input_img = get_preprocessed_img("../data/squirrel_cls.jpg")
+
+    # get ImageNet labels
+    imagenet_labels = get_imagenet_labels("../data/dnn/classification_classes_ILSVRC2012.txt")
+
+    # obtain OpenCV DNN predictions
+    get_opencv_dnn_prediction(opencv_net, input_img, imagenet_labels)
+
+    # obtain original PyTorch ResNet50 predictions
+    get_pytorch_dnn_prediction(original_model, input_img, imagenet_labels)
+
+
+if __name__ == "__main__":
+    main()

samples/dnn/dnn_model_runner/dnn_conversion/pytorch/classification/py_to_py_resnet50_onnx.py

@@ -0,0 +1,50 @@
+import os
+
+import torch
+import torch.onnx
+from torch.autograd import Variable
+from torchvision import models
+
+
+def get_pytorch_onnx_model(original_model):
+    # define the directory for further converted model save
+    onnx_model_path = "models"
+    # define the name of further converted model
+    onnx_model_name = "resnet50.onnx"
+
+    # create directory for further converted model
+    os.makedirs(onnx_model_path, exist_ok=True)
+
+    # get full path to 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)
+    )
+
+    # model export into ONNX format
+    torch.onnx.export(
+        original_model,
+        generated_input,
+        full_model_path,
+        verbose=True,
+        input_names=["input"],
+        output_names=["output"],
+        opset_version=11
+    )
+
+    return full_model_path
+
+
+def main():
+    # initialize PyTorch ResNet-50 model
+    original_model = models.resnet50(pretrained=True)
+
+    # get the path to the converted into ONNX PyTorch model
+    full_model_path = get_pytorch_onnx_model(original_model)
+    print("PyTorch ResNet-50 model was successfully converted: ", full_model_path)
+
+
+if __name__ == "__main__":
+    main()

samples/dnn/dnn_model_runner/dnn_conversion/pytorch/pytorch_model.py

@@ -0,0 +1,98 @@
+import os
+
+import cv2
+import torch.onnx
+from torch.autograd import Variable
+
+from ..common.abstract_model import AbstractModel, Framework
+from ..common.utils import DNN_LIB, get_full_model_path
+
+CURRENT_LIB = "PyTorch"
+MODEL_FORMAT = ".onnx"
+
+
+class PyTorchModelPreparer(AbstractModel):
+
+    def __init__(
+            self,
+            height,
+            width,
+            model_name="default",
+            original_model=object,
+            batch_size=1,
+            default_input_name="input",
+            default_output_name="output"
+    ):
+        self._height = height
+        self._width = width
+        self._model_name = model_name
+        self._original_model = original_model
+        self._batch_size = batch_size
+        self._default_input_name = default_input_name
+        self._default_output_name = default_output_name
+
+        self.model_path = self._set_model_path()
+        self._dnn_model = self._set_dnn_model()
+
+    def _set_dnn_model(self):
+        generated_input = Variable(torch.randn(
+            self._batch_size, 3, self._height, self._width)
+        )
+        os.makedirs(self.model_path["path"], exist_ok=True)
+        torch.onnx.export(
+            self._original_model,
+            generated_input,
+            self.model_path["full_path"],
+            verbose=True,
+            input_names=[self._default_input_name],
+            output_names=[self._default_output_name],
+            opset_version=11
+        )
+
+        return cv2.dnn.readNetFromONNX(self.model_path["full_path"])
+
+    def _set_model_path(self):
+        model_to_save = self._model_name + MODEL_FORMAT
+        return get_full_model_path(CURRENT_LIB.lower(), model_to_save)
+
+    def get_prepared_models(self):
+        return {
+            CURRENT_LIB + " " + self._model_name: self._original_model,
+            DNN_LIB + " " + self._model_name: self._dnn_model
+        }
+
+
+class PyTorchModelProcessor(Framework):
+    def __init__(self, prepared_model, model_name):
+        self._prepared_model = prepared_model
+        self._name = model_name
+
+    def get_output(self, input_blob):
+        tensor = torch.FloatTensor(input_blob)
+        self._prepared_model.eval()
+
+        with torch.no_grad():
+            model_out = self._prepared_model(tensor)
+
+        # segmentation case
+        if len(model_out) == 2:
+            model_out = model_out['out']
+
+        out = model_out.detach().numpy()
+        return out
+
+    def get_name(self):
+        return self._name
+
+
+class PyTorchDnnModelProcessor(Framework):
+    def __init__(self, prepared_dnn_model, model_name):
+        self._prepared_dnn_model = prepared_dnn_model
+        self._name = model_name
+
+    def get_output(self, input_blob):
+        self._prepared_dnn_model.setInput(input_blob, '')
+        return self._prepared_dnn_model.forward()
+
+    def get_name(self):
+        return self._name