YOLOv7 to Tensorflow

Guangzhou, China

YOLO PyTorch to ONNX
Converting the Onnx Model to Tensorflow
Converting the TensorFlow model to tflite
Predictions using TFlite

The YOLOv7 model created is based on PyTorch. The YOLOv7 Repository already provides 3 export options to CoreML, ONNX and TensorRT. We can use those to - indirectly - transfer our YOLO model to Tensorflow.

YOLO PyTorch to ONNX

Now we can use the export script from the YOLOv7 repository to convert the YOLOv7 model. Since I want to end up at Tensorflow Lite I will try out the yolov7-tiny.pt weights:

wget https://github.com/WongKinYiu/yolov7/releases/download/v0.1/yolov7-tiny.pt

python export.py --weights 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

grid: The grid parameter is an option allowing the export of the detection layer grid.
end2end: This option allows the export of end-to-end ONNX graph which does both bounding box prediction and NMS.
simplify: It is the option by which we can select whether we want to simplify the ONNX graph using reparameterization.
topk-all: It's the option to select the top k object per image using IOU and confidence threshold.
iou-thres: It is the option to set the IOU threshold for NMS.
conf-thres: It is the option to select the confidence threshold score.
img_size / max-wh: These parameters are related to the size of the input image.

Starting TorchScript export with torch 1.13.1+cu117...
Starting TorchScript-Lite export with torch 1.13.1+cu117...
TorchScript-Lite export success, saved as weights/yolov7-tiny.torchscript.ptl
Starting ONNX export with onnx 1.13.0...
Starting to simplify ONNX...
ONNX export success, saved as weights/yolov7-tiny.onnx
Export complete (5.26s)

Converting the Onnx Model to Tensorflow

We use the onnx-tf module to perform the conversion between ONNX and Tensorflow. For the conversion to Tensorflow we need the following ONNX dependencies (I already had Tensorflow installed but was missing tensorflow-probability):

pip install onnx-tf tensorflow-probability

I also had onnx, onnxruntime and onnxsim installed - I am not sure if they are necessary for the conversion.

onnx-tf convert -i weights/yolov7-tiny.onnx -o weights

INFO:onnx-tf:Converting completes successfully.

This now generated the Tensorflow files we need inside the weights directory:

weights
├── assets
├── fingerprint.pb
├── saved_model.pb
├── variables
│   ├── variables.data-00000-of-00001
│   └── variables.index
├── yolov7-tiny.onnx
├── yolov7-tiny.pt

Converting the TensorFlow model to tflite

We use the TensorFlow module’s TFlite converter to convert the model from TF saved model format to a TFLite graph:

tf_to_tflite.py

import tensorflow as tf
converter = tf.lite.TFLiteConverter.from_saved_model('weights/')
tflite_model = converter.convert()

with open('weights/yolov7-tiny.tflite', 'wb') as f:
  f.write(tflite_model)

Predictions using TFlite

tflite_predictions.py

# https://github.com/VikasOjha666/yolov7_to_tflite/blob/main/yoloV7_to_TFlite%20.ipynb

import cv2
import random
import numpy as np
from PIL import Image
import tensorflow as tf

# Load the TFLite model
interpreter = tf.lite.Interpreter(model_path="weights/yolov7-tiny.tflite")


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)

#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)}

#Load and preprocess the image.
img = cv2.imread('guangzhou.jpg')
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

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] - 2),cv2.FONT_HERSHEY_SIMPLEX,0.75,[225, 255, 255],thickness=2)  

prediction = Image.fromarray(ori_images[0])
prediction.show()

YOLOv7 Data Conversion