Views
No views yet
Files and versions tab. You need both the .onnx and the .onnx.data files to inference the model.git clone https://github.com/hellozhuo/pidinet.git1"""
2Export a PiDiNet checkpoint to ONNX.
3
4Example:
5python pidinet_to_onnx.py \
6 --checkpoint table5_pidinet.pth \
7 --output pidinet_table5.onnx \
8 --config carv4 --sa --dil --height 512 --width 512
9"""
10
11import argparse
12from types import SimpleNamespace
13import torch
14
15from pidinet.models import (
16 pidinet_converted,
17 pidinet_small_converted,
18 pidinet_tiny_converted,
19)
20from pidinet.models.convert_pidinet import convert_pidinet
21
22
23MODEL_BUILDERS = {
24 "base": pidinet_converted,
25 "small": pidinet_small_converted,
26 "tiny": pidinet_tiny_converted,
27}
28
29
30def build_model(config: str, sa: bool, dil: bool, variant: str) -> torch.nn.Module:
31 """Create the converted PiDiNet model (uses vanilla convs)."""
32 if variant not in MODEL_BUILDERS:
33 raise ValueError(f"Unsupported variant '{variant}' (choose from {list(MODEL_BUILDERS)})")
34
35 args = SimpleNamespace(config=config, sa=sa, dil=dil)
36 return MODEL_BUILDERS[variant](args)
37
38
39def _read_checkpoint(ckpt_path: str):
40 checkpoint = torch.load(ckpt_path, map_location="cpu")
41 state = checkpoint.get("state_dict", checkpoint)
42 return _strip_module_prefix(state)
43
44
45def _infer_flags_from_state(state_dict):
46 """Infer sa/dil from checkpoint contents."""
47 has_sa = any(k.startswith("attentions.") for k in state_dict)
48 has_dil = any(k.startswith("dilations.") for k in state_dict)
49 return has_sa, has_dil
50
51
52def _strip_module_prefix(state_dict):
53 """Remove a leading 'module.' (from DataParallel) if present."""
54 if not any(k.startswith("module.") for k in state_dict.keys()):
55 return state_dict
56 return {k.replace("module.", "", 1): v for k, v in state_dict.items()}
57
58
59def export_onnx(model, dummy, output_path: str, opset: int):
60 output_names = ["side1", "side2", "side3", "side4", "fused"]
61 dynamic_axes = {
62 "image": {0: "batch", 2: "height", 3: "width"},
63 "side1": {0: "batch", 2: "height", 3: "width"},
64 "side2": {0: "batch", 2: "height", 3: "width"},
65 "side3": {0: "batch", 2: "height", 3: "width"},
66 "side4": {0: "batch", 2: "height", 3: "width"},
67 "fused": {0: "batch", 2: "height", 3: "width"},
68 }
69 torch.onnx.export(
70 model,
71 dummy,
72 output_path,
73 opset_version=opset,
74 input_names=["image"],
75 output_names=output_names,
76 dynamic_axes=dynamic_axes,
77 do_constant_folding=True,
78 )
79
80
81def parse_args():
82 parser = argparse.ArgumentParser(description="Convert PiDiNet checkpoint to ONNX.")
83 parser.add_argument(
84 "--checkpoint",
85 type=str,
86 default="pidinet_model/table5_pidinet.pth",
87 help="Path to PiDiNet checkpoint (.pth).",
88 )
89 parser.add_argument(
90 "--output",
91 type=str,
92 default="pidinet_table5.onnx",
93 help="Path to write ONNX file.",
94 )
95 parser.add_argument(
96 "--config",
97 type=str,
98 default="carv4",
99 help="Model config name (see pidinet/models/config.py).",
100 )
101 parser.add_argument("--sa", action="store_true", help="Use CSAM.")
102 parser.add_argument("--dil", action="store_true", help="Use CDCM.")
103 parser.add_argument("--height", type=int, default=512, help="Dummy input height.")
104 parser.add_argument("--width", type=int, default=512, help="Dummy input width.")
105 parser.add_argument("--batch", type=int, default=1, help="Dummy batch size.")
106 parser.add_argument(
107 "--opset",
108 type=int,
109 default=18,
110 help="ONNX opset version (>=18 recommended to avoid converter errors).",
111 )
112 parser.add_argument(
113 "--cuda",
114 action="store_true",
115 help="Export with the model on CUDA (optional).",
116 )
117 parser.add_argument(
118 "--variant",
119 choices=["base", "small", "tiny"],
120 default="base",
121 help="Width of the PiDiNet: 'base' (table5_pidinet), 'small' (table5_pidinet-small), or 'tiny' (table5_pidinet-tiny).",
122 )
123 parser.add_argument(
124 "--strict-flags",
125 action="store_true",
126 help="Do not auto-adjust --sa/--dil based on checkpoint contents.",
127 )
128 return parser.parse_args()
129
130
131def main():
132 args = parse_args()
133
134 raw_state = _read_checkpoint(args.checkpoint)
135 inferred_sa, inferred_dil = _infer_flags_from_state(raw_state)
136
137 sa = inferred_sa or args.sa
138 dil = inferred_dil or args.dil
139 if not args.strict_flags:
140 if args.sa and not inferred_sa:
141 print("Checkpoint lacks attention layers; disabling --sa for this export.")
142 sa = False
143 if args.dil and not inferred_dil:
144 print("Checkpoint lacks dilation modules; disabling --dil for this export.")
145 dil = False
146
147 device = torch.device("cuda" if args.cuda and torch.cuda.is_available() else "cpu")
148 print(f"Export settings -> variant: {args.variant}, sa: {sa}, dil: {dil}, config: {args.config}")
149 model = build_model(args.config, sa, dil, args.variant)
150 model.load_state_dict(convert_pidinet(raw_state, args.config))
151 model.eval().to(device)
152
153 dummy = torch.randn(args.batch, 3, args.height, args.width, device=device)
154 export_onnx(model, dummy, args.output, args.opset)
155
156 print(f"Exported ONNX to {args.output}")
157
158
159if __name__ == "__main__":
160 main()1"""
2Run the PiDiNet ONNX model on one image and save the fused edge map.
3
4Example:
5python test_pidinet_onnx.py \
6 --onnx model_PIDINET/pidinet_table5.onnx \
7 --image Images/example.jpg \
8 --output Results/example_edges.png
9"""
10
11import argparse
12from pathlib import Path
13
14import numpy as np
15import onnxruntime as ort
16from PIL import Image
17
18
19MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)[:, None, None]
20STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)[:, None, None]
21
22
23def preprocess(img_path: Path) -> np.ndarray:
24 img = Image.open(img_path).convert("RGB")
25 arr = np.asarray(img, dtype=np.float32) / 255.0 # HWC in [0,1]
26 arr = arr.transpose(2, 0, 1) # CHW
27 arr = (arr - MEAN) / STD
28 return arr[None, ...] # BCHW
29
30
31def postprocess(edge_map: np.ndarray, out_path: Path):
32 out_path.parent.mkdir(parents=True, exist_ok=True)
33 edge_map = np.clip(edge_map, 0.0, 1.0)
34 edge_img = (edge_map * 255.0).astype(np.uint8)
35 Image.fromarray(edge_img).save(out_path)
36
37
38def parse_args():
39 parser = argparse.ArgumentParser(description="Test PiDiNet ONNX on a single image.")
40 parser.add_argument(
41 "--onnx",
42 type=Path,
43 default=Path("model_PIDINET/pidinet_table5.onnx"),
44 help="Path to the PiDiNet ONNX file.",
45 )
46 parser.add_argument(
47 "--image",
48 type=Path,
49 required=True,
50 help="Input image path.",
51 )
52 parser.add_argument(
53 "--output",
54 type=Path,
55 default=Path("Results/pidinet_edges.png"),
56 help="Where to save the fused edge map.",
57 )
58 parser.add_argument(
59 "--provider",
60 type=str,
61 default="CPUExecutionProvider",
62 help="ONNX Runtime provider (e.g., CPUExecutionProvider or CUDAExecutionProvider).",
63 )
64 return parser.parse_args()
65
66
67def main():
68 args = parse_args()
69
70 session = ort.InferenceSession(
71 str(args.onnx),
72 providers=[args.provider],
73 )
74
75 inp = preprocess(args.image)
76 outputs = session.run(None, {"image": inp})
77
78 fused = np.array(outputs[-1])[0, 0] # fused edge map
79 postprocess(fused, args.output)
80
81 print(f"Saved edge map to {args.output}")
82
83
84if __name__ == "__main__":
85 main()
86