Update README.md

README.md

@@ -1,4 +1,172 @@
-How do inference the pidinet onnx:
+## How to convert to ONNX:
+1) download the model from https://github.com/hellozhuo/pidinet#:~:text=vary%20too%20much)%3A-,Model,-ODS
+2) Git clone the pidinet repo `git clone https://github.com/hellozhuo/pidinet.git`
+3) run the following code via CLI:
+
+```python
+"""
+Export a PiDiNet checkpoint to ONNX.
+
+Example:
+python pidinet_to_onnx.py \
+  --checkpoint pidinet_model/table5_pidinet.pth \
+  --output pidinet_table5.onnx \
+  --config carv4 --sa --dil --height 512 --width 512
+"""
+
+import argparse
+from types import SimpleNamespace
+import torch
+
+from pidinet.models import (
+    pidinet_converted,
+    pidinet_small_converted,
+    pidinet_tiny_converted,
+)
+from pidinet.models.convert_pidinet import convert_pidinet
+
+
+MODEL_BUILDERS = {
+    "base": pidinet_converted,
+    "small": pidinet_small_converted,
+    "tiny": pidinet_tiny_converted,
+}
+
+
+def build_model(config: str, sa: bool, dil: bool, variant: str) -> torch.nn.Module:
+    """Create the converted PiDiNet model (uses vanilla convs)."""
+    if variant not in MODEL_BUILDERS:
+        raise ValueError(f"Unsupported variant '{variant}' (choose from {list(MODEL_BUILDERS)})")
+
+    args = SimpleNamespace(config=config, sa=sa, dil=dil)
+    return MODEL_BUILDERS[variant](args)
+
+
+def _read_checkpoint(ckpt_path: str):
+    checkpoint = torch.load(ckpt_path, map_location="cpu")
+    state = checkpoint.get("state_dict", checkpoint)
+    return _strip_module_prefix(state)
+
+
+def _infer_flags_from_state(state_dict):
+    """Infer sa/dil from checkpoint contents."""
+    has_sa = any(k.startswith("attentions.") for k in state_dict)
+    has_dil = any(k.startswith("dilations.") for k in state_dict)
+    return has_sa, has_dil
+
+
+def _strip_module_prefix(state_dict):
+    """Remove a leading 'module.' (from DataParallel) if present."""
+    if not any(k.startswith("module.") for k in state_dict.keys()):
+        return state_dict
+    return {k.replace("module.", "", 1): v for k, v in state_dict.items()}
+
+
+def export_onnx(model, dummy, output_path: str, opset: int):
+    output_names = ["side1", "side2", "side3", "side4", "fused"]
+    dynamic_axes = {
+        "image": {0: "batch", 2: "height", 3: "width"},
+        "side1": {0: "batch", 2: "height", 3: "width"},
+        "side2": {0: "batch", 2: "height", 3: "width"},
+        "side3": {0: "batch", 2: "height", 3: "width"},
+        "side4": {0: "batch", 2: "height", 3: "width"},
+        "fused": {0: "batch", 2: "height", 3: "width"},
+    }
+    torch.onnx.export(
+        model,
+        dummy,
+        output_path,
+        opset_version=opset,
+        input_names=["image"],
+        output_names=output_names,
+        dynamic_axes=dynamic_axes,
+        do_constant_folding=True,
+    )
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Convert PiDiNet checkpoint to ONNX.")
+    parser.add_argument(
+        "--checkpoint",
+        type=str,
+        default="pidinet_model/table5_pidinet.pth",
+        help="Path to PiDiNet checkpoint (.pth).",
+    )
+    parser.add_argument(
+        "--output",
+        type=str,
+        default="pidinet_table5.onnx",
+        help="Path to write ONNX file.",
+    )
+    parser.add_argument(
+        "--config",
+        type=str,
+        default="carv4",
+        help="Model config name (see pidinet/models/config.py).",
+    )
+    parser.add_argument("--sa", action="store_true", help="Use CSAM.")
+    parser.add_argument("--dil", action="store_true", help="Use CDCM.")
+    parser.add_argument("--height", type=int, default=512, help="Dummy input height.")
+    parser.add_argument("--width", type=int, default=512, help="Dummy input width.")
+    parser.add_argument("--batch", type=int, default=1, help="Dummy batch size.")
+    parser.add_argument(
+        "--opset",
+        type=int,
+        default=18,
+        help="ONNX opset version (>=18 recommended to avoid converter errors).",
+    )
+    parser.add_argument(
+        "--cuda",
+        action="store_true",
+        help="Export with the model on CUDA (optional).",
+    )
+    parser.add_argument(
+        "--variant",
+        choices=["base", "small", "tiny"],
+        default="base",
+        help="Width of the PiDiNet: 'base' (table5_pidinet), 'small' (table5_pidinet-small), or 'tiny' (table5_pidinet-tiny).",
+    )
+    parser.add_argument(
+        "--strict-flags",
+        action="store_true",
+        help="Do not auto-adjust --sa/--dil based on checkpoint contents.",
+    )
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+
+    raw_state = _read_checkpoint(args.checkpoint)
+    inferred_sa, inferred_dil = _infer_flags_from_state(raw_state)
+
+    sa = inferred_sa or args.sa
+    dil = inferred_dil or args.dil
+    if not args.strict_flags:
+        if args.sa and not inferred_sa:
+            print("Checkpoint lacks attention layers; disabling --sa for this export.")
+            sa = False
+        if args.dil and not inferred_dil:
+            print("Checkpoint lacks dilation modules; disabling --dil for this export.")
+            dil = False
+
+    device = torch.device("cuda" if args.cuda and torch.cuda.is_available() else "cpu")
+    print(f"Export settings -> variant: {args.variant}, sa: {sa}, dil: {dil}, config: {args.config}")
+    model = build_model(args.config, sa, dil, args.variant)
+    model.load_state_dict(convert_pidinet(raw_state, args.config))
+    model.eval().to(device)
+
+    dummy = torch.randn(args.batch, 3, args.height, args.width, device=device)
+    export_onnx(model, dummy, args.output, args.opset)
+
+    print(f"Exported ONNX to {args.output}")
+
+
+if __name__ == "__main__":
+    main()
+```
+
+## How do inference the pidinet onnx:
 
 ```python
 """