Add ONNX models for 1.1.0

README.md

@@ -8,12 +8,13 @@ language:
 tags:
 - document-scanning
 - document-dewarping
+- onnx
 - pytorch
 ---
 
 # Flattener
 
-Version 1.0.0 of the models for [Flattener](https://github.com/deveworld/flattener).
+Version 1.1.0 of the models for [Flattener](https://github.com/deveworld/flattener).
 Flattener turns document photos and open books into flat, upright scans.
 Processing runs locally.
 
@@ -21,33 +22,34 @@ Processing runs locally.
 
 | File | Purpose |
 |---|---|
-| `dewarp.pt` | Correct page perspective and curvature |
-| `gate.pt` | Detect the page region |
-| `orient.pt` | Predict text orientation |
-| `bundle.json` | Load the three models together |
+| `dewarp.onnx`, `dewarp.pt` | Correct page perspective and curvature |
+| `gate.onnx`, `gate.pt` | Detect the page region |
+| `orient.onnx`, `orient.pt` | Predict text orientation |
+| `manifest.json` | Load the three ONNX models together |
+| `bundle.json` | Load the three PyTorch checkpoints together |
 
-The checkpoints contain inference weights and model settings.
-The orientation checkpoint also includes its calibrated temperature.
+The ONNX models are exported from the PyTorch checkpoints and hold the same weights.
+The scanner runs them with ONNX Runtime on the CPU by default.
+The PyTorch checkpoints run on a CUDA GPU and are the starting point for training.
+`manifest.json` and the orientation checkpoint include the calibrated orientation temperature and threshold.
 
 ## Use
 
-From a Flattener 1.0.0 source checkout:
-
 ```bash
-uv sync --python 3.12
-uv run python -m flattener.scan photo.jpg --out scans
+uvx flattener-scan -i photo.jpg scan.jpg
 ```
 
-The first scan downloads this bundle at revision `1.0.0`.
+The first scan downloads the ONNX models at revision `1.1.0`.
 Later scans reuse the local Hugging Face cache, including offline.
+With PyTorch installed and a CUDA GPU present (`uvx --with torch flattener-scan ...`), the scanner downloads and uses the PyTorch checkpoints instead.
 
 Download the files to a chosen directory with:
 
 ```bash
-hf download DevWorld/flattener bundle.json dewarp.pt gate.pt orient.pt --revision 1.0.0 --local-dir models
+hf download DevWorld/flattener manifest.json dewarp.onnx gate.onnx orient.onnx --revision 1.1.0 --local-dir models
 ```
 
-Then pass `--bundle models/bundle.json` to the scanner.
+Then pass `--bundle models/manifest.json` to the scanner.
 
 ## Training and use
 

dewarp.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:34dae2641b43d1f6a7961b6842517e6d22f4115f711abcc9f0bc9c5c8b7b2604
+size 31602437

gate.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9c2391d02701f59998c6b7f5e878c62e63d714ad9b42e43ccf128108589bf322
+size 2254621

manifest.json

@@ -0,0 +1,178 @@
+{
+  "schema": "flattener-web-bundle-1",
+  "version": "1.1.0",
+  "precision": "float32",
+  "opset": 18,
+  "models": {
+    "dewarp": {
+      "file": "dewarp.onnx",
+      "bytes": 31602437,
+      "inputs": {
+        "image": [
+          1,
+          3,
+          712,
+          488
+        ]
+      },
+      "outputs": [
+        "grid"
+      ],
+      "input_size": [
+        712,
+        488
+      ],
+      "grid_size": [
+        45,
+        31
+      ],
+      "grid_convention": "flattener",
+      "tta": "flip",
+      "model_config": {
+        "input_size": [
+          712,
+          488
+        ],
+        "grid_size": [
+          45,
+          31
+        ],
+        "num_filter": 32,
+        "kernel_size": 5,
+        "with_3d": true,
+        "grid_convention": "flattener",
+        "architecture": "uvdocnet"
+      }
+    },
+    "gate": {
+      "file": "gate.onnx",
+      "bytes": 2254621,
+      "inputs": {
+        "image": [
+          1,
+          3,
+          192,
+          144
+        ]
+      },
+      "outputs": [
+        "probability"
+      ],
+      "input_size": [
+        192,
+        144
+      ],
+      "model_config": {
+        "input_size": [
+          192,
+          144
+        ],
+        "channels": [
+          16,
+          32,
+          64,
+          128
+        ],
+        "model_version": "region-context-unet-2.0"
+      }
+    },
+    "orient": {
+      "file": "orient.onnx",
+      "bytes": 1780334,
+      "inputs": {
+        "global_view": [
+          1,
+          1,
+          256,
+          256
+        ],
+        "patches": [
+          1,
+          4,
+          1,
+          128,
+          128
+        ]
+      },
+      "outputs": [
+        "logits",
+        "probabilities"
+      ],
+      "config": {
+        "global_size": 256,
+        "patch_size": 128,
+        "patches": 4,
+        "page_long_side": 1024,
+        "energy_grid": 8,
+        "pad_value": 128,
+        "stem": 16,
+        "blocks": [
+          [
+            32,
+            1,
+            3,
+            false
+          ],
+          [
+            48,
+            2,
+            3,
+            false
+          ],
+          [
+            48,
+            1,
+            3,
+            false
+          ],
+          [
+            96,
+            2,
+            3,
+            false
+          ],
+          [
+            96,
+            1,
+            3,
+            false
+          ],
+          [
+            192,
+            2,
+            5,
+            false
+          ],
+          [
+            192,
+            1,
+            5,
+            false
+          ],
+          [
+            192,
+            1,
+            5,
+            true
+          ],
+          [
+            320,
+            2,
+            5,
+            true
+          ],
+          [
+            320,
+            1,
+            5,
+            true
+          ]
+        ],
+        "hidden": 128,
+        "dropout": 0.1
+      },
+      "temperature": 0.5192675672115961,
+      "threshold": 0.9
+    }
+  }
+}

orient.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9632f72031f8d54b4322110a4674f30a1a070d05e29843dcefb4dbbbe65aad69
+size 1780334