singam96/ShadeNet-5-35M: upload hf_draft

README.md

@@ -94,7 +94,7 @@ The six-pass average is better than a single pass but costs six forward passes.
 - **Against the labels:** see the Results table: normals recover about 94% of the six-pass gain, albedo about 60%, depth about 8%. The refiner does not reproduce the six-pass average's changes to the overall depth levels, so depth is essentially the single pass's.
 - **Cost:** about 12 ms per image on top of the single pass (PyTorch, A4000) and about 0.26 s on a CPU for the 448x448 ONNX version.
 
-Files: `refiner/refiner.pt` (PyTorch weights) and `onnx/shadenet5_refiner.onnx` (fp32, fixed 448x448). Run it after the main model: `python inference/run_onnx.py photos/ --refine`.
+Files: `refiner/refiner.pt` (PyTorch weights), `onnx/shadenet5_refiner.onnx` (the refiner alone, fp32, fixed 448x448) and `onnx/shadenet5_mix_refined.onnx` (35 MB: the MIX model, the refiner, the clamp and the unit-length normals in **one file**, photo in and final maps out; its mean L1 is identical to running the two files one after the other). Run them with `python inference/run_onnx.py photos/ --refine` (two files) or `--merged` (the single file).
 
 On the shipped 4/8-bit MIX ONNX model (CPU, same 23 images, unit-length normals): single pass 0.1755, single pass + refiner **0.1676**, six passes 0.1645. The refiner was trained on full-precision outputs, so this is the number for the file you would download.
 
