| """ORMBG (Open Remove Background Model, an IS-Net from xuebinqin/DIS) in MLX: a port of schirrmacher/ormbg's |
| `ormbg/models/ormbg.py` (Apache-2.0). |
| |
| import mlx_model |
| model = mlx_model.load("fp16") # a folder with config.json + model.safetensors |
| alpha = mlx_model.matte(model, rgb) # rgb: HxWx3 uint8 -> HxW uint8 alpha (255 = foreground) |
| |
| Runs on its own next to its weights (mlx and numpy). NHWC; each BatchNorm is folded into the conv before it (exact |
| in fp32, done at conversion), so the weights are conv kernels and biases only. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import json |
| from pathlib import Path |
|
|
| import mlx.core as mx |
| import mlx.nn as nn |
| import numpy as np |
|
|
| CONFIG = {"model_type": "ormbg", "input_size": 1024} |
| PRECISIONS = ["bf16", "fp16"] |
|
|
|
|
| def resize_bilinear(x, size): |
| """torch F.interpolate(mode="bilinear", align_corners=False) for NHWC.""" |
| b, h, w, c = x.shape |
| oh, ow = size |
| if (oh, ow) == (h, w): |
| return x |
|
|
| def axis(n_in, n_out): |
| s = (np.arange(n_out, dtype=np.float64) + 0.5) * (n_in / n_out) - 0.5 |
| s = np.clip(s, 0, None) |
| i0 = np.minimum(np.floor(s).astype(np.int32), n_in - 1) |
| i1 = np.minimum(i0 + 1, n_in - 1) |
| f = (s - i0).astype(np.float32) |
| return mx.array(i0), mx.array(i1), mx.array(f) |
|
|
| y0, y1, fy = axis(h, oh) |
| x0, x1, fx = axis(w, ow) |
| fy = fy.astype(x.dtype)[None, :, None, None] |
| fx = fx.astype(x.dtype)[None, None, :, None] |
| top, bot = x[:, y0], x[:, y1] |
| top = top[:, :, x0] * (1 - fx) + top[:, :, x1] * fx |
| bot = bot[:, :, x0] * (1 - fx) + bot[:, :, x1] * fx |
| return top * (1 - fy) + bot * fy |
|
|
|
|
| def maxpool2(x): |
| """MaxPool2d(2, 2, ceil_mode=True), NHWC.""" |
| b, h, w, c = x.shape |
| ph, pw = h % 2, w % 2 |
| if ph or pw: |
| x = mx.pad(x, [(0, 0), (0, ph), (0, pw), (0, 0)], constant_values=-float("inf")) |
| h, w = h + ph, w + pw |
| return x.reshape(b, h // 2, 2, w // 2, 2, c).max(axis=(2, 4)) |
|
|
|
|
| class REBNCONV(nn.Module): |
| def __init__(self, cin, cout, d=1): |
| super().__init__() |
| self.conv_s1 = nn.Conv2d(cin, cout, 3, padding=d, dilation=d) |
|
|
| def __call__(self, x): |
| return nn.relu(self.conv_s1(x)) |
|
|
|
|
| class RSU(nn.Module): |
| """RSU-L (L = depth: 7, 6, 5, 4): an encoder-decoder with pooling; dilation 2 at the bottom.""" |
|
|
| def __init__(self, depth, cin, mid, cout): |
| super().__init__() |
| self.depth = depth |
| self.rebnconvin = REBNCONV(cin, cout) |
| setattr(self, "rebnconv1", REBNCONV(cout, mid)) |
| for i in range(2, depth): |
| setattr(self, f"rebnconv{i}", REBNCONV(mid, mid)) |
| setattr(self, f"rebnconv{depth}", REBNCONV(mid, mid, d=2)) |
| for i in range(depth - 1, 1, -1): |
| setattr(self, f"rebnconv{i}d", REBNCONV(mid * 2, mid)) |
| self.rebnconv1d = REBNCONV(mid * 2, cout) |
|
|
| def __call__(self, x): |
| hxin = self.rebnconvin(x) |
| enc, hx = [], hxin |
| for i in range(1, self.depth): |
| hx = getattr(self, f"rebnconv{i}")(hx) |
| enc.append(hx) |
| if i < self.depth - 1: |
| hx = maxpool2(hx) |
| hx = getattr(self, f"rebnconv{self.depth}")(hx) |
| for i in range(self.depth - 1, 0, -1): |
| hx = getattr(self, f"rebnconv{i}d")(mx.concatenate([hx, enc[i - 1]], -1)) |
| if i > 1: |
| hx = resize_bilinear(hx, enc[i - 2].shape[1:3]) |
| return hx + hxin |
|
|
|
|
| class RSU4F(nn.Module): |
| """RSU-4F: no pooling, dilations 1, 2, 4, 8.""" |
|
|
| def __init__(self, cin, mid, cout): |
| super().__init__() |
| self.rebnconvin = REBNCONV(cin, cout) |
| self.rebnconv1 = REBNCONV(cout, mid) |
| self.rebnconv2 = REBNCONV(mid, mid, 2) |
| self.rebnconv3 = REBNCONV(mid, mid, 4) |
| self.rebnconv4 = REBNCONV(mid, mid, 8) |
| self.rebnconv3d = REBNCONV(mid * 2, mid, 4) |
| self.rebnconv2d = REBNCONV(mid * 2, mid, 2) |
| self.rebnconv1d = REBNCONV(mid * 2, cout) |
|
|
| def __call__(self, x): |
| hxin = self.rebnconvin(x) |
| hx1 = self.rebnconv1(hxin) |
| hx2 = self.rebnconv2(hx1) |
| hx3 = self.rebnconv3(hx2) |
| hx4 = self.rebnconv4(hx3) |
| hx3d = self.rebnconv3d(mx.concatenate([hx4, hx3], -1)) |
| hx2d = self.rebnconv2d(mx.concatenate([hx3d, hx2], -1)) |
| return self.rebnconv1d(mx.concatenate([hx2d, hx1], -1)) + hxin |
|
|
|
|
| class Model(nn.Module): |
| def __init__(self, config: dict = CONFIG): |
| super().__init__() |
| self.size = config["input_size"] |
| self.conv_in = nn.Conv2d(3, 64, 3, stride=2, padding=1) |
| self.stage1, self.stage2 = RSU(7, 64, 32, 64), RSU(6, 64, 32, 128) |
| self.stage3, self.stage4 = RSU(5, 128, 64, 256), RSU(4, 256, 128, 512) |
| self.stage5, self.stage6 = RSU4F(512, 256, 512), RSU4F(512, 256, 512) |
| self.stage5d, self.stage4d = RSU4F(1024, 256, 512), RSU(4, 1024, 128, 256) |
| self.stage3d, self.stage2d, self.stage1d = RSU(5, 512, 64, 128), RSU(6, 256, 32, 64), RSU(7, 128, 16, 64) |
| self.side1 = nn.Conv2d(64, 1, 3, padding=1) |
|
|
| def __call__(self, x): |
| """x: (B, S, S, 3) in [0, 1] -> (B, S, S, 1) foreground probability (side output 1, the one used).""" |
| hxin = self.conv_in(x) |
| hx1 = self.stage1(hxin) |
| hx2 = self.stage2(maxpool2(hx1)) |
| hx3 = self.stage3(maxpool2(hx2)) |
| hx4 = self.stage4(maxpool2(hx3)) |
| hx5 = self.stage5(maxpool2(hx4)) |
| hx6 = self.stage6(maxpool2(hx5)) |
| up = lambda a, b: resize_bilinear(a, b.shape[1:3]) |
| hx5d = self.stage5d(mx.concatenate([up(hx6, hx5), hx5], -1)) |
| hx4d = self.stage4d(mx.concatenate([up(hx5d, hx4), hx4], -1)) |
| hx3d = self.stage3d(mx.concatenate([up(hx4d, hx3), hx3], -1)) |
| hx2d = self.stage2d(mx.concatenate([up(hx3d, hx2), hx2], -1)) |
| hx1d = self.stage1d(mx.concatenate([up(hx2d, hx1), hx1], -1)) |
| return mx.sigmoid(resize_bilinear(self.side1(hx1d), x.shape[1:3])) |
|
|
|
|
| def load(folder) -> Model: |
| folder = Path(folder) |
| model = Model(json.loads((folder / "config.json").read_text())) |
| model.load_weights(str(folder / "model.safetensors"), strict=True) |
| mx.eval(model.parameters()) |
| return model |
|
|
|
|
| def predict(model: Model, rgb: np.ndarray) -> np.ndarray: |
| """The raw model output at the original size, float32 HxW in [0, 1] (as ormbg/inference.py, before min-max).""" |
| dt = model.conv_in.weight.dtype |
| h, w, _ = rgb.shape |
| x = mx.array(rgb.astype(np.float32))[None] |
| x = mx.floor(resize_bilinear(x, (model.size, model.size))) / 255.0 |
| y = model(x.astype(dt)).astype(mx.float32) |
| return np.array(resize_bilinear(y, (h, w)))[0, :, :, 0] |
|
|
|
|
| def matte(model: Model, rgb: np.ndarray) -> np.ndarray: |
| """HxWx3 uint8 -> HxW uint8 alpha, min-max normalized like ormbg/inference.py.""" |
| p = predict(model, rgb) |
| p = (p - p.min()) / max(p.max() - p.min(), 1e-8) |
| return (p * 255).astype(np.uint8) |
|
|
|
|
| def convert(source: Path) -> dict: |
| """runs/source/models/ormbg.pth -> conv weights with each BatchNorm folded in (fp32), MLX layout.""" |
| import torch |
| sd = {k: v.double() for k, v in torch.load(Path(source) / "models" / "ormbg.pth", map_location="cpu", |
| weights_only=True).items()} |
| out = {} |
| for k, w in sd.items(): |
| if not k.endswith(".weight") or w.ndim != 4: |
| continue |
| base = k[:-len(".weight")] |
| b = sd.get(base + ".bias", torch.zeros(w.shape[0], dtype=torch.float64)) |
| bn = base.replace("conv_s1", "bn_s1") |
| if bn != base and bn + ".running_mean" in sd: |
| g = sd[bn + ".weight"] / torch.sqrt(sd[bn + ".running_var"] + 1e-5) |
| w = w * g[:, None, None, None] |
| b = (b - sd[bn + ".running_mean"]) * g + sd[bn + ".bias"] |
| if base.startswith("side") and base != "side1": |
| continue |
| out[base + ".weight"] = w.float().numpy().transpose(0, 2, 3, 1) |
| out[base + ".bias"] = b.float().numpy() |
| return out |
|
|
