Export Yolo V7 to Tensorflow Lite
Export Yolo V7 to Tensorflow Lite
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/activateInstall basic requirements:
pip install onnx onnxruntime onnxsim onnx-tfClone 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.ptRun requirements.txt install within yolov7 folder to install pytorch and the rest of libraries:
pip install -r requirements.txtPyTorch 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 640My output (missing onnx_graphsurgeon but it seem it didn't stop the export).
Install required tensorflow libraries
pip install tensorflow tensorflow_probabilityConvert 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:
f.write(tflite_model)Run the script:
python3 tf_model_to_tf_lite.pyThe 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 = [
"person",
"bicycle",
"car",
"motorcycle",
"airplane",
"bus",
"train",
"truck",
"boat",
"traffic light",
"fire hydrant",
"stop sign",
"parking meter",
"bench",
"bird",
"cat",
"dog",
"horse",
"sheep",
"cow",
"elephant",
"bear",
"zebra",
"giraffe",
"backpack",
"umbrella",
"handbag",
"tie",
"suitcase",
"frisbee",
"skis",
"snowboard",
"sports ball",
"kite",
"baseball bat",
"baseball glove",
"skateboard",
"surfboard",
"tennis racket",
"bottle",
"wine glass",
"cup",
"fork",
"knife",
"spoon",
"bowl",
"banana",
"apple",
"sandwich",
"orange",
"broccoli",
"carrot",
"hot dog",
"pizza",
"donut",
"cake",
"chair",
"couch",
"potted plant",
"bed",
"dining table",
"toilet",
"tv",
"laptop",
"mouse",
"remote",
"keyboard",
"cell phone",
"microwave",
"oven",
"toaster",
"sink",
"refrigerator",
"book",
"clock",
"vase",
"scissors",
"teddy bear",
"hair drier",
"toothbrush",
]
# 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.
interpreter.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)
interpreter.invoke()
# 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)
cv2.putText(
image,
name,
(box[0], box[1] + 20),
cv2.FONT_HERSHEY_SIMPLEX,
0.75,
[225, 255, 255],
thickness=2,
)
# plt.imshow(ori_images[0])
cv2.imshow("test", ori_images[0])
cv2.waitKey(0)
# Destroys all the windows created
cv2.destroyAllWindows()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)