# Copyright 2023-2024 NXP # SPDX-License-Identifier: BSD-3-Clause import os import numpy as np import struct import tensorflow as tf from tensorflow.keras.layers import Conv2D, Input, LeakyReLU from tensorflow.keras.layers import ZeroPadding2D, UpSampling2D from tensorflow.keras.layers import MaxPool2D, add, concatenate from tensorflow.keras.models import Model import argparse import PIL.Image as im import random random.seed(42) N_CALIBRATION_IMAGES = 100 parser = argparse.ArgumentParser() parser.add_argument('--weights_path', help='path to darknet weights') parser.add_argument('--output_path', help='path to save tflite model') parser.add_argument('--images_path', help='path to representative images for quantization', default=None) args = parser.parse_args() def _conv_block(inp, convs, skip=False): x = inp count = 0 for conv in convs: if count == (len(convs) - 2) and skip: skip_connection = x count += 1 if conv['stride'] > 1: x = ZeroPadding2D(((1, 0), (1, 0)), name='zerop_' + str(conv['layer_idx']))( x) # peculiar padding as darknet prefer left and top x = Conv2D(conv['filter'], conv['kernel'], strides=conv['stride'], # peculiar padding as darknet prefer left and top padding='valid' if conv['stride'] > 1 else 'same', name='convn_' + str(conv['layer_idx']) \ if conv['bnorm'] else 'conv_' + str(conv['layer_idx']), activation=None, use_bias=True)(x) if conv['activ'] == 1: x = LeakyReLU(alpha=0.1, name='leaky_' + str(conv['layer_idx']))(x) return add([skip_connection, x], name='add_' + str(conv['layer_idx'] + 1)) if skip else x def _split_block(input_layer, layer_idx): s = tf.split(input_layer, num_or_size_splits=2, axis=-1, name=f"split_{layer_idx}") return s[1] def make_yolov4_tiny_model(): input_image = Input(shape=(416, 416, 3), batch_size=1, name='input_0') # Layer 0 x = _conv_block(input_image, [{'filter': 32, 'kernel': 3, 'stride': 2, 'bnorm': True, 'activ': 1, 'layer_idx': 0}]) layer_0 = x # Layer 1 x = _conv_block(x, [{'filter': 64, 'kernel': 3, 'stride': 2, 'bnorm': True, 'activ': 1, 'layer_idx': 1}]) layer_1 = x # Layer 2, concat1 x = _conv_block(x, [{'filter': 64, 'kernel': 3, 'stride': 1, 'bnorm': True, 'activ': 1, 'layer_idx': 2}]) layer_2 = x # Layer 3, route group x = _split_block(x, layer_idx=3) # Layer 4, concat_route_1 x = _conv_block(x, [{'filter': 32, 'kernel': 3, 'stride': 1, 'bnorm': True, 'activ': 1, 'layer_idx': 4}]) layer_4 = x # Layer 5, concat_route_2 x = _conv_block(x, [{'filter': 32, 'kernel': 3, 'stride': 1, 'bnorm': True, 'activ': 1, 'layer_idx': 5}]) layer_5 = x # Layer 6, concat route x = concatenate([layer_5, layer_4], axis=-1, name='concat_6') # Layer 7, concat2 x = _conv_block(x, [{'filter': 64, 'kernel': 1, 'stride': 1, 'bnorm': True, 'activ': 1, 'layer_idx': 7}]) layer_7 = x # Layer 8, concat x = concatenate([layer_2, layer_7], axis=-1, name='concat_8') # Layer 9 x = MaxPool2D(pool_size=(2, 2), padding='same', name='layer_9')(x) # Layer 10, concat 1 x = _conv_block(x, [{'filter': 128, 'kernel': 3, 'stride': 1, 'bnorm': True, 'activ': 1, 'layer_idx': 10}]) layer_10 = x # Layer 11 x = _split_block(x, layer_idx=11) # Layer 12, concat route 1 x = _conv_block(x, [{'filter': 64, 'kernel': 3, 'stride': 1, 'bnorm': True, 'activ': 1, 'layer_idx': 12}]) layer_12 = x # Layer 13, concat route 2 x = _conv_block(x, [{'filter': 64, 'kernel': 3, 'stride': 1, 'bnorm': True, 'activ': 1, 'layer_idx': 13}]) layer_13 = x # Layer 14 x = concatenate([layer_13, layer_12], axis=-1, name='concat_14') # Layer 15, concat 2 x = _conv_block(x, [{'filter': 128, 'kernel': 1, 'stride': 1, 'bnorm': True, 'activ': 1, 'layer_idx': 15}]) layer_15 = x # Layer 16 x = concatenate([layer_10, layer_15], axis=-1, name='concat_16') # Layer 17 x = MaxPool2D(pool_size=(2, 2), padding='same', name='layer_17')(x) # Layer 18, concat 1 x = _conv_block(x, [{'filter': 256, 'kernel': 3, 'stride': 1, 'bnorm': True, 'activ': 1, 'layer_idx': 18}]) layer_18 = x # Layer 19 x = _split_block(x, layer_idx=19) # Layer 20, concat route 1 x = _conv_block(x, [{'filter': 128, 'kernel': 3, 'stride': 1, 'bnorm': True, 'activ': 1, 'layer_idx': 20}]) layer_20 = x # Layer 21, concat route 2 x = _conv_block(x, [{'filter': 128, 'kernel': 3, 'stride': 1, 'bnorm': True, 'activ': 1, 'layer_idx': 21}]) layer_21 = x # Layer 22 x = concatenate([layer_21, layer_20], axis=-1, name='concat_22') # Layer 23, concat 2, output 1 of cspdarknet x = _conv_block(x, [{'filter': 256, 'kernel': 1, 'stride': 1, 'bnorm': True, 'activ': 1, 'layer_idx': 23}]) layer_23 = x # Layer 24 x = concatenate([layer_18, layer_23], axis=-1, name='concat_24') # Layer 25 x = MaxPool2D(pool_size=(2, 2), padding='same', name='layer_25')(x) # Layer 26, output 2 of cspdarknet x = _conv_block(x, [{'filter': 512, 'kernel': 3, 'stride': 1, 'bnorm': True, 'activ': 1, 'layer_idx': 26}]) layer_26 = x # After backbone # Layer 27, concat 1, branch 1 x = _conv_block(layer_26, [{'filter': 256, 'kernel': 1, 'stride': 1, 'bnorm': True, 'activ': 1, 'layer_idx': 27}]) layer_27 = x # Layer 28 x = _conv_block(x, [{'filter': 512, 'kernel': 3, 'stride': 1, 'bnorm': True, 'activ': 1, 'layer_idx': 28}]) layer_28 = x # Layer 29, output of large grid x = _conv_block(x, [{'filter': 255, 'kernel': 1, 'stride': 1, 'bnorm': True, 'activ': 0, 'layer_idx': 29}]) layer_29 = x # Layer 30, continue from layer_27 x = _conv_block(layer_27, [{'filter': 128, 'kernel': 1, 'stride': 1, 'bnorm': True, 'activ': 1, 'layer_idx': 30}]) layer_30 = x # Layer 31 x = UpSampling2D(size=(2, 2), name='upsamp_31', interpolation='bilinear')(x) layer_31 = x # Layer 32 x = concatenate([layer_31, layer_23], axis=-1, name='concat_32') # Layer 33 x = _conv_block(x, [{'filter': 256, 'kernel': 3, 'stride': 1, 'bnorm': True, 'activ': 1, 'layer_idx': 33}]) # Layer 34, output of medium grid x = _conv_block(x, [{'filter': 255, 'kernel': 1, 'stride': 1, 'bnorm': True, 'activ': 0, 'layer_idx': 34}]) layer_34 = x # End model = Model(input_image, [layer_34, layer_29], name='Yolov4-tiny') model.summary() return model # Define the model model = make_yolov4_tiny_model() model.summary() # load weights in keras class WeightReader: def __init__(self, weight_file): with open(weight_file, 'rb') as w_f: major, = struct.unpack('i', w_f.read(4)) minor, = struct.unpack('i', w_f.read(4)) revision, = struct.unpack('i', w_f.read(4)) if (major * 10 + minor) >= 2 and major < 1000 and minor < 1000: print("reading 64 bytes") w_f.read(8) else: print("reading 32 bytes") w_f.read(4) transpose = (major > 1000) or (minor > 1000) binary = w_f.read() self.offset = 0 self.all_weights = np.frombuffer(binary, dtype='float32') print(f"weight total length {len(self.all_weights)}") def read_bytes(self, size): self.offset = self.offset + size return self.all_weights[self.offset - size:self.offset] def load_weights(self, model): count = 0 ncount = 0 for i in range(35): try: conv_layer = model.get_layer('convn_' + str(i)) filter = conv_layer.kernel.shape[-1] # kernel*kernel*c*filter nweights = np.prod(conv_layer.kernel.shape) print(f"loading weights of convolution #" + str(i) + "- nb parameters: " + str(nweights + filter)) if i in [29, 34]: bias = self.read_bytes(filter) # bias weights = self.read_bytes(nweights) # weights else: bias = self.read_bytes(filter) # bias scale = self.read_bytes(filter) # scale mean = self.read_bytes(filter) # mean var = self.read_bytes(filter) # variance weights = self.read_bytes(nweights) # weights # normalize bias bias = bias - scale * mean / (np.sqrt(var + 0.00001)) # normalize weights weights = np.reshape(weights, (filter, int(nweights / filter))) A = scale / (np.sqrt(var + 0.00001)) A = np.expand_dims(A, axis=0) weights = weights * A.T weights = np.reshape(weights, (nweights)) shp = list(reversed(conv_layer.get_weights()[0].shape)) weights = weights.reshape(shp) weights = weights.transpose([2, 3, 1, 0]) if len(conv_layer.get_weights()) > 1: a = conv_layer.set_weights([weights, bias]) else: a = conv_layer.set_weights([weights]) count = count + 1 ncount = ncount + nweights + filter except ValueError: print("no convolution #" + str(i)) print(count, "Convolution Normalized Layers are loaded with ", ncount, " parameters") def reset(self): self.offset = 0 darknet_model = args.weights_path + '/yolov4-tiny.weights' weight_reader = WeightReader(darknet_model) weight_reader.load_weights(model) def image_resize(image, resize_shape): image_copy = np.copy(image) resize_h, resize_w = resize_shape orig_h, orig_w, _ = image_copy.shape scale = min(resize_h / orig_h, resize_w / orig_w) temp_w, temp_h = int(scale * orig_w), int(scale * orig_h) image_resized = image.resize((temp_w, temp_h), im.BILINEAR) image_paded = np.full(shape=[resize_h, resize_w, 3], fill_value=128.0) r_w = (resize_w - temp_w) // 2 # real_w r_h = (resize_h - temp_h) // 2 # real_h image_paded[r_h:temp_h + r_h, r_w:temp_w + r_w, :] = image_resized image_paded = image_paded / 255. return image_paded def representative_dataset(): _, h, w, _ = model.input_shape image_folder = args.images_path image_files = os.listdir(image_folder) random.shuffle(image_files) image_files = image_files[:N_CALIBRATION_IMAGES] for image_file in image_files: image_path = os.path.join(image_folder, image_file) original_image = im.open(image_path) if original_image.mode != "RGB": continue image_data = image_resize(original_image, [h, w]) img_in = image_data[np.newaxis, ...].astype(np.float32) yield [img_in] def dummy_dataset(): _, h, w, _ = model.input_shape for i in range(N_CALIBRATION_IMAGES): # Tensorflow basic format : NHWC img_in = np.random.randn(1, h, w, 3).astype('float32') yield [img_in] converter = tf.lite.TFLiteConverter.from_keras_model(model) # quantized model tflite_quant = args.output_path + '/yolov4-tiny_416_quant.tflite' converter.optimizations = [tf.lite.Optimize.DEFAULT] if args.images_path is not None: converter.representative_dataset = representative_dataset else: # Dummy dataset if no representative dataset is given converter.representative_dataset = dummy_dataset converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8] converter.inference_input_type = tf.int8 converter.inference_output_type = tf.float32 tflite_model = converter.convert() with open(tflite_quant, 'wb') as f: f.write(tflite_model) # float32 model converter = tf.lite.TFLiteConverter.from_keras_model(model) tflite_float = args.output_path + '/yolov4-tiny_416_float32.tflite' tflite_model = converter.convert() with open(tflite_float, 'wb') as f: f.write(tflite_model)