Add : model conversion script

wechat_qr/model_conversion.py

@@ -0,0 +1,86 @@
+
+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.")