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:
- download the model from https://github.com/hellozhuo/pidinet#:~:text=vary%20too%20much)%3A-,Model,-ODS
- Git clone the pidinet repo
git clone https://github.com/hellozhuo/pidinet.git - run the following code via CLI:
"""
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:
"""
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()