coreml-hawp

HAWP-v2 line and junction detector converted to CoreML for Apple Silicon.

License attribution

This repository converts HAWP-v2 by Nan Xue et al. to CoreML format.

Original work: HAWP (Holistically Attracted Wireframe Parser) Repository: https://github.com/cherubicXN/hawp Paper: "HAWP: Holistically-Attracted Wireframe Parsing" (TPAMI 2023) Authors: Nan Xue, Tianfu Wu, Song Bai, Fudong Wang, Gui-Song Xia, Liangpei Zhang, Philip H.S. Torr License: MIT (see LICENSE file)

This conversion repository is also MIT-licensed. The CoreML artefacts in artefacts/ are derived from the upstream PyTorch model weights.

What's shipped

path	contents
artefacts/hawp_v2_512.mlpackage	FP16 CoreML program, input 512×512 RGB, backbone-only (see "Model split")
artefacts/hawp_v2_postprocess.pt	PyTorch state dict for fc2/fc2_res/fc2_head + HAT decoding hyperparams
convert.py	PyTorch → CoreML conversion script
postprocess.py	PyTorch module that turns the 4 backbone tensors into lines + junctions
test_ane.swift	Swift smoke test; loads the mlpackage and runs one inference
benchmarks.json	M4 latency table + per-op ANE/GPU/CPU dispatch + PyTorch-parity diff
test_images/	Three reference drawings; expected_output.json is the PyTorch baseline

Model split (important)

HAWP-v2's end-to-end forward_test cannot be traced into a single mlprogram. Three pieces block conversion:

1. post_jheatmap uses scipy.argsort + NumPy with a .cpu().numpy() call inside the forward pass.
2. torch.unique(..., return_inverse=True, dim=0) is not in CoreML MIL.
3. topk = int((jloc_nms > th).float().sum().item()) feeds .item() back into the graph, which is a data-dependent break.

The shipped mlpackage therefore covers the backbone + three projection 1×1 convs. It outputs four dense tensors:

output	shape	role
heads	(1, 9, 128, 128)	md (3) · dis (1) · res (1) · jloc (2) · joff (2)
loi_features	(1, 128, 128, 128)	128-ch feature map for LOI pooling (junction endpoints)
loi_features_thin	(1, 4, 128, 128)	4-ch feature map for sampled line midpoints
loi_features_aux	(1, 4, 128, 128)	4-ch auxiliary feature map for raw HAT-decoded lines

The HAT-field decoding + line refinement ships as postprocess.py (pure PyTorch, CPU, ~30-160 ms depending on line count). This matches PLAN.md's documented fallback and keeps 99.7% of compute ops on the Apple Neural Engine.

Quick start (Python)

import coremltools as ct
import cv2, torch
from PIL import Image
from postprocess import HAWPDecoder

# 1. Load the two artefacts
ml = ct.models.MLModel("artefacts/hawp_v2_512.mlpackage")
decoder = HAWPDecoder("artefacts/hawp_v2_postprocess.pt")

# 2. Preprocess: BGR → RGB → 512x512 PIL image
img_bgr = cv2.imread("test_images/telefono.png")
rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
pil = Image.fromarray(cv2.resize(rgb, (512, 512)), mode="RGB")

# 3. Forward: ANE backbone + CPU decoder
tensors = ml.predict({"image": pil})
out = decoder(
    heads=torch.from_numpy(tensors["heads"]),
    loi_features=torch.from_numpy(tensors["loi_features"]),
    loi_features_thin=torch.from_numpy(tensors["loi_features_thin"]),
    loi_features_aux=torch.from_numpy(tensors["loi_features_aux"]),
    orig_height=img_bgr.shape[0],
    orig_width=img_bgr.shape[1],
)

# out["lines_pred"]: (N, 4)  xyxy in original image pixels
# out["lines_score"]: (N,)   in [0, 1]
# out["juncs_pred"]: (M, 2)  xy in original image pixels
# out["juncs_score"]: (M,)   in [0, 1]
print(f"{len(out['lines_pred'])} lines, {len(out['juncs_pred'])} junctions")

Swift usage

test_ane.swift demonstrates loading and a 20-iteration latency loop.

swift test_ane.swift test_images/telefono.png ane

Full decoding in Swift is not yet implemented — the post-processing tensors are compact enough that the expected pattern is to hand them off to a Python or C++ implementation of the HAWP matcher.

Re-converting from scratch

python -m venv .venv && . .venv/bin/activate
pip install -r requirements.txt

git clone https://github.com/cherubicXN/hawp.git
(cd hawp && git checkout 92ae446)
pip install --no-deps -e hawp
curl -LO https://github.com/cherubicXN/hawp-torchhub/releases/download/HAWPv2/hawpv2-edb9b23f.pth

python convert.py --ckpt hawpv2-edb9b23f.pth --hawp-repo hawp \
                  --out artefacts/hawp_v2_512.mlpackage

Benchmarks (Apple M4, 32 GB, macOS 15, coremltools 9.0, FP16)

compute units	mean ms	p50 ms	p90 ms
ANE (cpuAndNeuralEngine)	15.7	15.7	16.0
ALL (scheduler → ANE)	16.0	15.9	16.6
GPU (cpuAndGPU)	40.0	39.9	41.3
CPU only	84.1	84.0	91.9

Per-op dispatch (from MLComputePlan): 322 of 323 non-const ops run on the Apple Neural Engine; the single CPU op is an output cast. Post-processing on CPU (postprocess.py) adds 30-160 ms depending on the number of candidate lines.

See benchmarks.json for full numbers including PyTorch parity diff.

Source pin

• Upstream: cherubicXN/hawp commit 92ae446 (main branch, Feb 2024)
• Checkpoint: hawpv2-edb9b23f.pth, SHA-256 can be verified via shasum -a 256 against the file pulled from the upstream Releases page.