File size: 1,644 Bytes
e0bfda2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 | import sys, torch, torch.nn as nn
import coremltools as ct
sys.path.insert(0, "/tmp/clothseg/repo")
from network import U2NET
CKPT = "/tmp/clothseg/cloth_segm_u2net_latest.pth"
OUT = "/tmp/clothseg/ClothSegmentation.mlpackage"
# Build + load (the checkpoint was saved as a plain state_dict, possibly with a module. prefix).
net = U2NET(in_ch=3, out_ch=4)
sd = torch.load(CKPT, map_location="cpu")
if isinstance(sd, dict) and "model_state_dict" in sd:
sd = sd["model_state_dict"]
sd = { (k[7:] if k.startswith("module.") else k): v for k, v in sd.items() }
net.load_state_dict(sd)
net.eval()
# Wrapper: return only the fused output d0, as per-class softmax probabilities (1,4,768,768).
class Wrap(nn.Module):
def __init__(self, m): super().__init__(); self.m = m
def forward(self, x):
d0 = self.m(x)[0] # (1,4,768,768) logits
return torch.softmax(d0, dim=1) # per-class probabilities
wrap = Wrap(net).eval()
dummy = torch.randn(1, 3, 768, 768)
with torch.no_grad():
traced = torch.jit.trace(wrap, dummy)
# CoreML image input: pixels [0,255] -> model wants (p/255-0.5)/0.5 = p/127.5 - 1.
mlmodel = ct.convert(
traced,
inputs=[ct.ImageType(name="image", shape=(1, 3, 768, 768),
scale=1.0/127.5, bias=[-1.0, -1.0, -1.0],
color_layout=ct.colorlayout.RGB)],
outputs=[ct.TensorType(name="probs")],
minimum_deployment_target=ct.target.iOS16,
compute_precision=ct.precision.FLOAT16,
)
mlmodel.short_description = "U2NET cloth segmentation (bg/upper/lower/full) — 768x768, softmax probs"
mlmodel.save(OUT)
print("SAVED", OUT)
|