LlamaCasty/tools/selfie-seg/convert_tflite_to_onnx.py

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