Retrain YOLOv5 on Custom Dataset

This guide will walk through how to retrain YOLOv5 on a custom dataset, in particular a subset of coco-2017 that contains only “person”, “car” and “bicycle” classes. For best practices with different datasets refer to:

Tips for Best Training Results

Downloading a subset of COCO 2017

Prerequisites:

  • pycotools
pip install git+https://github.com/waleedka/cocoapi.git#egg=pycocotools&subdirectory=PythonAPI
  • fiftyone
pip install fiftyone

Download the subset of COCO 2017 with “person”, “car” and “bicycle” classes

import fiftyone.zoo as foz
dataset = foz.load_zoo_dataset(
“coco-2017”,
splits=[“train”, “validation”, “test”],
label_type=“detections”,
classes=[“person”, “car”, “bicycle”],
)

Note that this will result in a dataset with annotations that contains at least one of the desired classes, but that may also have undesired ones. Therefore, extra processing of the dataset is required to get rid of these extra labels.

Preparing Dataset for YOLOv5

The dataset must be pre-processed to fit yolov5 requirements.

YOLOv5 dataset format

There should be two folders:

  • Labels: one file named labels containing all categories and one .txt file per image containing one row per object. Each row must have the format: label x_center y_center width
  • Images: all image files.

Formatting

The dataset can be formatted either automatically with Roboflow, or manually. I chose to do it manually because the dataset needed further cleaning of the annotations.

Roboflow

Follow steps from Train Custom Data

Manually

The coco 2017 dataset annotations contain the following format for the bounding box annotation: top left x position, top left y position, width, height

The following functions perform the conversion to yolov5’s bbox format, create a a label file containing all categories and create a .txt file for each image. Each txt file contains a row for each of the image’s annotation, as in the format description.

import os
import json
import shutil
from tqdm import tqdm

def convert_bbox_coco2yolo(img_width, img_height, bbox):
    """
    Convert bounding box from COCO  format to YOLO format

    Parameters
    ----------
    img_width : int
        width of image
    img_height : int
        height of image
    bbox : list[int]
        bounding box annotation in COCO format: 
        [top left x position, top left y position, width, height]

    Returns
    -------
    list[float]
        bounding box annotation in YOLO format: 
        [x_center_rel, y_center_rel, width_rel, height_rel]
    """
    
    # YOLO bounding box format: [x_center, y_center, width, height]
    # (float values relative to width and height of image)
    x_tl, y_tl, w, h = bbox

    dw = 1.0 / img_width
    dh = 1.0 / img_height

    x_center = x_tl + w / 2.0
    y_center = y_tl + h / 2.0

    x = x_center * dw
    y = y_center * dh
    w = w * dw
    h = h * dh

    return [x, y, w, h]
    
def convert_coco_json_to_yolo_txt(output_path, json_file):
    if os.path.exists(output_path):
        shutil.rmtree(output_path)
    path = os.makedirs(output_path)

    with open(json_file) as f:
        json_data = json.load(f)

    label_file = os.path.join(output_path, "labels")
    categories_list = ["person", "car", "bicycle"]
    with open(label_file, "w") as f:
        for category in tqdm(categories_list, desc="Categories"):
            f.write(f"{category}\n")

    for image in tqdm(json_data["images"], desc="Annotation txt for each image"):
        
        img_id = image["id"]
        img_name = image["file_name"]
        img_width = image["width"]
        img_height = image["height"]

        anno_in_image = [anno for anno in json_data["annotations"] if anno["image_id"] == img_id]
        anno_txt = os.path.join(output_path, img_name.split(".")[0] + ".txt")
        with open(anno_txt, "w") as f:
            for anno in anno_in_image:
                cur_cat = anno["category_id"]
                cat_dict = {1:0, 2:1, 3:2}   # creates a dict for the classes 
                if cur_cat in [1,2,3]:       # only keep relevant labels
                    bbox_COCO = anno["bbox"]
                    x, y, w, h = convert_bbox_coco2yolo(img_width, img_height, bbox_COCO)
                    f.write(f"{cat_dict[cur_cat]} {x:.6f} {y:.6f} {w:.6f} {h:.6f}\n")

    print("Converting COCO Json to YOLO txt finished!")
    
