o i e|4@sdZddlZddlZddlZddlZddlZddlZddlmZ ddlm Z ddl m Z dZ ddZd d Zd d Zd(ddZddZddZddZd)ddZddZddZddZddZd d!Zd"d#Zd$d%Zd&d'ZdS)*zUtility functions for FlatBuffers. All functions that are commonly used to work with FlatBuffers. Refer to the tensorflow lite flatbuffer schema here: tensorflow/lite/schema/schema.fbs N)schema_py_generated) schema_util)gfilesTFL3cCstj|d}tj|S)zBConverts a tflite model from a bytearray to an object for parsing.r) schema_fbModelGetRootAsModelModelT InitFromObj)model_bytearray model_objectr U/var/www/myenv/lib/python3.10/site-packages/tensorflow/lite/tools/flatbuffer_utils.pyconvert_bytearray_to_object&s rcCspt|s td|t|d}t|}Wdn1s"wYt|}tjdkr6t |dd|S)a2Reads a tflite model as a python object. Args: input_tflite_file: Full path name to the input tflite file Raises: RuntimeError: If input_tflite_file path is invalid. IOError: If input_tflite_file cannot be opened. Returns: A python object corresponding to the input tflite file. zInput file not found at %r rbNbiglittle) rExists RuntimeErrorGFile bytearrayreadrsys byteorderbyte_swap_tflite_model_obj)input_tflite_fileinput_file_handler modelr r r read_model,s   rcCstt|S)aCReads a tflite model as a python object with mutable tensors. Similar to read_model() with the addition that the returned object has mutable tensors (read_model() returns an object with immutable tensors). NOTE: This API only works for TFLite generated with _experimental_use_buffer_offset=false Args: input_tflite_file: Full path name to the input tflite file Raises: RuntimeError: If input_tflite_file path is invalid. IOError: If input_tflite_file cannot be opened. Returns: A mutable python object corresponding to the input tflite file. )copydeepcopyr)rr r r read_model_with_mutable_tensorsCsr cCs:td}||}|j|tdt|}||}|S)z@Converts a tflite model from an object to a immutable bytearray.i)file_identifier) flatbuffersBuilderPackFinish_TFLITE_FILE_IDENTIFIERbytesOutput)r extra_bufferbuilder model_offsetr r r r convert_object_to_bytearrayYs   r-cCsdtjdkrt|}t|ddt|}t|d}||WddS1s+wYdS)asWrites the tflite model, a python object, into the output file. NOTE: This API only works for TFLite generated with _experimental_use_buffer_offset=false Args: model_object: A tflite model as a python object output_tflite_file: Full path name to the output tflite file. Raises: IOError: If output_tflite_file path is invalid or cannot be opened. rrwbN) rrrrrr-rrwrite)r output_tflite_filer output_file_handler r r write_modelds   "r2cCs4d|_|jD]}d|_|jD]}d|_qqd|_dS)ayStrips all nonessential strings from the model to reduce model size. We remove the following strings: (find strings by searching ":string" in the tensorflow lite flatbuffer schema) 1. Model description 2. SubGraph name 3. Tensor names We retain OperatorCode custom_code and Metadata name. Args: model: The model from which to remove nonessential strings. N) description subgraphsnametensors signatureDefs)rsubgraphtensorr r r strip_stringsys   r:cCs*tjjD] \}}||kr|SqdS)z4Converts a numerical enum to a readable tensor type.N)r TensorType__dict__items) tensor_typer5valuer r r type_to_names r@cs6t||j}tdt|}durfdd|D}i}|jD]"}|jD]}|jdur/n|jD]}|j|} t | j || j <q2q&q!|D]R} || j } | durSdn| j } | dkr[qF|| d} | dr| dkrld nd }td| t|D]}td d }t|| ||qwqFt| D] }tdd | |<qqFdS)aRandomize weights in a model. Args: model: The model in which to randomize weights. random_seed: The input to the random number generator (default value is 0). buffers_to_skip: The list of buffer indices to skip. The weights in these buffers are left unmodified. Ncsg|]}|vr|qSr r ).0idxbuffers_to_skipr r z%randomize_weights..rINT8FLOATFLOAT16efgg?)randomseedbuffersrangelenr4 operatorsinputsr6r@typebufferdatasizeget startswithstructcalcsizeuniform pack_intorandint)r random_seedrErP buffer_ids buffer_typesgraphop input_idxr9i buffer_i_data buffer_i_size buffer_type format_codeoffsetr?jr rDr randomize_weightss>           rmcCs:|jD]}|jr|jd}||vr||d|_qdS)zRename custom ops so they use the same naming style as builtin ops. Args: model: The input tflite model. map_custom_op_renames: A mapping from old to new custom op names. asciiN) operatorCodes customCodedecodeencode)rmap_custom_op_renamesop_code op_code_strr r r rename_custom_opss  rvcCsD|j|}t|j|j}ttjD] \}}||kr|SqdS)zConverts a TFLite op_code to the human readable name. Args: model: The input tflite model. op_code: The op_code to resolve to a readable name. Returns: A string containing the human readable op name, or None if not resolvable. N)romax builtinCodedeprecatedBuiltinCodevarsrBuiltinOperatorr=)rrtrdcoder5r?r r r opcode_to_names r}c Cstd}t}t|3}|D]%}||}|durq|d}td|d}dd|D}||qWdt |S1sBwYt |S)a@Converts xxd output C++ source file to bytes (immutable). Args: input_cc_file: Full path name to th C++ source file dumped by xxd Raises: RuntimeError: If input_cc_file path is invalid. IOError: If input_cc_file cannot be opened. Returns: A bytearray corresponding to the input cc file array. z\W*(0x[0-9a-fA-F,x ]+).*NrA,cSsg|]}t|ddqS))base)int)rBxr r r rF rGz'xxd_output_to_bytes..) recompileropenmatchgroupfiltersplitextendr() input_cc_filepatternr file_handleline values_match list_text values_textvaluesr r r xxd_output_to_bytess      rcCst|}t|S)a7Converts xxd output C++ source file to object. Args: input_cc_file: Full path name to th C++ source file dumped by xxd Raises: RuntimeError: If input_cc_file path is invalid. IOError: If input_cc_file cannot be opened. Returns: A python object corresponding to the input tflite file. )rr)r model_bytesr r r xxd_output_to_objects rcsDfddtdtjD}dfdd|D_dS)z4Helper function for byte-swapping the buffers field.csg|] }j||qSr )rW)rBrf)rV chunksizer r rF%sz,byte_swap_buffer_content..rr!cs g|] }t|qSr )r from_bytesto_bytes)rBbyteswap)r from_endiness to_endinessr r rF*s  N)rQrRrWjoin)rVrrrto_swapr )rVrrrr byte_swap_buffer_content#s  rc Cs$|durdSg}tjjtjjtjjg}tjjtjjtjjtjjg}tjj tjj tjj tjj g}|j D]]}|jD]W}|jdkr|jt|jkr|j|vr|j|jjdur|j|vret|j|jd||n#|j|vrvt|j|jd||n|j|vrt|j|jd||nq7||jq7q2dS)aByte swaps the buffers field in a TFLite model. Args: model: TFLite model object of from_endiness format. from_endiness: The original endianness format of the buffers in model. to_endiness: The destined endianness format of the buffers in model. Nr)rr;rJINT16UINT16FLOAT32INT32 COMPLEX64UINT32INT64FLOAT64 COMPLEX128UINT64r4r6rVrRrPrWrUrappend) rrrbuffer_swappedtypes_of_16_bitstypes_of_32_bitstypes_of_64_bitsr8r9r r r r3sP        rcCs(|durdSt|}t|||t|S)aGenerates a new model byte array after byte swapping its buffers field. Args: tflite_model: TFLite flatbuffer in a byte array. from_endiness: The original endianness format of the buffers in tflite_model. to_endiness: The destined endianness format of the buffers in tflite_model. Returns: TFLite flatbuffer in a byte array, after being byte swapped to to_endiness format. N)rrr-) tflite_modelrrrr r r byte_swap_tflite_bufferis   rcCsrt|tjs t|}t}|jD]$}|jdurq|jD]}t|j |j }|tj j kr3| |jjqqt|S)zCalculates the number of unique resource variables in a model. Args: model: the input tflite model, either as bytearray or object. Returns: An integer number representing the number of unique resource variables. N) isinstancerrrsetr4rSr#get_builtin_code_from_operator_codero opcodeIndexr{ VAR_HANDLEaddbuiltinOptions sharedNamerR)runique_shared_namesr8rd builtin_coder r r count_resource_variabless      r)r!)rN)__doc__rrNrr[rr#tensorflow.lite.pythonrrrtensorflow.python.platformrr'rrr r-r2r:r@rmrvr}rrrrrrr r r r s6      3'6