Upload convert.py with huggingface_hub

convert.py

@@ -0,0 +1,42 @@
+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)