File size: 3,048 Bytes
b4ddbbf | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 |
import sys
import os
import argparse
import hashlib
# Only use sys.path if caffe2onnx isn't installed in your pip/conda env
sys.path.insert(0, "/path/to/cloned/caffe2onnx")
def compute_sha256(filepath):
"""Compute SHA-256 hash of a file for integrity verification."""
sha256 = hashlib.sha256()
with open(filepath, "rb") as f:
for chunk in iter(lambda: f.read(8192), b""):
sha256.update(chunk)
return sha256.hexdigest()
def convert_caffe_to_onnx(prototxt, caffemodel, onnx_path):
"""
Convert a single Caffe model (.prototxt + .caffemodel) to ONNX.
Parameters
----------
prototxt : str — path to Caffe .prototxt (network architecture)
caffemodel : str — path to Caffe .caffemodel (weights)
onnx_path : str — output .onnx file path
"""
from caffe2onnx.src.load_save_model import loadcaffemodel, saveonnxmodel
from caffe2onnx.src.caffe2onnx import Caffe2Onnx
assert os.path.isfile(prototxt), "Prototxt not found: {}".format(prototxt)
assert os.path.isfile(caffemodel), "Caffemodel not found: {}".format(caffemodel)
print("Converting: {} -> {}".format(prototxt, onnx_path))
graph, params = loadcaffemodel(prototxt, caffemodel)
converter = Caffe2Onnx(graph, params, onnx_path)
onnx_model = converter.createOnnxModel()
saveonnxmodel(onnx_model, onnx_path)
sha = compute_sha256(onnx_path)
print(" Saved: {} (SHA-256: {})".format(onnx_path, sha))
return sha
MODELS = {
"detect": ("dnn-models/dnn/wechat_2021-01/detect.prototxt", "dnn-models/dnn/wechat_2021-01/detect.caffemodel"),
"sr": ("dnn-models/dnn/wechat_2021-01/sr.prototxt", "dnn-models/dnn/wechat_2021-01/sr.caffemodel"),
}
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description="Convert WeChatQR Caffe models to ONNX.")
parser.add_argument("--input_dir", default=".",
help="Directory with Caffe model files (default: cwd)")
parser.add_argument("--output_dir", default=".",
help="Directory for output ONNX files (default: cwd)")
args = parser.parse_args()
try:
from caffe2onnx.src.load_save_model import loadcaffemodel, saveonnxmodel
from caffe2onnx.src.caffe2onnx import Caffe2Onnx
except ImportError:
print("Error: caffe2onnx is not installed.")
print(" pip install caffe2onnx")
sys.exit(1)
os.makedirs(args.output_dir, exist_ok=True)
print("=" * 60)
print("WeChatQR Caffe -> ONNX Conversion")
print(" Input : {}".format(args.input_dir))
print(" Output : {}".format(args.output_dir))
print("=" * 60)
for name, (proto_file, caffe_file) in MODELS.items():
proto_path = os.path.join(args.input_dir, proto_file)
caffe_path = os.path.join(args.input_dir, caffe_file)
onnx_path = os.path.join(args.output_dir, name + "_2026april.onnx.onnx")
convert_caffe_to_onnx(proto_path, caffe_path, onnx_path)
print("=" * 60)
print("Done.")
|