"""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])