@@ -275,6 +275,16 @@ python inference/run_onnx.py photos/ --refine --out out/
 
 Each photo becomes a panel image: input, albedo, depth, normal. `--refine` runs the single pass and then the small refiner (about 12 ms more; the recommended mode, see The refiner). `--enhanced` is the six-pass post-processing described above in one switch (mirror average, rotations of ±4 degrees, unit-length normals; six forward passes); leave both flags out for the plain single pass. The whole photo is kept: its long side is scaled to 448 and the rest of the 448x448 input is filled by mirroring the photo's own edge, then the padding is cut off the maps again, so the outputs have the photo's own shape. A photo that is not square therefore uses fewer than 448 pixels on its short side. (The PyTorch model runs at 448 and 896 square; the ONNX files are fixed at 448x448.) `torch` is only used for the rotation resampling of `--enhanced`. Set `--threads` to the number of CPU cores you have: more threads than cores is slower, not faster. A hosted demo runs the same code: [ShadeNet-5 Space](https://huggingface.co/spaces/singam96/ShadeNet-5).
 
+Or call the single merged file directly:
+
+```python
+import numpy as np, onnxruntime as ort
+sess = ort.InferenceSession("shadenet5_mix_refined.onnx")
+# photo: float32 [1, 3, 448, 448] in [-1, 1]
+maps = sess.run(None, {"input_rgb": photo})[0]   # [1, 7, 448, 448]
+albedo, depth, normal = maps[:, 0:3], maps[:, 3:4], maps[:, 4:7]
+```
+
 The PyTorch checkpoint (`checkpoints/best-l1-pruned10-ep19.ckpt`) is the full-precision reference. Loading it needs the training code, which is not published yet (the `--enhanced` switch of its `infer.py` is the same recipe), so use the ONNX file unless you have that code.
 
 `config.json` in this repository lists the architecture, the input and output convention (7 maps in [-1, 1]: albedo, depth, normal), the supported sizes and the recommended inference settings.

config.json

@@ -47,7 +47,8 @@
   "onnx": {
     "onnx/shadenet5_mix_int4_int8.onnx": "ViT fc1/fc2 int8, everything else int4",
     "onnx/shadenet5_mix2_int4_int8.onnx": "ViT fc2 int8, everything else int4",
-    "onnx/shadenet5_refiner.onnx": "the small refiner, fp32, fixed 448x448 (inputs photo and maps, output refined maps)"
+    "onnx/shadenet5_refiner.onnx": "the small refiner, fp32, fixed 448x448 (inputs photo and maps, output refined maps)",
+    "onnx/shadenet5_mix_refined.onnx": "the MIX model + the refiner + clamp + unit-length normals in one file: input `input_rgb` [1,3,448,448] in [-1,1], output `refined` [1,7,448,448]"
   },
   "recommended_inference": {
     "mode": "refine",
@@ -58,7 +59,9 @@
       "unit_normals": true,
       "forward_passes": 6,
       "command": "python inference/run_onnx.py <photos> --enhanced --out out/"
-    }
+    },
+    "single_file": "onnx/shadenet5_mix_refined.onnx",
+    "single_file_command": "python inference/run_onnx.py <photos> --merged --out out/"
   },
   "validation_mean_l1": {
     "single_pass_unit_normals": 0.1757,
@@ -66,7 +69,8 @@
     "six_passes": 0.1628,
     "mix_onnx_single": 0.1755,
     "mix_onnx_plus_refiner": 0.1676,
-    "mix_onnx_six_passes": 0.1645
+    "mix_onnx_six_passes": 0.1645,
+    "mix_onnx_single_merged_file": 0.1676
   },
   "refiner": {
     "pytorch": "refiner/refiner.pt",

inference/run_onnx.py

@@ -27,6 +27,7 @@ from framing import crop_to_box, display_size, letterbox, resize_maps  # noqa: E
 REPO = "singam96/ShadeNet-5-35M"
 ONNX_FILE = "onnx/shadenet5_mix_int4_int8.onnx"
 REFINER_FILE = "onnx/shadenet5_refiner.onnx"
+MERGED_FILE = "onnx/shadenet5_mix_refined.onnx"
 SIZE = 448                       # the ONNX files are fixed at 448x448
 ROT_DEG = 4.0
 
@@ -50,6 +51,9 @@ def main():
     ap.add_argument("--onnx", default=None, help="local .onnx (default: download)")
     ap.add_argument("--size", type=int, default=SIZE, choices=[448],
                     help="the ONNX files are fixed at 448x448")
+    ap.add_argument("--merged", action="store_true",
+                    help="the single file with the MIX model, refiner and unit normals in one graph "
+                         "(same result as --refine, one session)")
     ap.add_argument("--refine", action="store_true",
                     help="single pass, then the small refiner (recommended)")
     ap.add_argument("--refiner-onnx", default=None, help="local refiner .onnx (default: download)")
@@ -60,13 +64,13 @@ def main():
     ap.add_argument("--threads", type=int, default=8,
                     help="ONNX Runtime threads (more than your cores is slower)")
     args = ap.parse_args()
-    if args.refine and args.enhanced:
-        ap.error('use --refine or --enhanced, not both')
+    if sum([args.refine, args.enhanced, args.merged]) > 1:
+        ap.error('use one of --refine, --enhanced, --merged')
 
     path = args.onnx
     if path is None:
         from huggingface_hub import hf_hub_download
-        path = hf_hub_download(REPO, ONNX_FILE)
+        path = hf_hub_download(REPO, MERGED_FILE if args.merged else ONNX_FILE)
     so = ort.SessionOptions()
     so.intra_op_num_threads = args.threads
     sess = ort.InferenceSession(path, so, providers=["CPUExecutionProvider"])
@@ -94,7 +98,9 @@ def main():
         square, box = letterbox(img, args.size)
         x = to_tensor(square)
         with torch.no_grad():
-            if args.refine:
+            if args.merged:
+                maps = refine(x)            # already refined, clamped, unit-length normals
+            elif args.refine:
                 single = refine(x)
                 out = ref_sess.run(None, {"photo": np.ascontiguousarray(x.numpy()),
                                           "maps": np.ascontiguousarray(single.numpy())})[0]

onnx/shadenet5_mix_refined.onnx

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

refiner/eval_onnx_mix_cpu.json

@@ -1,23 +1,30 @@
 {
  "MIX single pass": {
-  "albedo": 0.1782670166829358,
-  "depth": 0.17611026375190073,
-  "normal": 0.17200583511072656,
-  "mean": 0.17546103851518768,
-  "ms": 1041.167746890675
+  "albedo": 0.1782614407332047,
+  "depth": 0.17611568907032843,
+  "normal": 0.1720099980416505,
+  "mean": 0.17546237594839462,
+  "ms": 1329.7054225748234
  },
  "MIX + refiner": {
-  "albedo": 0.17474121535601822,
-  "depth": 0.17591332482255023,
-  "normal": 0.15225264818772025,
-  "mean": 0.1676357294554296,
-  "ms": 1300.2352497794413
+  "albedo": 0.17473657526399777,
+  "depth": 0.1759216370789901,
+  "normal": 0.15225956880527994,
+  "mean": 0.16763926038275595,
+  "ms": 1649.118575182828
  },
  "MIX six passes": {
-  "albedo": 0.1728826994481294,
-  "depth": 0.17006371332251508,
-  "normal": 0.1506343720399815,
-  "mean": 0.16452692827020868,
-  "ms": 6070.682460611517
+  "albedo": 0.17288330618454062,
+  "depth": 0.1700636342815731,
+  "normal": 0.15063375072634738,
+  "mean": 0.1645268970641537,
+  "ms": 8918.771418658169
+ },
+ "one merged file": {
+  "albedo": 0.17473587231791537,
+  "depth": 0.175918893645639,
+  "normal": 0.15225798798644025,
+  "mean": 0.16763758464999826,
+  "ms": 1721.9356731934981
  }
 }