[Object Detection Algorithm] Reproducing YOLOv8

1. Environment Setup

Create a virtual environment.

conda create -n yolov8 python=3.7
conda activate yolov8

Install PyTorch1.8.0.

conda install pytorch==1.8.0 torchvision==0.9.0 torchaudio==0.8.0 cudatoolkit=11.1 -c pytorch -c conda-forge

Download the program open-sourced by the authors and install the other dependencies.

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

2. Dataset Preparation (Instance Segmentation)

(0) Use a camera to record target object data

This example uses Realsense 相机 + Ubuntu 20.04 + ROS Noetic. After connecting the camera to the computer via USB, run the following Python program to capture the target object. Press s to save an image, and press q or esc to stop recording.

The images in the dataset should be as diverse as possible in terms of the number of objects, angles, backgrounds, and other factors.

(1) Use Labelme to create an instance segmentation dataset

Install labelme.

pip install labelme

After installation, enter labelme directly on the command line to open it.

Use labels to annotate the images. Place the generated json files and the original jpg images in the same folder.

(2) Convert Labelme format to YOLO format

Refer to the labelme2yolo package on PyPI.

pip install labelme2yolo
# 或者使用清华源
pip install -i https://pypi.tuna.tsinghua.edu.cn/simple labelme2yolo

labelme2coco --json_dir path/to/labelme/dir

(3) Create a YAML file for the dataset

Open the ultralytics/datasets directory, make a copy of coco128-seg.yaml, rename it to custom-seg.yaml, and modify it for your dataset.

For example:

# 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: /media/mahaofei/OneTouch/Dataset/Program_data/image_processing/ultralytics/20230223_Phone_4Obj_YOLO  # dataset root dir
train: images/train2017  # train images (relative to 'path') 128 images
val: images/train2017  # val images (relative to 'path') 128 images
test:  # test images (optional)

# Classes
names:
  0: ammeter
  1: coffeebox
  2: realsensebox
  3: sucker

3. Python Tutorial

This method is generally used for training and testing.

(1) Training

from ultralytics import YOLO

model = YOLO('yolov8n.pt') # 从预训练模型开始
model.train(epochs=5)

(2) Evaluation

from ultralytics import YOLO

model = YOLO("model.pt")
model.val()  # 使用model.pt的data yaml进行评价
model.val(data='coco128.yaml')  # 或指定数据进行评价

(3) Prediction

Obtain the prediction results.

from ultralytics import YOLO
from PIL import Image
import cv2

model = YOLO("model.pt")
# 接受所有类型 - image/dir/Path/URL/video/PIL/ndarray. 0 for webcam
# 从摄像头
results = model.predict(source="0")
# 从文件夹
results = model.predict(source="folder", show=True) # Display preds. Accepts all YOLO predict arguments

# 从PIL图像
im1 = Image.open("bus.jpg")
results = model.predict(source=im1, save=True)  # save plotted images

# 从ndarray
im2 = cv2.imread("bus.jpg")
results = model.predict(source=im2, save=True, save_txt=True)  # save predictions as labels

# 从PIL/ndarray的列表
results = model.predict(source=[im1, im2])

Analyze the prediction results (results contains a list of all prediction results. When processing many images, take care to avoid running out of memory, especially during instance segmentation).

# 1. return as a list
results = model.predict(source="folder")

# 2.  return as a generator (stream=True)
results = model.predict(source=0, stream=True)

for result in results:
    # Detection
    result.boxes.xyxy   # box with xyxy format, (N, 4)
    result.boxes.xywh   # box with xywh format, (N, 4)
    result.boxes.xyxyn  # box with xyxy format but normalized, (N, 4)
    result.boxes.xywhn  # box with xywh format but normalized, (N, 4)
    result.boxes.conf   # confidence score, (N, 1)
    result.boxes.cls    # cls, (N, 1)

    # Segmentation
    result.masks.data      # masks, (N, H, W)
    result.masks.xy        # x,y segments (pixels), List[segment] * N
    result.masks.xyn       # x,y segments (normalized), List[segment] * N

    # Classification
    result.probs     # cls prob, (num_class, )

# Each result is composed of torch.Tensor by default, 
# in which you can easily use following functionality:
result = result.cuda()
result = result.cpu()
result = result.to("cpu")
result = result.numpy()

4. Start Training

Create a new Python file such as train.py and add the following content:

from ultralytics import YOLO

# Load a model
# model = YOLO('yolov8n-seg.yaml')  # build a new model from YAML
model = YOLO('yolov8n-seg.pt')  # load a pretrained model (recommended for training)
# model = YOLO('yolov8n-seg.yaml').load('yolov8n.pt')  # build from YAML and transfer weights

# Train the model
model.train(data='custom-seg.yaml', epochs=100, imgsz=3904, batch=1)

2.5 Result Prediction