YOEO — You Only Encode Once

A CNN for Embedded Object Detection and Semantic Segmentation

This project is based upon PyTorch-YOLOv3 and will continuously be modified to our needs.

Installation

Installing from source

For normal training and evaluation we recommend installing the package from source using a poetry virtual environment.

git clone https://github.com/bit-bots/YOEO
cd YOEO/
pip3 install poetry --user
poetry install

You need to join the virtual environment by running poetry shell in this directory before running any of the following commands without the poetry run prefix. Also have a look at the other installing method, if you want to use the commands everywhere without opening a poetry-shell.

Download pretrained weights

./weights/download_weights.sh

Test

Evaluates the model on the test dataset. See help page for more details.

poetry run yoeo-test -h

Inference

Uses pretrained weights to make predictions on images.

poetry run yoeo-detect --images data/samples/

API

You are able to import the modules of this repo in your own project if you install this repo as a python package.

An example prediction call from a simple OpenCV python script would look like this:

import cv2
from yoeo import detect, models

# Load the YOEO model
model = models.load_model(
  "<PATH_TO_YOUR_CONFIG_FOLDER>/yoeo.cfg",
  "<PATH_TO_YOUR_WEIGHTS_FOLDER>/yoeo.pth")

# Load the image as a numpy array
img = cv2.imread("<PATH_TO_YOUR_IMAGE>")

# Convert OpenCV bgr to rgb
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

# Runs the YOEO model on the image
boxes, segmentation = detect.detect_image(model, img)

print(boxes)
# Output will be a numpy array in the following format:
# [[x1, y1, x2, y2, confidence, class]]

print(segmentation)
# Output will be a 2d numpy array with the coresponding class id in each cell

For more advanced usage look at the method's doc strings.

Convert your YOEO model

Convert your YOEO model to an ONNX model

To convert your YOEO model to an ONNX model, you can use the following command:

poetry run yoeo-to-onnx config/yoeo.cfg config/yoeo.pth # Replace path with your .cfg and .pth file

For more information on ONNX, read the ONNX runtime website.

Testing the ONNX Model

poetry run yoeo-run-onnx data/samples/frame6554.jpg config/yoeo.onnx

Parameters:

data/samples/frame6554.jpg: This is the path to the input image you want to test the model with.
config/yoeo.onnx: Make sure to replace this with the actual path to your ONNX model file.

Convert ONNX model to TensorRT

After successful conversion of your YOEO model to an ONNX model using this guide, you can move on with the next conversion to an TensorRT model model using the following command. However, before converting to TensorRT, you need to modify the ONNX model due to compatibility issues with certain operations, the a Cast operation from float to int, which TensorRT does not support.

Here is how you can work around this issue:

poetry run yoeo-fix-cast config/yoeo.onnx config/yoeo_fixed.onnx

Parameters:

config/yoeo.onnx: Make sure to replace this with the actual path to your ONNX model to be fixed.
config/yoeo_fixed.onnx: Path to save the fixed .onnx model.

The onnx2trt.sh script invokes the trtexec, which is a tool provided by TensorRT for converting ONNX models into optimized TensorRT engines. Use the following command to convert the yoeo_fixed.onnx model to TensorRT:

$ sh onnx2trt.sh

For more information on trtexec, read the NVIDIA documentation.

Publication