model conversion to onnx format

Author: clibdev

.gitignore

@@ -2,3 +2,4 @@
 /.idea
 *.pth
 *.onnx
+saved_*.png

README.md

@@ -3,6 +3,8 @@
 Differences between original repository and fork:
 
 * Compatibility with PyTorch >=2.0. (🔥)
+* Original pretrained models and converted ONNX models from GitHub [releases page](https://github.com/clibdev/colorization/releases). (🔥)
+* Model conversion to ONNX format using the [export.py](export.py) file. (🔥)
 * Installation with updated [requirements.txt](requirements.txt) file.
 * Additional command line options for specifying model weights in the [demo_release.py](demo_release.py) file.
 
@@ -12,8 +14,25 @@ Differences between original repository and fork:
 pip install -r requirements.txt
 ```
 
+# Pretrained models
+
+| Name        | Link                                                                                                                                                                                                                             |
+|-------------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
+| ECCV 16     | [PyTorch](https://github.com/clibdev/colorization/releases/latest/download/colorization_release_v2-9b330a0b.pth), [ONNX](https://github.com/clibdev/colorization/releases/latest/download/colorization_release_v2-9b330a0b.onnx) |
+| SIGGRAPH 17 | [PyTorch](https://github.com/clibdev/colorization/releases/latest/download/siggraph17-df00044c.pth), [ONNX](https://github.com/clibdev/colorization/releases/latest/download/siggraph17-df00044c.onnx)                           |
+
 # Inference
 
 ```shell
 python demo_release.py --eccv16_weights colorization_release_v2-9b330a0b.pth --siggraph17_weights siggraph17-df00044c.pth -i imgs/ansel_adams3.jpg
 ```
+
+# Export to ONNX format
+
+```shell
+pip install onnx
+```
+```shell
+python export.py --weights colorization_release_v2-9b330a0b.pth --net_type eccv16
+python export.py --weights siggraph17-df00044c.pth --net_type siggraph17
+```

export.py

@@ -0,0 +1,41 @@
+import argparse
+import os
+from colorizers import *
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument('--weights', type=str, default='./colorization_release_v2-9b330a0b.pth', help='Weights path')
+    parser.add_argument('--net_type', default='eccv16', type=str, help='The network architecture: eccv16 or siggraph17')
+    parser.add_argument('--device', default='cpu', type=str, help='cuda:0 or cpu')
+    args = parser.parse_args()
+
+    if not os.path.exists(args.weights):
+        print('Cannot find weights: {0}'.format(args.weights))
+        exit()
+
+    if args.net_type == 'eccv16':
+        colorizer = eccv16(pretrained=True, weights=args.weights)
+    elif args.net_type == 'siggraph17':
+        colorizer = siggraph17(pretrained=True, weights=args.weights)
+    else:
+        print('Unsupported network type: {0}'.format(args.net_type))
+        exit()
+
+    colorizer.eval()
+    colorizer.to(args.device)
+
+    model_path = os.path.splitext(args.weights)[0] + '.onnx'
+    print(model_path)
+
+    dummy_input = torch.randn(1, 1, 256, 256).to(args.device)
+    torch.onnx.export(
+        colorizer,
+        dummy_input,
+        model_path,
+        verbose=False,
+        input_names=['input'],
+        output_names=['output'],
+        opset_version=18
+    )

requirements.txt

@@ -1,3 +1,6 @@
 torch>=2.0.0
 scikit-image>=0.21.0
 matplotlib>=3.7.0
+
+### Optional
+# onnx>=1.14.0