Fork of ZHKKKe/MODNet

Differences between original repository and fork:

• Compatibility with PyTorch >=2.5. (🔥)
• Original pretrained models and converted ONNX models from GitHub releases page. (🔥)
• Installation with requirements.txt file.
• ONNX Simplifier integration in the export_onnx.py file.
• Minor modifications in the inference.py file.
• The following deprecations and errors has been fixed:
  • FutureWarning: You are using 'torch.load' with 'weights_only=False'.
  • AttributeError: module 'onnx' has no attribute 'load_from_string'.

Installation

pip install -r requirements.txt

Pretrained models

• Download links:

Name	Model Size (MB)	Link	SHA-256
MODNet (Photographic)	25.0 24.7	PyTorch ONNX	7c22235f0925deba15d4d63e53afcb654c47055bbcd98f56e393ab2584007ed8 d381d71c0e85f0edc28df99a6dd59544b0c33cc2bf234b1fc22aa417fec127be
MODNet (Webcam)	25.0 24.7	PyTorch ONNX	913b82b66558db39b6286c150f809017d7528c872b156eb14333c9c6cb52108b 52ace7b5c455f527542b496151996bed429b67ca8efc8bd793485a631363c688

Inference

python -m demo.image_matting.colab.inference --ckpt-path pretrained/modnet-photographic.pt --image-path data/images/test.jpg --output-path result.jpg

Export to ONNX format

pip install onnx onnxsim

python -m onnx_model.export_onnx --ckpt-path pretrained/modnet-photographic.pt --output-path pretrained/modnet-photographic.onnx
python -m onnx_model.export_onnx --ckpt-path pretrained/modnet-webcam.pt --output-path pretrained/modnet-webcam.onnx

ONNX inference

python -m onnx_model.inference_onnx --model-path pretrained/modnet-photographic.onnx --image-path data/images/test.jpg --output-path result.jpg