Add ONNX Export and MNN/TNN/ONNXRuntime C++ Demo

Features:

export_onnx.py to export dynamic and static onnx files.
infer_onnx.py, test demo.
C++ inference for MGMatting

export onnx:

# cd to code-base, then run
PYTHONPATH=. python3 ./export_onnx.py --dynamic --simplify --config ./config/MGMatting-DIM-100k.toml --checkpoint ./checkpoints/MGMatting-DIM-100k/latest_model.pth --output_path ./checkpoints/MGMatting-DIM-100k/latest_model.onnx

log:

Torch Version:  1.8.1 ONNX Version:  1.8.0
Namespace(checkpoint='./checkpoints/MGMatting-DIM-100k/latest_model.pth', config='./config/MGMatting-DIM-100k.toml', dynamic=True, output_path='./checkpoints/MGMatting-DIM-100k/latest_model.onnx', simplify=True)
Generate Model Done.
Load ./checkpoints/MGMatting-DIM-100k/latest_model.pth Done.
Export the dynamic onnx model ...
export onnx done.
graph torch-jit-export (
  %image[FLOAT, 1x3xheightxwidth]
  %mask[FLOAT, 1x1xheightxwidth]
) initializers (
  %encoder.conv1.module.weight_u[FLOAT, 32]
...
export onnx sim done.
NodeArg(name='image', type='tensor(float)', shape=[1, 3, 'height', 'width'])
NodeArg(name='mask', type='tensor(float)', shape=[1, 1, 'height', 'width'])
NodeArg(name='alpha_os1', type='tensor(float)', shape=[1, 1, None, None])
NodeArg(name='alpha_os4', type='tensor(float)', shape=[1, 1, None, None])
NodeArg(name='alpha_os8', type='tensor(float)', shape=[1, 1, None, None])

infer_onnx

PYTHONPATH=. python3 ./infer_onnx.py --post-process --onnx ./checkpoints/MGMatting-DIM-100k/latest_model.onnx --image_path ./test_input.jpg --mask_path test_mask.png --output_path ./test_onnx_output.jpg

log:

ONNXRuntime Version:  1.7.0
Namespace(guidance_thres=128, image_path='./test_input.jpg', mask_path='test_mask.png', onnx='./checkpoints/MGMatting-DIM-100k/latest_model.onnx', output_path='./test_onnx_output.jpg', post_process=True)
Load ./checkpoints/MGMatting-DIM-100k/latest_model.onnx Done!
Image: ./test_input.jpg
 Mask: test_mask.png
0.0 1.0
Inference ONNX Done! Saved to ./test_onnx_output.jpg !

the output is:
test_input test_mask test_onnx_output

@yucornetto

yucornetto / MGMatting

Add ONNX Export and MNN/TNN/ONNXRuntime C++ Demo #31