wehub-resource-sync
162 lines
6.6 KiB
Python
162 lines
6.6 KiB
Python
# Licensed to the Apache Software Foundation (ASF) under one
|
|
# or more contributor license agreements. See the NOTICE file
|
|
# distributed with this work for additional information
|
|
# regarding copyright ownership. The ASF licenses this file
|
|
# to you under the Apache License, Version 2.0 (the
|
|
# "License"); you may not use this file except in compliance
|
|
# with the License. You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing,
|
|
# software distributed under the License is distributed on an
|
|
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
# KIND, either express or implied. See the License for the
|
|
# specific language governing permissions and limitations
|
|
# under the License.
|
|
# pylint: disable=invalid-name, unused-argument, too-many-lines, import-outside-toplevel
|
|
# pylint: disable=broad-exception-raised, use-list-literal
|
|
"""Tensorflow lite frontend helper to parse custom options in Flexbuffer format."""
|
|
|
|
import struct
|
|
from enum import IntEnum
|
|
|
|
|
|
class BitWidth(IntEnum):
|
|
"""Flexbuffer bit width schema from flexbuffers.h"""
|
|
|
|
BIT_WIDTH_8 = 0
|
|
BIT_WIDTH_16 = 1
|
|
BIT_WIDTH_32 = 2
|
|
BIT_WIDTH_64 = 3
|
|
|
|
|
|
class FlexBufferType(IntEnum):
|
|
"""Flexbuffer type schema from flexbuffers.h"""
|
|
|
|
FBT_NULL = 0
|
|
FBT_INT = 1
|
|
FBT_UINT = 2
|
|
FBT_FLOAT = 3
|
|
# Types above stored inline, types below store an offset.
|
|
FBT_KEY = 4
|
|
FBT_STRING = 5
|
|
FBT_INDIRECT_INT = 6
|
|
FBT_INDIRECT_UINT = 7
|
|
FBT_INDIRECT_FLOAT = 8
|
|
FBT_MAP = 9
|
|
FBT_VECTOR = 10 # Untyped.
|
|
FBT_VECTOR_INT = 11 # Typed any size (stores no type table).
|
|
FBT_VECTOR_UINT = 12
|
|
FBT_VECTOR_FLOAT = 13
|
|
FBT_VECTOR_KEY = 14
|
|
FBT_VECTOR_STRING = 15
|
|
FBT_VECTOR_INT2 = 16 # Typed tuple (no type table, no size field).
|
|
FBT_VECTOR_UINT2 = 17
|
|
FBT_VECTOR_FLOAT2 = 18
|
|
FBT_VECTOR_INT3 = 19 # Typed triple (no type table, no size field).
|
|
FBT_VECTOR_UINT3 = 20
|
|
FBT_VECTOR_FLOAT3 = 21
|
|
FBT_VECTOR_INT4 = 22 # Typed quad (no type table, no size field).
|
|
FBT_VECTOR_UINT4 = 23
|
|
FBT_VECTOR_FLOAT4 = 24
|
|
FBT_BLOB = 25
|
|
FBT_BOOL = 26
|
|
FBT_VECTOR_BOOL = 36 # To Allow the same type of conversion of type to vector type
|
|
|
|
|
|
class FlexBufferDecoder:
|
|
"""
|
|
This implements partial flexbuffer deserialization to be able
|
|
to read custom options. It is not intended to be a general
|
|
purpose flexbuffer deserializer and as such only supports a
|
|
limited number of types and assumes the data is a flat map.
|
|
"""
|
|
|
|
def __init__(self, buffer):
|
|
self.buffer = buffer
|
|
|
|
def indirect_jump(self, offset, byte_width):
|
|
"""Helper function to read the offset value and jump"""
|
|
unpack_str = {1: "<B", 2: "<H", 4: "<I", 8: "<Q"}[byte_width]
|
|
back_jump = struct.unpack(unpack_str, self.buffer[offset : offset + byte_width])[0]
|
|
return offset - back_jump
|
|
|
|
def decode_keys(self, end, size, byte_width):
|
|
"""Decodes the flexbuffer type vector. Map keys are stored in this form"""
|
|
# Keys are strings here. The format is all strings separated by null, followed by back
|
|
# offsets for each of the string. For example, (str1)\0(str1)\0(offset1)(offset2) The end
|
|
# pointer is pointing at the end of all strings
|
|
keys = list()
|
|
for i in range(0, size):
|
|
offset_pos = end + i * byte_width
|
|
start_index = self.indirect_jump(offset_pos, byte_width)
|
|
str_size = self.buffer[start_index:].find(b"\0")
|
|
assert str_size != -1
|
|
s = self.buffer[start_index : start_index + str_size].decode("utf-8")
|
|
keys.append(s)
|
|
return keys
|
|
|
|
def decode_vector(self, end, size, byte_width):
|
|
"""Decodes the flexbuffer vector"""
|
|
# Each entry in the vector can have different datatype. Each entry is of fixed length. The
|
|
# format is a sequence of all values followed by a sequence of datatype of all values. For
|
|
# example - (4)(3.56)(int)(float) The end here points to the start of the values.
|
|
# Each type byte contains: (type << 2) | bit_width, where bit_width determines actual size.
|
|
values = list()
|
|
for i in range(0, size):
|
|
value_type_pos = end + size * byte_width + i
|
|
value_type_packed = self.buffer[value_type_pos]
|
|
value_type = FlexBufferType(value_type_packed >> 2)
|
|
value_bit_width = BitWidth(value_type_packed & 3)
|
|
value_byte_width = 1 << value_bit_width
|
|
value_bytes = self.buffer[
|
|
end + i * byte_width : end + i * byte_width + value_byte_width
|
|
]
|
|
if value_type == FlexBufferType.FBT_BOOL:
|
|
value = bool(value_bytes[0])
|
|
elif value_type == FlexBufferType.FBT_INT:
|
|
fmt = {1: "<b", 2: "<h", 4: "<i", 8: "<q"}[value_byte_width]
|
|
value = struct.unpack(fmt, value_bytes)[0]
|
|
elif value_type == FlexBufferType.FBT_UINT:
|
|
fmt = {1: "<B", 2: "<H", 4: "<I", 8: "<Q"}[value_byte_width]
|
|
value = struct.unpack(fmt, value_bytes)[0]
|
|
elif value_type == FlexBufferType.FBT_FLOAT:
|
|
fmt = {4: "<f", 8: "<d"}[value_byte_width]
|
|
value = struct.unpack(fmt, value_bytes)[0]
|
|
else:
|
|
raise Exception
|
|
values.append(value)
|
|
return values
|
|
|
|
def decode_map(self, end, byte_width, parent_byte_width):
|
|
"""Decodes the flexbuffer map and returns a dict"""
|
|
mid_loc = self.indirect_jump(end, parent_byte_width)
|
|
size_fmt = {1: "<b", 2: "<h", 4: "<i", 8: "<q"}[byte_width]
|
|
map_size = struct.unpack(size_fmt, self.buffer[mid_loc - byte_width : mid_loc])[0]
|
|
|
|
# Find keys
|
|
keys_offset = mid_loc - byte_width * 3
|
|
keys_end = self.indirect_jump(keys_offset, byte_width)
|
|
keys = self.decode_keys(keys_end, map_size, 1)
|
|
|
|
# Find values
|
|
values_end = self.indirect_jump(end, parent_byte_width)
|
|
values = self.decode_vector(values_end, map_size, byte_width)
|
|
return dict(zip(keys, values))
|
|
|
|
def decode(self):
|
|
"""Decode the buffer. Decoding is partially implemented"""
|
|
root_end = len(self.buffer) - 1
|
|
root_byte_width = self.buffer[root_end]
|
|
root_end -= 1
|
|
root_packed_type = self.buffer[root_end]
|
|
root_end -= root_byte_width
|
|
|
|
root_type = FlexBufferType(root_packed_type >> 2)
|
|
byte_width = 1 << BitWidth(root_packed_type & 3)
|
|
|
|
if root_type == FlexBufferType.FBT_MAP:
|
|
return self.decode_map(root_end, byte_width, root_byte_width)
|
|
raise NotImplementedError("Flexbuffer Decoding is partially imlpemented.")
|