Vision Model Training

SpireCV Wiki

Official SpireCV Vision Model Training Tutorial

How to Train a Detector with a Custom Dataset

1 Data Annotation

1. Download the SpireView annotation software from https://gitee.com/jario-jin/SpireView. We recommend downloading the latest version from Baidu Netdisk and extracting it.
2. Double-click SpireView.exe to open the annotation software. Click Tools → Settings..., then enter an English-only path in the saving path field. Annotation files will be saved to this path.

Note

All annotation files will be stored in this folder.

3. Click Input → Image Dir, select the folder containing the images to annotate, and import the images.
4. Click Tools → Annotate Image and select an annotation type. Box Label creates bounding-box annotations, while Instance Label creates segmentation annotations.

5. In the dialog box that appears, change label to the target class you are annotating. You do not need to close this dialog box. After annotating one image, use the left/right arrow keys (←/→) to move to the previous/next image. Annotation results are saved automatically.
6. Use the mouse wheel to zoom the image. Hold down the left mouse button and drag to move the visible image area. For bounding-box annotations, left-click two points to enclose the target in a box.

7. While annotating, if you click the wrong point, right-click to cancel it. After annotation is complete, if you are not satisfied, left-click the green border (it will turn red, as shown below), then press Delete to remove it.

8. To annotate multiple classes, change the label as needed.

9. After completing all annotations, press Ctrl+O to select an annotation format and export the results. Here, select Yolo detection format. If you are training a segmentation network, also select output segs. Click OK to generate YOLO-format training files in the save path. YOLO training requires the scaled_images folder, the Yolo_labels folder, and the Yolo_categories.names file.

Note

The following two folders are required for training.

2 Train the Model

1. On the training server/workstation (with an NVIDIA GPU that has at least 12 GB of VRAM; Ubuntu is recommended, with CUDA 10.2+ and PyTorch 1.7+ installed), create the folder [PATH-TO-YOUR-DATA]/images/train. Store the training images from the scaled_images folder generated in the preceding step in this folder. Create the folder [PATH-TO-YOUR-DATA]/labels/train, and store the training annotations from the Yolo_labels folder in it.
2. Download the YOLOv5-v7 code.

git clone https://gitee.com/jario-jin/yolov5-v7.git
cd yolov5-v7
pip3 install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple/ 

3. To train an object-detection network, copy [PATH-TO-yolov5-v7]/data/coco128.yaml and give the copy a custom name. To train an instance-segmentation network, copy [PATH-TO-yolov5-v7]/data/coco128-seg.yaml and give the copy a custom name. In the .yaml file, modify path, train, and val; then modify names according to the class names in the Yolo_categories.names file, as shown below.

# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
path: [PATH-TO-YOUR-DATA]  # dataset root dir
train: images/train  # train images (relative to 'path') 128 images
val: images/val  # val images (relative to 'path') 128 images
test:  # test images (optional)

# Classes
names:
  0: person
  1: car

4. To train an object bounding-box detection network, run:

cd [PATH-TO-yolov5-v7]
python3 train.py \
--weights weights/yolov5s6.pt \
--cfg models/hub/yolov5s6.yaml \
--data [PATH-TO-YOUR-yaml] \
--hyp data/hyps/hyp.scratch-med.yaml \
--epochs 50 \
--batch-size 8 \
--imgsz 1280

5. To train an instance-segmentation network, run:

cd [PATH-TO-yolov5-v7]
python3 segment/train.py \
--weights weights/yolov5s-seg.pt \
--cfg models/segment/yolov5s-seg.yaml \
--data [PATH-TO-YOUR-yaml] \
--hyp data/hyps/hyp.scratch-low.yaml \
--epochs 50 \
--batch-size 8 \
--imgsz 640

6. The parameters in steps 4 and 5 are described below:
   • --weights: The pretrained network model used to initialize the network weights. Download the weights from the following URLs and place them in [PATH-TO-yolov5-v7]/weights.
     • https://download.amovlab.com/model/yolov5-7/yolov5s.pt
     • https://download.amovlab.com/model/yolov5-7/yolov5s6.pt
     • https://download.amovlab.com/model/yolov5-7/yolov5s-seg.pt
   • --cfg: Defines the network architecture, such as models/yolov5.yaml, models/hub/yolov5s6.yaml, or models/segment/yolov5s-seg.yaml.
   • --data: The training-data description file—the custom .yaml file created above.
   • --hyp: Hyperparameter settings for network training.
   • --epochs: Number of training epochs.
   • --batch-size: Number of images input to the network simultaneously during training.
   • --imgsz: Input image size during training, either 640 or 1280.

