README.md

metadata

license: apache-2.0
base_model:
  - cherubicxn/scalelsd
tags:
  - mlx
  - mlx-swift
  - line-segment-detection
  - wireframe-parsing
pipeline_tag: image-feature-extraction
library_name: mlx

ScaleLSD — MLX weights

MLX-format weights for ScaleLSD, converted for use with mlx-swift-ScaleLSD on Apple Silicon.

These are converted redistributions of the original checkpoints published at cherubicxn/scalelsd. No retraining or fine-tuning was performed — the numerical content of the network is unchanged.

Contents

folder	LayerScale	parameters	source checkpoint
scalelsd-vitbase-v1/	no	122,525,833	scalelsd-vitbase-v1-train-sa1b.pt
scalelsd-vitbase-v2/	yes	122,544,265	scalelsd-vitbase-v2-train-sa1b.pt

Each folder holds config.json + model.safetensors. Upstream recommends v2 by default.

Usage

import MLXScaleLSD

// Downloads from this repo on first use, then caches locally.
let directory = try await ModelStore.download(.v2)
let session = try ScaleLSDSession.load(directory: directory)

let image = try ScaleLSDSession.loadImage(at: imageURL)
let result = try session.detect(image)
for segment in result.segments(minimumScore: 10) {
    print(segment.x1, segment.y1, segment.x2, segment.y2, segment.score)
}

Or from the command line:

scalelsd detect -m <directory> -i image.jpg -e png --save-to out/

What was changed in conversion

The original checkpoints are PyTorch pickles, which MLX cannot read. Scripts/convert.py in the Swift repo performs a format conversion plus several inference-only graph simplifications, each of which is numerically equivalent (verified to ~2e-6 relative against the PyTorch reference):

• Weight standardisation baked in. timm's StdConv2dSame re-standardises its weight on every forward pass; inference weights are frozen, so the standardised tensor is stored directly. (Note: the hybrid ViT uses eps=1e-8, not the class default 1e-6.)
• conv + BatchNorm folded. The 16 conv(bias=False) -> BatchNorm2d pairs in DPT's ResidualConvUnit_custom collapse into single biased convolutions.
• nn.Sequential indices renamed to named submodules, so keys read structurally.
• Conv weights transposed from PyTorch (O, I, kH, kW) to MLX (O, kH, kW, I).
• The 1000-class ImageNet classifier head dropped — ScaleLSD never calls it.

Accuracy

Verified stage by stage against the PyTorch reference. The final 9-channel HAT field matches to 1.1e-05 (v1) / 1.4e-05 (v2) maximum relative error. End-to-end on assets/indoor.jpg:

	v1	v2
junctions matched within 0.01 px	512/512	511/512
segments matched within 0.01 px	1879/1880	1581/1590

Detections are not bit-exact by construction: the 512-junction cap and the nearest-junction assignment are discrete choices that a sub-noise perturbation can flip. See docs/PARITY.md.

Performance

Apple M5 Max, 512×512 input, Release build, median of 20 runs:

runtime	per image
mlx-swift (this port)	57 ms
PyTorch 2.13, MPS	88 ms
PyTorch 2.13, CPU	652 ms

License and attribution

Apache-2.0, inherited from the original checkpoints at cherubicxn/scalelsd. The upstream ScaleLSD source is MIT (Copyright © 2023 Nan Xue). Original work and all model credit belong to the ScaleLSD authors; this repository contributes only a format conversion.

@inproceedings{ScaleLSD,
    title = {ScaleLSD: Scalable Deep Line Segment Detection Streamlined},
    author = {Zeran Ke and Bin Tan and Xianwei Zheng and Yujun Shen and Tianfu Wu and Nan Xue},
    booktitle = {IEEE Conference on Computer Vision and Pattern Recognition (CVPR)},
    year = {2025},
}