File size: 3,950 Bytes
e2f3b24 | 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 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 | """
Visual end-to-end check of the SHIPPING path: run the .mlpackage exactly as the
iOS app does (center-square crop -> 512 input -> fp16 logits -> bilinear upsample
-> 19-way argmax -> palette LUT -> alpha mix) and write a contact sheet.
This mirrors SegmentationShaders.metal's math on the CPU, so a good-looking
sheet means the model AND the app's decode contract AND the palette are right.
Usage: verify_coreml_visual.py --model <.mlpackage> --out <dir> [--source <dir>]
"""
import argparse, glob, os
import numpy as np
from PIL import Image, ImageDraw
import coremltools as ct
from palette import CLASS_NAMES, PALETTE
PAL = np.array(PALETTE, dtype=np.float32)
def center_square(im):
w, h = im.size
s = min(w, h)
return im.crop(((w - s) // 2, (h - s) // 2, (w + s) // 2, (h + s) // 2))
def bilinear_upsample(logits, out):
"""(C,g,g) -> (C,out,out), matching the shader's bilinear sampler."""
C, g, _ = logits.shape
# sample centers, same convention as a normalized-UV bilinear sampler
coord = (np.arange(out) + 0.5) / out * g - 0.5
coord = np.clip(coord, 0, g - 1)
i0 = np.floor(coord).astype(int)
i1 = np.minimum(i0 + 1, g - 1)
w1 = (coord - i0).astype(np.float32)
w0 = 1.0 - w1
tmp = logits[:, :, i0] * w0 + logits[:, :, i1] * w1 # x interp
return tmp[:, i0, :] * w0[None, :, None] + tmp[:, i1, :] * w1[None, :, None]
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--model", required=True)
ap.add_argument("--out", default="/private/tmp/claude-501/-Users-ari-Documents-XcodeProjects-facesegmentation/149f7917-f682-4847-ba5b-1e726cecc20c/scratchpad/visual")
ap.add_argument("--source", default="/Users/ari/FaceSegmentation/dataset_celebamaskhq_semantic/images/val")
ap.add_argument("--n", type=int, default=8)
ap.add_argument("--alpha", type=float, default=0.55)
ap.add_argument("--tile", type=int, default=320)
args = ap.parse_args()
os.makedirs(args.out, exist_ok=True)
m = ct.models.MLModel(args.model, compute_units=ct.ComputeUnit.CPU_AND_NE)
spec_in = list(m.input_description)[0]
R = 512
files = sorted(glob.glob(os.path.join(args.source, "**", "*.jpg"), recursive=True))[: args.n]
if not files:
files = sorted(glob.glob(os.path.join(args.source, "**", "*.png"), recursive=True))[: args.n]
tiles, class_hits = [], np.zeros(len(CLASS_NAMES))
for p in files:
im = center_square(Image.open(p).convert("RGB")).resize((R, R), Image.BILINEAR)
out = m.predict({spec_in: im})
logits = np.asarray(out["logits"], dtype=np.float32)[0] # (19,g,g)
up = bilinear_upsample(logits, R) # (19,R,R)
cls = up.argmax(0).astype(np.int32)
class_hits += np.bincount(cls.ravel(), minlength=len(CLASS_NAMES)) / cls.size
rgb = np.asarray(im, dtype=np.float32)
col = PAL[cls]
a = (cls != 0).astype(np.float32)[..., None] * args.alpha
blend = (rgb * (1 - a) + col * a).astype(np.uint8)
pair = Image.new("RGB", (args.tile * 2, args.tile))
pair.paste(im.resize((args.tile, args.tile)), (0, 0))
pair.paste(Image.fromarray(blend).resize((args.tile, args.tile)), (args.tile, 0))
tiles.append(pair)
cols = 2
rows = (len(tiles) + cols - 1) // cols
W, H = tiles[0].size
sheet = Image.new("RGB", (cols * W, rows * H), (16, 16, 16))
for i, t in enumerate(tiles):
sheet.paste(t, ((i % cols) * W, (i // cols) * H))
sheet_path = os.path.join(args.out, "verify_sheet.jpg")
sheet.save(sheet_path, quality=92)
# legend of what the model actually predicted
order = np.argsort(-class_hits)
print("mean class coverage (top 10):")
for i in order[:10]:
if class_hits[i] > 0:
print(f" {CLASS_NAMES[i]:14s} {class_hits[i]/len(tiles):6.2%}")
print("wrote", sheet_path)
if __name__ == "__main__":
main()
|