convert_coco_json_to_yolo_txt("train_labels_output", "/local/path/fiftyone/coco-2017/train/labels.json")
convert_coco_json_to_yolo_txt("val_labels_output","/local/path/fiftyone/coco-2017/validation/labels.json")

Environment Preparation

Prerequisites:

git clone https://github.com/hailo-ai/hailo_model_zoo.git
  1. Build the docker image:
cd hailo_model_zoo/training/yolov5
docker build --build-arg timezone=`cat /etc/timezone` -t yolov5:v0 .
  • This command will build the docker image with the necessary requirements using the Dockerfile in the yolov5 directory.
  1. Start your docker:

docker run -it --ipc=host -v /path/to/local/drive:/path/to/docker/dir yolov5:v0 --name container_name

  • If there is GPU available for use, use the flag --gpus all to assign all available gpus to the docker container. To assign a specific GPU to the docker container (in case of multiple GPUs available in your machine) use --gpus device=device_number
  • /path/to/local/drive: path to a folder you might want access through the docker, for instance, the datasets folder.
  • /path/to/docker/dir: where the new folder in the docker will be located, it can be for example /workspace/datasets/

After exiting the container you can re-start it and attach with:

docker start -ai container_name
  • To list all containers:
docker ps -a

Folders Organization

For this example we have organized the folders in the following way:

Editing yaml files

Model YAML

cd /workspace/yolov5/models
vim yolov5.yaml
  • Change the number of classes (nc) to the new value:

Note that even though the dataset has 3 classes, the number of classes is 4 to include background.

Dataset YAML

Create a copy of the coco128 yaml, edit the copy by changing the number of classes, class names, and dataset paths.

cd /workspace/yolov5/data
cp coco128.yaml dataset.yaml
vim dataset.yaml

Training

From inside /workspace/yolov5/:

python train.py --img 640 --batch 16 --epochs 300 --data dataset.yaml --weights yolov5s.pt --cfg models/yolov5s.yaml
  • –img: COCO trains at the native resolution of --img 640. If there are many small objects then custom datasets will benefit from training at native or higher resolution.
  • –batch: use the largest amount the hardware allows for.
  • –epochs: make sure that the number is below the overfitting threshold.
  • –data: the name of the created dataset yaml file.
  • –weights:
    • yolov5s.pt - pretrained weights. You can find the pretrained weights for yolov5s, yolov5m, yolov5l, yolov5x in your working directory.
    • '' - without pretrained weights, recommended for large datasets.

You can find the training results in /workspace/yolov5/runs/expn/, where n is the number of the last training run.

Exporting to ONNX

In the folder /workspace/yolov5/runs/ it is possible to see the data related to all training runs

Export at 640x640 with batch size 16:

python models/export.py --weights /path/to/trained/best.pt --img 640 --batch 16
  • –weights: you can find it in /workspace/yolov5/runs/expn/weights, it’s better to use best.pt.

The resulting best.onnx file will be located in /yolov5/runs/expn/weights.

Compiling with Hailo Model Zoo

Prerequisites:

  • A working model zoo environment. If not yet setup, run from inside of the hailo_model_zoo folder:
pip install -e .

Configuration yaml creation

cd /hailo_model_zoo/cfg/networks
cp yolov5s.yaml yolov5s_3classes.yaml
vim yolov5s_3classes.yaml

Add the following lines to the file:

Compilation

Hailo Model Zoo will go through all the required stages to compile the model.

python hailo_model_zoo/main.py compile --ckpt yolov5s.onnx --calib-path /path/to/calibration/imgs/dir/ --yaml yolov5s.yaml
  • –ckpt: the onnx file resulting from the exporting stage.
  • –calib-path: a small set of images, preferably separate set from test, train, val
  • –yaml: path to the yaml that was created on the previous section

The resulting hef file will be saved as yolov5s.hef inside the hailo_model_zoo folder.

Sources for this post:

How can I download a specific part of Coco Dataset?

Annotation Conversion: COCO JSON to YOLO Txt | Haobin Tan

5 Likes