SAIL (Sophon Artificial Intelligent Library) 是 Sophon Inference 中的核心模块。SAIL 对 BMNNSDK 中的 BMLib、BMDecoder、BMCV、BMRuntime 进行了封装,将 BMNNSDK 中原有的“加载 bmodel 并驱动 TPU 推理”、“驱动 TPU 做图像处理”、“驱动 VPU 做图像和视频解码”等功能抽象成更为简单 的 C++ 接口对外提供;并且使用 pybind11 再次封装,提供简洁易用的 python 接口。
import sophon.sail as sail
engine = sail.Engine(tpu_id = 0)
engine.load(bmodel_path)
# engine = sail.Engine(bmodel_path,tpu_id,mode)
class PreProcessor:
def __init__(self, bmcv, scale):
self.bmcv = bmcv
self.ab = [x * scale for x in [1, -123, 1, -117, 1, -104]]
def process(self, input, output):
tmp = self.bmcv.vpp_resize(input, 300, 300)
self.bmcv.convert_to(tmp, output, ((self.ab[0], self.ab[1]), (self.ab[2], self.ab[3]), (self.ab[4], self.ab[5])))
bmcv = sail.Bmcv(handle) #图形处理加速模块
scale = engine.get_input_scale(graph_name, input_name)
pre_processor = PreProcessor(bmcv, scale) #预处理初始化
img0 = decoder.read(handle) #解码视频输出image
img1 = bmcv.tensor_to_bm_image(input) #将推理的输入地址挂载到image
pre_processor.process(img0, img1) #预处理
graph_name = engine.get_graph_names()[0]
engine.set_io_mode(graph_name, sail.IOMode.SYSO)
input_name = engine.get_input_names(graph_name)[0]
output_name = engine.get_output_names(graph_name)[0]
input_shape = [1, 3, 300, 300]
output_shape = [1, 1, 200, 7]
handle = engine.get_handle()
input_dtype = engine.get_input_dtype(graph_name, input_name)
output_dtype = engine.get_output_dtype(graph_name, output_name)
input = sail.Tensor(handle, input_shape, input_dtype, False, True)
output = sail.Tensor(handle, output_shape, output_dtype, True, True)
input_tensors = { input_name: input }
output_tensors = { output_name: output }
...
#此处省略 解码,预处理 代码
...
engine.process(graph_name, input_tensors, output_tensors) #推理
out = output.asnumpy()
dets = post_processor.process(out, img0.width(), img0.height()) #后处理
...