README.md

---
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](https://github.com/ant-research/scalelsd), converted for use
with [**mlx-swift-ScaleLSD**](https://github.com/mnmly/mlx-swift-ScaleLSD) on Apple Silicon.

These are **converted redistributions** of the original checkpoints published at
[cherubicxn/scalelsd](https://huggingface.co/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

```swift
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:

```bash
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](https://github.com/mnmly/mlx-swift-ScaleLSD/blob/main/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](https://huggingface.co/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.

```bibtex
@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},
}
```