mxnet inference 最后输出层结果

mxnet inference 最后输出层结果

import cv2
import numpy as np
import mxnet as mx
from collections import namedtuple

class mxnetPredictor:
    def __init__(self, prefix, epoch, input_shape, ctx=mx.cpu()):
        sym, arg_params, aux_params = mx.model.load_checkpoint(prefix=prefix, epoch=epoch)
        self.input_shape = input_shape
        self.outputs = sym.list_outputs()

        self.mod = mx.mod.Module(symbol=sym, context=ctx, label_names=None)
        self.mod.bind(for_training=False, data_shapes=[('data', self.input_shape)])
        self.mod.set_params(arg_params, aux_params, allow_missing=False, allow_extra=False)

    def forward(self, tensor):
        self.mod.forward(Batch([mx.nd.array(tensor)]))
        out = self.mod.get_outputs()
        out = list(map(lambda x:x.asnumpy(), out))
        return out
        
def img2tensor(mat):
    if mat is None:
        tensor = np.random.randint(0, 256, input_shape, dtype=np.uint8)
    else:
        mat = cv2.resize(mat, (input_shape[3], input_shape[2]))
        if input_shape[1] == 1:
            mat = cv2.cvtColor(mat, cv2.COLOR_BGR2GRAY)
            mat = mat[:,:,np.newaxis]
        tensor = np.expand_dims(np.transpose(mat, (2,0,1)), axis=0)
    tensor = tensor.astype(np.float32)
    return tensor    

if __name__ == "__main__":
    img_path = '01.jpg'
    mxnet_path = 'model-sim'
    input_shape=(1, 3, 320, 320)
    
    # data
    mat = None
    if len(img_path) != 0:
        print("imread:", img_path)
        mat = cv2.imread(img_path)
    tensor = img2tensor(mat)    

    # mxnet inference
    Batch = namedtuple('Batch',['data'])
    netmxnet = mxnetPredictor(mxnet_path, 0, input_shape)
    out_mxnet = netmxnet.forward(tensor)
    for out in out_mxnet:
        print('mxnet max',np.max(out),out.shape)