TASK 51 slice 1: background-effects engine — bundled MediaPipe Selfie Segmenter (ONNX) + Remove (Unit: bg-effects engine)
- Model: Google's official MediaPipe Selfie Segmenter (landscape, Apache-2.0 per model card), converted to ONNX by tools/selfie-seg/convert_tflite_to_onnx.py (custom op Convolution2DTransposeBias -> ConvTranspose) and verified against MediaPipe's own runtime (verify_against_mediapipe.py: max |diff| 5e-5, 100% binary agreement). Embedded resource, 430 KB. - Services/BackgroundEffects: IPersonSegmenter seam, OnnxPersonSegmenter (onnxruntime 1.31.0 CPU, never throws), BackgroundEffectProcessor (None passthrough, Remove = alpha mask). Not wired yet. - Tests: Remove integration test (fake segmenter) + bundled-model contract test vs MediaPipe reference mask. Full suite 407/408 — the one failure is the documented interactive-desktop flake RealMouseDrag_OnTheLayerList_PersistsTheReorder (passed in the run before; code path untouched). - Docs: ai.md reversal + section + licensing guard, THIRD-PARTY-NOTICES 11-12, TASK 51, indexes, MyMistakes recipe, HANDOFF. References: obs-backgroundremoval (approach only — ONNX Runtime on CPU, 2 threads; GPL-3.0, no code used) https://github.com/royshil/obs-backgroundremoval ; MediaPipe image segmenter https://ai.google.dev/edge/mediapipe/solutions/vision/image_segmenter
This commit is contained in:
@@ -0,0 +1,218 @@
|
||||
"""Convert Google's MediaPipe Selfie Segmenter (.tflite) to ONNX for LlamaCasty (TASK 51).
|
||||
|
||||
Why a hand-rolled converter: the model ends in MediaPipe's custom op
|
||||
`Convolution2DTransposeBias`, which tf2onnx / tflite2onnx don't support. The graph only
|
||||
uses 11 op types, so a direct, auditable translation is smaller than any workaround.
|
||||
|
||||
Source model (Apache-2.0, model card "MediaPipe Selfie Segmentation", Google, 2021-05-06):
|
||||
https://storage.googleapis.com/mediapipe-models/image_segmenter/selfie_segmenter_landscape/float16/latest/selfie_segmenter_landscape.tflite
|
||||
|
||||
Layout: TFLite is NHWC. The ONNX graph keeps the SAME external contract
|
||||
(input [1,144,256,3] float RGB in [0,1], output [1,144,256,1] person probability in [0,1])
|
||||
and transposes to NCHW internally.
|
||||
|
||||
Usage: python convert_tflite_to_onnx.py in.tflite out.onnx
|
||||
Requires: tflite, onnx, numpy. Verify with verify_against_mediapipe.py.
|
||||
"""
|
||||
import sys
|
||||
|
||||
import numpy as np
|
||||
import onnx
|
||||
import tflite
|
||||
from onnx import TensorProto, helper, numpy_helper
|
||||
from tflite.ActivationFunctionType import ActivationFunctionType as Act
|
||||
from tflite.BuiltinOperator import BuiltinOperator
|
||||
from tflite.Padding import Padding
|
||||
from tflite.TensorType import TensorType
|
||||
|
||||
OPSET = 17
|
||||
OP_NAMES = {v: k for k, v in BuiltinOperator.__dict__.items() if not k.startswith("_")}
|
||||
NP_TYPES = {TensorType.FLOAT32: np.float32, TensorType.FLOAT16: np.float16, TensorType.INT32: np.int32}
|
||||
|
||||
|
||||
def main(src, dst):
|
||||
model = tflite.Model.GetRootAsModel(open(src, "rb").read(), 0)
|
||||
g = model.Subgraphs(0)
|
||||
|
||||
def shape(ix):
|
||||
return g.Tensors(ix).ShapeAsNumpy().tolist()
|
||||
|
||||
def const(ix):
|
||||
"""Constant tensor value, or None for runtime tensors."""
|
||||
t = g.Tensors(ix)
|
||||
buf = model.Buffers(t.Buffer())
|
||||
if buf is None or buf.DataLength() == 0:
|
||||
return None
|
||||
raw = buf.DataAsNumpy().tobytes()
|
||||
return np.frombuffer(raw, dtype=NP_TYPES[t.Type()]).reshape(shape(ix))
|
||||
|
||||
nodes, inits = [], []
|
||||
consts = {} # tensor index -> float32 numpy (after DEQUANTIZE folding)
|
||||
uid = [0]
|
||||
|
||||
def name(ix):
|
||||
return f"t{ix}"
|
||||
|
||||
def fresh(prefix):
|
||||
uid[0] += 1
|
||||
return f"{prefix}_{uid[0]}"
|
||||
|
||||
def add_init(arr, nm=None):
|
||||
nm = nm or fresh("c")
|
||||
inits.append(numpy_helper.from_array(np.ascontiguousarray(arr), nm))
|
||||
return nm
|
||||
|
||||
def value(ix):
|
||||
if ix in consts:
|
||||
return consts[ix]
|
||||
c = const(ix)
|
||||
if c is not None:
|
||||
consts[ix] = c.astype(np.float32) if c.dtype == np.float16 else c
|
||||
return consts[ix]
|
||||
return None
|
||||
|
||||
def activation(out_name, act, final_name):
|
||||
if act == Act.NONE:
|
||||
return out_name
|
||||
if act == Act.RELU:
|
||||
nodes.append(helper.make_node("Relu", [out_name], [final_name]))
|
||||
elif act == Act.RELU6:
|
||||
nodes.append(helper.make_node("Clip", [out_name, add_init(np.float32(0)), add_init(np.float32(6))], [final_name]))
|
||||
else:
|
||||
raise NotImplementedError(f"fused activation {act}")
|
||||
return final_name
|
||||
|
||||
def operand(ix):
|
||||
"""Runtime tensor name, or a constant reshaped from NHWC broadcast form to NCHW."""
|
||||
v = value(ix)
|
||||
if v is None:
|
||||
return name(ix)
|
||||
if v.ndim == 1:
|
||||
v = v.reshape(1, -1, 1, 1)
|
||||
elif v.ndim == 4:
|
||||
v = v.transpose(0, 3, 1, 2)
|
||||
return add_init(v.astype(np.float32))
|
||||
|
||||
# NHWC -> NCHW at the input.
|
||||
inp = g.Inputs(0)
|
||||
nodes.append(helper.make_node("Transpose", ["input"], [name(inp)], perm=[0, 3, 1, 2]))
|
||||
|
||||
for i in range(g.OperatorsLength()):
|
||||
op = g.Operators(i)
|
||||
oc = model.OperatorCodes(op.OpcodeIndex())
|
||||
code = max(oc.BuiltinCode(), oc.DeprecatedBuiltinCode())
|
||||
kind = "CUSTOM:" + oc.CustomCode().decode() if code == BuiltinOperator.CUSTOM else OP_NAMES[code]
|
||||
ins = [op.Inputs(j) for j in range(op.InputsLength())]
|
||||
out = op.Outputs(0)
|
||||
o = name(out)
|
||||
|
||||
if kind == "DEQUANTIZE":
|
||||
consts[out] = const(ins[0]).astype(np.float32)
|
||||
continue
|
||||
|
||||
if kind in ("CONV_2D", "DEPTHWISE_CONV_2D"):
|
||||
opt = (tflite.Conv2DOptions() if kind == "CONV_2D" else tflite.DepthwiseConv2DOptions())
|
||||
bo = op.BuiltinOptions()
|
||||
opt.Init(bo.Bytes, bo.Pos)
|
||||
w = value(ins[1])
|
||||
b = value(ins[2]) if len(ins) > 2 and ins[2] >= 0 else None
|
||||
if kind == "CONV_2D":
|
||||
W = w.transpose(0, 3, 1, 2) # OHWI -> OIHW
|
||||
group = 1
|
||||
else:
|
||||
# [1,H,W,C*mult] -> [C*mult,1,H,W], group = input channels
|
||||
W = w.transpose(3, 0, 1, 2)
|
||||
group = shape(ins[0])[3]
|
||||
attrs = dict(
|
||||
strides=[opt.StrideH(), opt.StrideW()],
|
||||
dilations=[opt.DilationHFactor(), opt.DilationWFactor()],
|
||||
kernel_shape=list(W.shape[2:]),
|
||||
group=group,
|
||||
auto_pad="SAME_UPPER" if opt.Padding() == Padding.SAME else "VALID",
|
||||
)
|
||||
conv_in = [name(ins[0]), add_init(W)] + ([add_init(b)] if b is not None else [])
|
||||
tmp = fresh("conv") if opt.FusedActivationFunction() != Act.NONE else o
|
||||
nodes.append(helper.make_node("Conv", conv_in, [tmp], **attrs))
|
||||
activation(tmp, opt.FusedActivationFunction(), o)
|
||||
continue
|
||||
|
||||
if kind in ("ADD", "MUL"):
|
||||
opt = tflite.AddOptions() if kind == "ADD" else tflite.MulOptions()
|
||||
bo = op.BuiltinOptions()
|
||||
act = Act.NONE
|
||||
if bo is not None:
|
||||
opt.Init(bo.Bytes, bo.Pos)
|
||||
act = opt.FusedActivationFunction()
|
||||
tmp = fresh("ew") if act != Act.NONE else o
|
||||
nodes.append(helper.make_node("Add" if kind == "ADD" else "Mul", [operand(ins[0]), operand(ins[1])], [tmp]))
|
||||
activation(tmp, act, o)
|
||||
continue
|
||||
|
||||
if kind == "RELU":
|
||||
nodes.append(helper.make_node("Relu", [name(ins[0])], [o]))
|
||||
continue
|
||||
if kind == "HARD_SWISH":
|
||||
nodes.append(helper.make_node("HardSwish", [name(ins[0])], [o]))
|
||||
continue
|
||||
if kind == "LOGISTIC":
|
||||
nodes.append(helper.make_node("Sigmoid", [name(ins[0])], [o]))
|
||||
continue
|
||||
|
||||
if kind == "MEAN":
|
||||
opt = tflite.ReducerOptions()
|
||||
bo = op.BuiltinOptions()
|
||||
opt.Init(bo.Bytes, bo.Pos)
|
||||
nhwc_to_nchw = {0: 0, 1: 2, 2: 3, 3: 1}
|
||||
axes = [nhwc_to_nchw[int(a)] for a in value(ins[1]).reshape(-1)]
|
||||
nodes.append(helper.make_node("ReduceMean", [name(ins[0])], [o], axes=axes, keepdims=1 if opt.KeepDims() else 0))
|
||||
continue
|
||||
|
||||
if kind == "RESIZE_BILINEAR":
|
||||
opt = tflite.ResizeBilinearOptions()
|
||||
bo = op.BuiltinOptions()
|
||||
opt.Init(bo.Bytes, bo.Pos)
|
||||
h, w = (int(x) for x in value(ins[1]).reshape(-1))
|
||||
c = shape(ins[0])[3]
|
||||
mode = "align_corners" if opt.AlignCorners() else ("half_pixel" if opt.HalfPixelCenters() else "asymmetric")
|
||||
sizes = add_init(np.array([1, c, h, w], dtype=np.int64))
|
||||
nodes.append(helper.make_node("Resize", [name(ins[0]), "", "", sizes], [o], mode="linear", coordinate_transformation_mode=mode))
|
||||
continue
|
||||
|
||||
if kind == "CUSTOM:Convolution2DTransposeBias":
|
||||
# Custom options = TfLiteTransposeConvParams {padding, stride_w, stride_h} (int32 each).
|
||||
# NOTE: `padding` is the C-API TfLitePadding enum (Unknown=0, SAME=1, VALID=2),
|
||||
# NOT the flatbuffer Padding enum (SAME=0, VALID=1).
|
||||
padding, stride_w, stride_h = np.frombuffer(bytes(op.CustomOptionsAsNumpy().tolist()), dtype=np.int32)[:3]
|
||||
w = value(ins[1]) # [O, kH, kW, I]
|
||||
b = value(ins[2])
|
||||
W = w.transpose(3, 0, 1, 2) # -> [I, O, kH, kW] (ONNX ConvTranspose)
|
||||
kh, kw = w.shape[1], w.shape[2]
|
||||
if padding != 1 or kh != stride_h or kw != stride_w:
|
||||
raise NotImplementedError("only SAME, kernel == stride transpose conv is supported")
|
||||
nodes.append(helper.make_node("ConvTranspose", [name(ins[0]), add_init(W), add_init(b)], [o],
|
||||
strides=[int(stride_h), int(stride_w)], kernel_shape=[kh, kw]))
|
||||
continue
|
||||
|
||||
raise NotImplementedError(kind)
|
||||
|
||||
# NCHW -> NHWC at the output.
|
||||
outp = g.Outputs(0)
|
||||
nodes.append(helper.make_node("Transpose", [name(outp)], ["mask"], perm=[0, 2, 3, 1]))
|
||||
|
||||
graph = helper.make_graph(
|
||||
nodes, "mediapipe_selfie_segmenter_landscape",
|
||||
[helper.make_tensor_value_info("input", TensorProto.FLOAT, shape(inp))],
|
||||
[helper.make_tensor_value_info("mask", TensorProto.FLOAT, shape(outp))],
|
||||
initializer=inits)
|
||||
m = helper.make_model(graph, opset_imports=[helper.make_opsetid("", OPSET)],
|
||||
producer_name="llamacasty-tflite2onnx",
|
||||
doc_string="MediaPipe Selfie Segmenter (landscape), Apache-2.0, Google. "
|
||||
"Converted for LlamaCasty by tools/selfie-seg/convert_tflite_to_onnx.py")
|
||||
m.ir_version = 8
|
||||
onnx.checker.check_model(m)
|
||||
onnx.save(m, dst)
|
||||
print(f"wrote {dst}: {len(nodes)} nodes, {len(inits)} initializers")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main(sys.argv[1], sys.argv[2])
|
||||
@@ -0,0 +1,59 @@
|
||||
"""Verify the converted ONNX model against Google's own MediaPipe runtime (TASK 51).
|
||||
|
||||
Feeds the SAME 256x144 RGB frame to (a) MediaPipe's ImageSegmenter running the original
|
||||
.tflite (with its custom op) and (b) onnxruntime running our converted .onnx, then compares
|
||||
the person-probability masks. Also writes side-by-side PNGs for an eyeball check.
|
||||
|
||||
Usage: python verify_against_mediapipe.py model.tflite model.onnx photo.jpg [out_prefix]
|
||||
Requires: mediapipe, onnxruntime, numpy, pillow.
|
||||
"""
|
||||
import sys
|
||||
|
||||
import mediapipe as mp
|
||||
import numpy as np
|
||||
import onnxruntime as ort
|
||||
from mediapipe.tasks import python as mp_python
|
||||
from mediapipe.tasks.python import vision
|
||||
from PIL import Image
|
||||
|
||||
W, H = 256, 144
|
||||
|
||||
|
||||
def main(tfl, onx, photo, prefix="verify"):
|
||||
img = Image.open(photo).convert("RGB")
|
||||
# Cover-crop to 16:9, then resize to the model's exact input so neither runtime resamples.
|
||||
sw, sh = img.size
|
||||
target = W / H
|
||||
if sw / sh > target:
|
||||
nw = int(sh * target); x = (sw - nw) // 2; img = img.crop((x, 0, x + nw, sh))
|
||||
else:
|
||||
nh = int(sw / target); y = (sh - nh) // 2; img = img.crop((0, y, sw, y + nh))
|
||||
img = img.resize((W, H), Image.BILINEAR)
|
||||
rgb = np.asarray(img, dtype=np.uint8)
|
||||
|
||||
seg = vision.ImageSegmenter.create_from_options(vision.ImageSegmenterOptions(
|
||||
base_options=mp_python.BaseOptions(model_asset_path=tfl), output_confidence_masks=True))
|
||||
res = seg.segment(mp.Image(image_format=mp.ImageFormat.SRGB, data=rgb))
|
||||
ref = res.confidence_masks[0].numpy_view().astype(np.float32).reshape(H, W)
|
||||
|
||||
sess = ort.InferenceSession(onx, providers=["CPUExecutionProvider"])
|
||||
x = (rgb.astype(np.float32) / 255.0)[None]
|
||||
got = sess.run(None, {"input": x})[0].reshape(H, W)
|
||||
|
||||
diff = np.abs(ref - got)
|
||||
agree = np.mean((ref > 0.5) == (got > 0.5))
|
||||
print(f"ref range {ref.min():.3f}..{ref.max():.3f} person%={np.mean(ref > 0.5) * 100:.1f}")
|
||||
print(f"onnx range {got.min():.3f}..{got.max():.3f} person%={np.mean(got > 0.5) * 100:.1f}")
|
||||
print(f"max|diff|={diff.max():.5f} mean|diff|={diff.mean():.6f} binary-agreement={agree * 100:.2f}%")
|
||||
|
||||
def to_img(m):
|
||||
return Image.fromarray((np.clip(m, 0, 1) * 255).astype(np.uint8)).convert("RGB")
|
||||
sheet = Image.new("RGB", (W * 3, H))
|
||||
sheet.paste(img, (0, 0)); sheet.paste(to_img(ref), (W, 0)); sheet.paste(to_img(got), (W * 2, 0))
|
||||
sheet.resize((W * 6, H * 2), Image.NEAREST).save(f"{prefix}_sheet.png")
|
||||
return diff.max()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
md = main(*sys.argv[1:])
|
||||
sys.exit(0 if md < 1e-3 else 1)
|
||||
Reference in New Issue
Block a user