README.md

ONNX models available in the `Files and versions` tab. You need both the .onnx and the .onnx.data files to inference the model.

## 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 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
"""
Run the PiDiNet ONNX model on one image and save the fused edge map.

Example:
python test_pidinet_onnx.py \
  --onnx model_PIDINET/pidinet_table5.onnx \
  --image Images/example.jpg \
  --output Results/example_edges.png
"""

import argparse
from pathlib import Path

import numpy as np
import onnxruntime as ort
from PIL import Image


MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)[:, None, None]
STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)[:, None, None]


def preprocess(img_path: Path) -> np.ndarray:
    img = Image.open(img_path).convert("RGB")
    arr = np.asarray(img, dtype=np.float32) / 255.0          # HWC in [0,1]
    arr = arr.transpose(2, 0, 1)                             # CHW
    arr = (arr - MEAN) / STD
    return arr[None, ...]                                    # BCHW


def postprocess(edge_map: np.ndarray, out_path: Path):
    out_path.parent.mkdir(parents=True, exist_ok=True)
    edge_map = np.clip(edge_map, 0.0, 1.0)
    edge_img = (edge_map * 255.0).astype(np.uint8)
    Image.fromarray(edge_img).save(out_path)


def parse_args():
    parser = argparse.ArgumentParser(description="Test PiDiNet ONNX on a single image.")
    parser.add_argument(
        "--onnx",
        type=Path,
        default=Path("model_PIDINET/pidinet_table5.onnx"),
        help="Path to the PiDiNet ONNX file.",
    )
    parser.add_argument(
        "--image",
        type=Path,
        required=True,
        help="Input image path.",
    )
    parser.add_argument(
        "--output",
        type=Path,
        default=Path("Results/pidinet_edges.png"),
        help="Where to save the fused edge map.",
    )
    parser.add_argument(
        "--provider",
        type=str,
        default="CPUExecutionProvider",
        help="ONNX Runtime provider (e.g., CPUExecutionProvider or CUDAExecutionProvider).",
    )
    return parser.parse_args()


def main():
    args = parse_args()

    session = ort.InferenceSession(
        str(args.onnx),
        providers=[args.provider],
    )

    inp = preprocess(args.image)
    outputs = session.run(None, {"image": inp})

    fused = np.array(outputs[-1])[0, 0]  # fused edge map
    postprocess(fused, args.output)

    print(f"Saved edge map to {args.output}")


if __name__ == "__main__":
    main()

```