Amazing-Python-Scripts

Форк
0
166 строк · 6.4 Кб
1
#!/usr/bin/env python3
2
"""
3
 Copyright (c) 2018 Intel Corporation.
4

5
 Permission is hereby granted, free of charge, to any person obtaining
6
 a copy of this software and associated documentation files (the
7
 "Software"), to deal in the Software without restriction, including
8
 without limitation the rights to use, copy, modify, merge, publish,
9
 distribute, sublicense, and/or sell copies of the Software, and to
10
 permit persons to whom the Software is furnished to do so, subject to
11
 the following conditions:
12

13
 The above copyright notice and this permission notice shall be
14
 included in all copies or substantial portions of the Software.
15

16
 THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
17
 EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
18
 MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
19
 NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
20
 LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
21
 OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
22
 WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
23
"""
24

25
import os
26
import sys
27
import logging as log
28
from openvino.inference_engine import IENetwork, IECore
29

30

31
class Network:
32
    """
33
    Load and configure inference plugins for the specified target devices 
34
    and performs synchronous and asynchronous modes for the specified infer requests.
35
    """
36

37
    def __init__(self):
38
        self.net = None
39
        self.plugin = None
40
        self.input_blob = None
41
        self.out_blob = None
42
        self.net_plugin = None
43
        self.infer_request_handle = None
44

45
    def load_model(self, model, device, input_size, output_size, num_requests, cpu_extension=None, plugin=None):
46
        """
47
         Loads a network and an image to the Inference Engine plugin.
48
        :param model: .xml file of pre trained model
49
        :param cpu_extension: extension for the CPU device
50
        :param device: Target device
51
        :param input_size: Number of input layers
52
        :param output_size: Number of output layers
53
        :param num_requests: Index of Infer request value. Limited to device capabilities.
54
        :param plugin: Plugin for specified device
55
        :return:  Shape of input layer
56
        """
57

58
        model_xml = model
59
        model_bin = os.path.splitext(model_xml)[0] + ".bin"
60
        # Plugin initialization for specified device
61
        # and load extensions library if specified
62
        if not plugin:
63
            log.info("Initializing plugin for {} device...".format(device))
64
            self.plugin = IECore()
65
        else:
66
            self.plugin = plugin
67

68
        if cpu_extension and 'CPU' in device:
69
            self.plugin.add_extension(cpu_extension, "CPU")
70

71
        # Read IR
72
        log.info("Reading IR...")
73
        self.net = self.plugin.read_network(model=model_xml, weights=model_bin)
74
        log.info("Loading IR to the plugin...")
75

76
        if "CPU" in device:
77
            supported_layers = self.plugin.query_network(self.net, "CPU")
78
            not_supported_layers = \
79
                [l for l in self.net.layers.keys() if l not in supported_layers]
80
            if len(not_supported_layers) != 0:
81
                log.error("Following layers are not supported by "
82
                          "the plugin for specified device {}:\n {}".
83
                          format(device,
84
                                 ', '.join(not_supported_layers)))
85
                log.error("Please try to specify cpu extensions library path"
86
                          " in command line parameters using -l "
87
                          "or --cpu_extension command line argument")
88
                sys.exit(1)
89

90
        if num_requests == 0:
91
            # Loads network read from IR to the plugin
92
            self.net_plugin = self.plugin.load_network(
93
                network=self.net, device_name=device)
94
        else:
95
            self.net_plugin = self.plugin.load_network(
96
                network=self.net, num_requests=num_requests, device_name=device)
97
            # log.error("num_requests != 0")
98

99
        self.input_blob = next(iter(self.net.inputs))
100
        self.out_blob = next(iter(self.net.outputs))
101
        assert len(self.net.inputs.keys()) == input_size, \
102
            "Supports only {} input topologies".format(len(self.net.inputs))
103
        assert len(self.net.outputs) == output_size, \
104
            "Supports only {} output topologies".format(len(self.net.outputs))
105

106
        return self.plugin, self.get_input_shape()
107

108
    def get_input_shape(self):
109
        """
110
        Gives the shape of the input layer of the network.
111
        :return: None
112
        """
113
        return self.net.inputs[self.input_blob].shape
114

115
    def performance_counter(self, request_id):
116
        """
117
        Queries performance measures per layer to get feedback of what is the
118
        most time consuming layer.
119
        :param request_id: Index of Infer request value. Limited to device capabilities
120
        :return: Performance of the layer  
121
        """
122
        perf_count = self.net_plugin.requests[request_id].get_perf_counts()
123
        return perf_count
124

125
    def exec_net(self, request_id, frame):
126
        """
127
        Starts asynchronous inference for specified request.
128
        :param request_id: Index of Infer request value. Limited to device capabilities.
129
        :param frame: Input image
130
        :return: Instance of Executable Network class
131
        """
132

133
        self.infer_request_handle = self.net_plugin.start_async(
134
            request_id=request_id, inputs={self.input_blob: frame})
135
        return self.net_plugin
136

137
    def wait(self, request_id):
138
        """
139
        Waits for the result to become available.
140
        :param request_id: Index of Infer request value. Limited to device capabilities.
141
        :return: Timeout value
142
        """
143
        wait_process = self.net_plugin.requests[request_id].wait(-1)
144
        return wait_process
145

146
    def get_output(self, request_id, output=None):
147
        """
148
        Gives a list of results for the output layer of the network.
149
        :param request_id: Index of Infer request value. Limited to device capabilities.
150
        :param output: Name of the output layer
151
        :return: Results for the specified request
152
        """
153
        if output:
154
            res = self.infer_request_handle.outputs[output]
155
        else:
156
            res = self.net_plugin.requests[request_id].outputs[self.out_blob]
157
        return res
158

159
    def clean(self):
160
        """
161
        Deletes all the instances
162
        :return: None
163
        """
164
        del self.net_plugin
165
        del self.plugin
166
        del self.net
167

Использование cookies

Мы используем файлы cookie в соответствии с Политикой конфиденциальности и Политикой использования cookies.

Нажимая кнопку «Принимаю», Вы даете АО «СберТех» согласие на обработку Ваших персональных данных в целях совершенствования нашего веб-сайта и Сервиса GitVerse, а также повышения удобства их использования.

Запретить использование cookies Вы можете самостоятельно в настройках Вашего браузера.