Prepraing Custom Dataset

Organize VOC-format dataset with this directory structure:

VOC_CUSTOM/
├── Annotations/
├── ImageSets/
│   └── Main/
└── JPEGImages/

Execute dataset splitting script:

# dataset_split.py
import os
import random
import argparse

parser = argparse.ArgumentParser()
parser.add_argument('--xml_path', default='Annotations/', type=str)
parser.add_argument('--txt_path', default='ImageSets/Main/', type=str)
args = parser.parse_args()

train_ratio = 0.8
val_ratio = 0.1
test_ratio = 0.1

xml_files = os.listdir(args.xml_path)
total_count = len(xml_files)
indices = list(range(total_count))

train_count = int(total_count * train_ratio)
val_count = int(total_count * val_ratio)

train_indices = random.sample(indices, train_count)
remaining = [i for i in indices if i not in train_indices]
val_indices = random.sample(remaining, val_count)
test_indices = [i for i in remaining if i not in val_indices]

with open(f'{args.txt_path}/train.txt', 'w') as f_train, \
     open(f'{args.txt_path}/val.txt', 'w') as f_val, \
     open(f'{args.txt_path}/test.txt', 'w') as f_test:
    
    for idx in train_indices:
        f_train.write(f'{xml_files[idx][:-4]}\n')
    for idx in val_indices:
        f_val.write(f'{xml_files[idx][:-4]}\n')
    for idx in test_indices:
        f_test.write(f'{xml_files[idx][:-4]}\n')

Convert VOC annotations to YOLO format:

# voc_to_yolo.py
import xml.etree.ElementTree as ET
import os

sets = ['train', 'val', 'test']
class_names = ["object1", "object2"]
base_path = '/path/to/dataset'

def convert(size, box):
    dw = 1. / size[0]
    dh = 1. / size[1]
    x = (box[0] + box[1]) / 2.0 - 1
    y = (box[2] + box[3]) / 2.0 - 1
    w = box[1] - box[0]
    h = box[3] - box[2]
    x *= dw
    w *= dw
    y *= dh
    h *= dh
    return x, y, w, h

def process_xml(image_id):
    in_file = open(f'{base_path}/VOC_CUSTOM/Annotations/{image_id}.xml')
    out_file = open(f'{base_path}/VOC_CUSTOM/labels/{image_id}.txt', 'w')
    
    tree = ET.parse(in_file)
    root = tree.getroot()
    size = root.find('size')
    w = int(size.find('width').text)
    h = int(size.find('height').text)
    
    for obj in root.iter('object'):
        cls_name = obj.find('name').text
        if cls_name not in class_names:
            continue
        cls_id = class_names.index(cls_name)
        bbox = obj.find('bndbox')
        box = (
            float(bbox.find('xmin').text),
            float(bbox.find('xmax').text),
            float(bbox.find('ymin').text),
            float(bbox.find('ymax').text)
        )
        bb = convert((w, h), box)
        out_file.write(f"{cls_id} {' '.join(map(str, bb))}\n")

os.makedirs(f'{base_path}/VOC_CUSTOM/labels/', exist_ok=True)
for subset in sets:
    with open(f'{base_path}/VOC_CUSTOM/ImageSets/Main/{subset}.txt') as f:
        image_ids = f.read().splitlines()
    with open(f'{base_path}/VOC_CUSTOM/{subset}.txt', 'w') as list_file:
        for img_id in image_ids:
            list_file.write(f'{base_path}/VOC_CUSTOM/JPEGImages/{img_id}.jpg\n')
            process_xml(img_id)

Model Training

Set up environment:

git clone https://github.com/ultralytics/yolov5
cd yolov5
pip install -r requirements.txt

Download pretarined weights:

wget https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5s.pt

Start training:

python train.py --weights yolov5s.pt \
                --cfg models/yolov5s.yaml \
                --data data/custom.yaml \
                --epochs 150 \
                --batch-size 16 \
                --device 0

Test inference:

python detect.py --source test_image.jpg \
                 --weights runs/train/exp/weights/best.pt \
                 --conf-thres 0.25

Model Conversion

Clone Rockchip-modified YOLOv5 repository:

git clone https://github.com/airockchip/yolov5.git
cd yolov5
pip install -r requirements.txt

Export to ONNX:

python export.py --rknpu --weight custom_model.pt

Prepare conversion files:

1. Create quantization dataset list dataset.txt
2. Prepare label file labels.txt
3. Modify conversion script paths

Convert to RKNN format:

python convert.py custom_model.onnx rk3588

Build executable:

./build-linux.sh -t rk3588 -a aarch64 -d yolov5

On-Device Deployment

Transfer and execute on ELF2 board:

tar -zxvf rknn_demo.tar.gz
cd rknn_demo
chmod +x yolov5_demo
./yolov5_demo model/custom.rknn test_image.jpg

Output example:

load label ./labels.txt
input: 640x640
output: detected objects
output image saved as result.jpg