a65f8cb709
|
||
|---|---|---|
| .. | ||
| Readme.md | ||
Readme.md
CV - Retrain Yolo v8 on custom data
Create environment
conda create -n yolov8_custom python=3.9
conda activate yolov8_custom
pip install simple_image_download==0.4
pip install ultralytics
To use the GPU also copy the pip install command listed here.
To check if CUDA is correctly set up:
import torch
torch.__version__
torch.cuda.is_available()
Download sample images from google
from simple_image_download import simple_image_download as simp
response = simpl.simple_image_download
keywords = ["building workers"]
for kw in keywords:
response().download(kw, 200)
Annotate images
To annotate the image labelImg tool can be used:
pip install labelImg
labelImg
Output folder structure:
.
├── train
│ ├── images
│ │ └── image.png
│ └── labels
│ └── image.txt
└── val
├── images
│ └── image.png
└── labels
└── image.txt
The folder structure must be declared in a .yaml file:
train: C:\<abs_path>\train
val: C:\<abs_path>\val
nc: 2
names: ["hat", "jacket"]
The names must be the same declared in a file called classes.txt.
Annotation format (BBox)
<class_number> <norm_box_center_h> <norm_box_center_v> <norm_box_height> <norm_box_width>
Train the net
yolo task=detect mode=train epochs=100 data=data_custom.yaml model=yolov8m.pt imgsz=600
The actual model and imgsz can be found here.
Out of memory error
In case of CUDA out of memory error a smaller batch size must be selected:
yolo task=detect mode=train epochs=100 data=data_custom.yaml model=yolov8m.pt imgsz=600 batch=4
Training generated files
The retrained net weights are stored in /run/detect/train/weights/best.pt.
Run the retrained network
From command line
yolo task=detect mode=predict model=best.pt show=True conf=0.5 source=image.png
From Python script
from ultralytics import YOLO
model = YOLO("best.pt")
model.predict(source="image.png", show=True, save=True, conf=0.5)
Export YOLO in ONNX format
yolo task=detect mode=export model=best.pt format=onnx