convert.py

"""
Convert HAWP-v2 (PyTorch) to a CoreML .mlpackage.

Scope: backbone + head convs only. Produces four dense output tensors
(`heads`, `loi_features`, `loi_features_thin`, `loi_features_aux`) that
downstream code in ``postprocess.py`` turns into lines and junctions.

The HAT-field decoder and line refinement stage use ``torch.unique``,
``.item()``, and NumPy/scipy calls inside ``forward_test``; these do not
survive ``torch.jit.trace`` / ``coremltools.convert`` and are handled
externally. See NOTES.md and the PLAN's "Known issues" for details.

Usage::

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

from __future__ import annotations

import argparse
import os
import sys
from pathlib import Path

import numpy as np
import torch
import torch.nn as nn


HAWP_MEAN = (109.73, 103.832, 98.681)
HAWP_STD = (22.275, 22.124, 23.229)
INPUT_SIZE = 512


class HAWPBackbone(nn.Module):
    """Wraps the HAWP backbone + projection convs in a traceable module.

    Input: raw 0-255 float RGB image of shape (1, 3, 512, 512). The module
    applies the HAWP normalisation internally so the CoreML ImageType can
    forward raw pixel values (scale=1.0, no bias).

    Output: four FP32 tensors ready for ``postprocess.decode``.
    """

    def __init__(self, wireframe_detector: nn.Module) -> None:
        super().__init__()
        self.backbone = wireframe_detector.backbone
        self.fc1 = wireframe_detector.fc1
        self.fc3 = wireframe_detector.fc3
        self.fc4 = wireframe_detector.fc4
        mean = torch.tensor(HAWP_MEAN, dtype=torch.float32).view(1, 3, 1, 1)
        std = torch.tensor(HAWP_STD, dtype=torch.float32).view(1, 3, 1, 1)
        self.register_buffer("mean", mean)
        self.register_buffer("std", std)

    def forward(self, image: torch.Tensor) -> tuple[torch.Tensor, ...]:
        x = (image - self.mean) / self.std
        outputs, features = self.backbone(x)
        heads = outputs[0]
        loi_features = self.fc1(features)
        loi_features_thin = self.fc3(features)
        loi_features_aux = self.fc4(features)
        return heads, loi_features, loi_features_thin, loi_features_aux


def _patch_cuda_in_init() -> None:
    """HAWP's ``WireframeDetector.__init__`` calls ``tspan.cuda()`` unconditionally.
    Neutralise ``.cuda()`` so the model can be built on a CPU-only host."""
    _original = torch.Tensor.cuda

    def _noop(self, *args, **kwargs):  # type: ignore[no-untyped-def]
        return self

    torch.Tensor.cuda = _noop  # type: ignore[assignment]
    return _original  # caller can restore if needed


def build_hawp_v2(ckpt_path: str, hawp_repo: str | None) -> nn.Module:
    if hawp_repo:
        sys.path.insert(0, hawp_repo)

    _patch_cuda_in_init()

    from hawp.fsl.config import cfg
    from hawp.fsl.model.build import build_model

    cfg_path = os.path.join(hawp_repo, "configs", "hawpv2.yaml") if hawp_repo else "configs/hawpv2.yaml"
    cfg.merge_from_file(cfg_path)
    cfg.MODEL.DEVICE = "cpu"

    model = build_model(cfg)
    state = torch.load(ckpt_path, map_location="cpu", weights_only=False)
    model.load_state_dict(state["model"])
    model.eval()
    return model


def trace_and_convert(wrapper: nn.Module, out_path: str) -> None:
    import coremltools as ct

    example = torch.randint(0, 256, (1, 3, INPUT_SIZE, INPUT_SIZE), dtype=torch.float32)
    with torch.no_grad():
        traced = torch.jit.trace(wrapper, example, strict=False)

    # Sanity check the traced graph matches eager execution.
    with torch.no_grad():
        eager_out = wrapper(example)
        traced_out = traced(example)
    for e, t in zip(eager_out, traced_out):
        torch.testing.assert_close(e, t, atol=1e-4, rtol=1e-4)
    print("Trace matches eager output within 1e-4.")

    mlmodel = ct.convert(
        traced,
        inputs=[
            ct.ImageType(
                name="image",
                shape=(1, 3, INPUT_SIZE, INPUT_SIZE),
                scale=1.0,
                bias=[0.0, 0.0, 0.0],
                color_layout=ct.colorlayout.RGB,
            )
        ],
        outputs=[
            ct.TensorType(name="heads"),
            ct.TensorType(name="loi_features"),
            ct.TensorType(name="loi_features_thin"),
            ct.TensorType(name="loi_features_aux"),
        ],
        convert_to="mlprogram",
        compute_precision=ct.precision.FLOAT16,
        compute_units=ct.ComputeUnit.ALL,
        minimum_deployment_target=ct.target.macOS15,
    )

    mlmodel.short_description = (
        "HAWP-v2 backbone (512x512 RGB input). Produces dense feature maps for "
        "external HAT-field decoding; see postprocess.py."
    )
    mlmodel.author = "gsdali (conversion) / Nan Xue et al. (original)"
    mlmodel.license = "MIT"
    mlmodel.version = "1.0.0"

    out_dir = Path(out_path).parent
    out_dir.mkdir(parents=True, exist_ok=True)
    mlmodel.save(out_path)
    print(f"Saved CoreML model to {out_path}")


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--ckpt", required=True, help="Path to hawpv2-edb9b23f.pth")
    parser.add_argument(
        "--hawp-repo",
        default=os.environ.get("HAWP_REPO"),
        help="Path to the cherubicXN/hawp source tree (for configs/).",
    )
    parser.add_argument(
        "--out",
        default="artefacts/hawp_v2_512.mlpackage",
        help="Output .mlpackage path.",
    )
    args = parser.parse_args()

    torch.manual_seed(0)
    np.random.seed(0)

    model = build_hawp_v2(args.ckpt, args.hawp_repo)
    wrapper = HAWPBackbone(model).eval()

    # Save the fc2 / fc2_res / fc2_head weights and tspan buffer so postprocess.py
    # can finish the computation end-to-end without re-building the full model.
    side = {
        "fc2": model.fc2.state_dict(),
        "fc2_res": model.fc2_res.state_dict(),
        "fc2_head": model.fc2_head.state_dict(),
        "tspan": model.tspan.cpu(),
        "loi_cls_type": model.loi_cls_type,
        "n_pts0": model.n_pts0,
        "dim_junction_feature": model.dim_junction_feature,
        "dim_edge_feature": model.dim_edge_feature,
        "dim_fc": model.dim_fc,
        "use_residual": model.use_residual,
        "j2l_threshold": model.j2l_threshold,
        "jhm_threshold": model.jhm_threshold,
        "n_out_junc": model.n_out_junc,
        "topk_junctions": model.topk_junctions,
        "hafm_dis_th": model.hafm_encoder.dis_th,
    }
    side_path = Path(args.out).parent / "hawp_v2_postprocess.pt"
    torch.save(side, side_path)
    print(f"Saved post-processing weights to {side_path}")

    trace_and_convert(wrapper, args.out)


if __name__ == "__main__":
    main()