# CV - Retrain Yolo v8 on custom data
## Create environment
```console
conda create -n yolov8_custom python=3.9
```
```console
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](https://pytorch.org/).
To check if CUDA is correctly set up:
```python
import torch
torch.__version__
torch.cuda.is_available()
```
## Download sample images from google
```python
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](https://pypi.org/project/labelImg/) can be used:
```console
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:
```yaml
train: C:\\train
val: C:\\val
nc: 2
names: ["hat", "jacket"]
```
The names must be the same declared in a file called *classes.txt*.
### Annotation format (BBox)
```
```
## Train the net
```console
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](https://github.com/ultralytics/ultralytics).
### Out of memory error
In case of CUDA out of memory error a smaller batch size must be selected:
```console
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
```console
yolo task=detect mode=predict model=best.pt show=True conf=0.5 source=image.png
```
## From Python script
```python
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
```console
yolo task=detect mode=export model=best.pt format=onnx
```
## References
- [Retrain yolo v8 for classification](https://www.youtube.com/watch?v=gRAyOPjQ9_s)
- [Yolo v8 for segmentation](https://www.youtube.com/watch?v=75LI9MI9eEo)
- [Segment objects with Yolov8](https://medium.com/@Mert.A/how-to-segment-with-yolov8-f33b1c63b6c6)