Export Yolo V7 to Tensorflow Lite

Export Yolo V7 to Tensorflow Lite

Export Yolo V7 to Tensorflow Lite
Photo by Jelleke Vanooteghem / Unsplash

The process of converting the PyTorch model to Tensorflow Lite

PyTorch model -> ONXX model -> Tensorflow Model -> Tensorflow Lite Model

Start by creating a new virtual environment:

python3 -m venv .venv
source .venv/bin/activate

Install basic requirements:

pip install onnx onnxruntime onnxsim onnx-tf

Clone YOLO v7 repository and download official YOLO v7 PyTorch weights:

git clone https://github.com/WongKinYiu/yolov7.git
cd yolov7
wget https://github.com/WongKinYiu/yolov7/releases/download/v0.1/yolov7.pt

Run requirements.txt install within yolov7 folder to install pytorch and the rest of libraries:

pip install -r requirements.txt

PyTorch model to ONXX

Run the export script from the yolov7 folder:

python export.py --weights yolov7-tiny.pt --grid --end2end --simplify \
        --topk-all 100 --iou-thres 0.65 --conf-thres 0.35 --img-size 640 640 --max-wh 640

My output (missing onnx_graphsurgeon but it seem it didn't stop the export).

Install required tensorflow libraries

pip install tensorflow tensorflow_probability

Convert onxx model to tensorflow

mkdir tfmodel
onnx-tf convert -i yolov7.onnx -o tfmodel/

the output should be stored in the newly created folder tfmodel

Convert model from Tensorflow to Tensorflow Lite

Create a small python script and name it tf_model_to_tf_lite.py, then copy paste the contents:

import tensorflow as tf

converter = tf.lite.TFLiteConverter.from_saved_model('tfmodel/')
tflite_model = converter.convert()

with open('tfmodel/yolov7_model.tflite', 'wb') as f:

Run the script:

python3 tf_model_to_tf_lite.py

The output of the script should be Tensorflow Lite model named tfmodel/yolov7_model.tflite

Run inference

import cv2
import random
import numpy as np
import tensorflow as tf
import matplotlib.pyplot as plt

# Load the TFLite model and allocate tensors.
interpreter = tf.lite.Interpreter(model_path="tfmodel/yolov7_model.tflite")

# Name of the classes according to class indices.
names = [
    "traffic light",
    "fire hydrant",
    "stop sign",
    "parking meter",
    "sports ball",
    "baseball bat",
    "baseball glove",
    "tennis racket",
    "wine glass",
    "hot dog",
    "potted plant",
    "dining table",
    "cell phone",
    "teddy bear",
    "hair drier",

# Creating random colors for bounding box visualization.
colors = {
    name: [random.randint(0, 255) for _ in range(3)] for i, name in enumerate(names)

img = cv2.imread("data/image2.jpg")
image = img.copy()
image, ratio, dwdh = letterbox(image, auto=False)
image = image.transpose((2, 0, 1))
image = np.expand_dims(image, 0)
image = np.ascontiguousarray(image)

im = image.astype(np.float32)
im /= 255

 Allocate tensors.
# Get input and output tensors.
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

# Test the model on random input data.
input_shape = input_details[0]["shape"]
interpreter.set_tensor(input_details[0]["index"], im)


# The function `get_tensor()` returns a copy of the tensor data.
# Use `tensor()` in order to get a pointer to the tensor.
output_data = interpreter.get_tensor(output_details[0]["index"])

ori_images = [img.copy()]

for i, (batch_id, x0, y0, x1, y1, cls_id, score) in enumerate(output_data):
    image = ori_images[int(batch_id)]
    box = np.array([x0, y0, x1, y1])
    box -= np.array(dwdh * 2)
    box /= ratio
    box = box.round().astype(np.int32).tolist()
    cls_id = int(cls_id)
    score = round(float(score), 3)
    name = names[cls_id]
    color = colors[name]
    name += " " + str(score)
    cv2.rectangle(image, box[:2], box[2:], color, 2)
        (box[0], box[1] + 20),
        [225, 255, 255],
# plt.imshow(ori_images[0])
cv2.imshow("test", ori_images[0])

# Destroys all the windows created

Helper method (resizing and letter-boxing the original image): 

def letterbox(
    im, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleup=True, stride=32
    # Resize and pad image while meeting stride-multiple constraints
    shape = im.shape[:2]  # current shape [height, width]
    if isinstance(new_shape, int):
        new_shape = (new_shape, new_shape)

    # Scale ratio (new / old)
    r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
    if not scaleup:  # only scale down, do not scale up (for better val mAP)
        r = min(r, 1.0)

    # Compute padding
    new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
    dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1]  # wh padding

    if auto:  # minimum rectangle
        dw, dh = np.mod(dw, stride), np.mod(dh, stride)  # wh padding

    dw /= 2  # divide padding into 2 sides
    dh /= 2

    if shape[::-1] != new_unpad:  # resize
        im = cv2.resize(im, new_unpad, interpolation=cv2.INTER_LINEAR)
    top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
    left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
    im = cv2.copyMakeBorder(
        im, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color
    )  # add border
    return im, r, (dw, dh)