3 Deploy the Application

3.1 Convert a .pt Model to .wts on the Training Platform

1. After training is complete:
   • Detection training results are located in [PATH-TO-yolov5-v7]/runs/train.
   • Segmentation training results are located in [PATH-TO-yolov5-v7]/runs/train-seg.
2. Convert the model by running the following code. Modify the path to the .pt weights file and the corresponding name of the generated .wts file as needed.

cd [PATH-TO-yolov5-v7]
python3 gen_wts.py -w runs/train/exp/weights/best.pt -o yolov5s.wts -t detect/seg

Note

Select `detect` for a detection model and `seg` for a classification model.

3.2 Convert a .wts Model to .engine on the Target Platform

3.2.1 Convert a Detection Model

Run the following on a device platform with the SpireCV SDK installed (from any path):

SpireCVDet -s [PATH-TO-YOUR-WTS] [PATH-TO-YOUR-ENGINE] CLS_NUM s/s6

Where:

• [PATH-TO-YOUR-WTS]: Path to the .wts file generated in the preceding step.
• [PATH-TO-YOUR-ENGINE]: Path to the .engine file to be generated.
• CLS_NUM: Number of target classes.
• s or s6: Different network models.

Note:

• s is a network with an input resolution of 640. Name the generated engine file using the format DatasetName.engine.
• s6 is a network with an input resolution of 1280. Name the generated engine file using the format DatasetName_HD.engine.

Example:

SpireCVDet -s yolov5s.wts COCO.engine 80 s
SpireCVDet -s yolov5s6.wts COCO_HD.engine 80 s6

3.2.2 Convert a Segmentation Model

Run the following on a device platform with the SpireCV SDK installed (from any path):

SpireCVSeg -s [PATH-TO-YOUR-WTS] [PATH-TO-YOUR-ENGINE] CLS_NUM s

Where:

• [PATH-TO-YOUR-WTS]: Path to the .wts file generated in the preceding step.
• [PATH-TO-YOUR-ENGINE]: Path to the .engine file to be generated (the name must end in _SEG).
• CLS_NUM: Number of target classes.
• s: Different network models (segmentation networks currently support only an input resolution of 640).

Example:

SpireCVSeg -s yolov5s-seg.wts COCO_SEG.engine 80 s

3.3 Deploy

Following the SpireCV SDK format, the process above produces the final DatasetName.engine, DatasetName_HD.engine, and DatasetName_SEG.engine weights files. Rename the weights files and place them in the ~/SpireCV/models folder.

For a custom dataset, modify the ~/SpireCV/sv_algorithm_params.json file. For example, suppose the custom dataset is named DatasetName; this name must match the name of the weights file. The dataset contains two target classes, person and car. The following example shows the required parameter changes:

{
    "CommonObjectDetector": {
        "dataset": "DatasetName",  // The specified Dataset selects the corresponding detection model (for example, PersonVehicle, Drone, COCO, or AnotherDatasetName).
        "inputSize": 640,
        "nmsThrs": 0.6,
        "scoreThrs": 0.4,
        "useWidthOrHeight": 1,
        "withSegmentation": false,
        "dataset**DatasetName**": {
            "person": [-1, -1],
            "car": [-1, -1]
        },
        "dataset**AnotherDatasetName**": {
            "another_category": [-1, -1]
        },
        "dataset**PersonVehicle**": {
            "person": [0.5, 1.8],
            "car": [4.1, 1.5],
            "bus": [10, 3],
            "truck": [-1, -1],
            "bike": [-1, -1],
            "train": [-1, -1],
            "boat": [-1, -1],
            "aeroplane": [-1, -1]
        },
        "dataset**Drone**": {
            "drone": [0.4, 0.2]
        },
        "dataset**COCO**": {
            "person": [-1, -1],
            "bicycle": [-1, -1],
            ...
        }
    }
}

Note

When inputSize=640 and withSegmentation=false, DatasetName.engine is run. When inputSize=1280 and withSegmentation=false, DatasetName_HD.engine is run. When inputSize=640 and withSegmentation=true, DatasetName_SEG.engine is run.