ultralytics-yolov8

yolov8

安装环境

装pytorch

conda install pytorch==1.12.1 torchvision==0.13.1 torchaudio==0.12.1 cudatoolkit=11.6 -c pytorch -c conda-forge

安装用到的包
requirements.txt

ultralytics
tqdm
scikit-learn
pyaml
shapely
pyclipper
scikit-image
imgaug
lmdb
tqdm
rapidfuzz
cython
Pillow
pyyaml
requests

pip install -r requirements.txt

数据标注

标注网站

得到via_region_data.csv，格式

{
  "42.243.238.32-202406181512034.jpg418013": {
    "fileref": "",
    "size": 418013,
    "filename": "42.243.238.32-202406181512034.jpg",
    "base64_img_data": "",
    "file_attributes": {},
    "regions": 
      [{
        "shape_attributes": {
          "name": "rect",
          "x": 324,
          "y": 179,
          "width": 279,
          "height": 108
        },
        "region_attributes": {}
      },
      {
        "shape_attributes": {
          "name": "rect",
          "x": 766,
          "y": 248,
          "width": 227,
          "height": 154
        },
        "region_attributes": {}
      }]
  },
}

其中(x,y)是左上角坐标， (x + width, y + height)是右下角坐标

经过统计：
平均：12142158
最小：7681360
最大：1440*2560

总 890, 408 background, 482 front
train 712, 325 background
val 89, 39 background
test 89, 41+？ background

转换格式

转成yolov8的格式, 格式为

label_id x y w h

其中

$$x,y,w,h\in \left[0,1\right]$$
转换

#!/usr/bin/env python
# -*- encoding: utf-8 -*-
import json
import os

import cv2
from tqdm import tqdm

img_dir = r"E:\ccc\origin_images"
json_path = r"E:\ccc\via_region_data.json"
label_dir = r"E:\ccc\origin_labels"

if __name__ == "__main__":
    with open(json_path, "r") as f:
        label_data: dict = json.load(f)

    os.makedirs(label_dir, exist_ok=True)

    cnt = 0
    tot = 0

    for label in tqdm(label_data.values()):
        filename = label["filename"]
        img = cv2.imread(os.path.join(img_dir, filename))
        h, w, _ = img.shape
        category = 0

        empty_region = True

        with open(
            os.path.join(label_dir, os.path.splitext(filename)[0] + ".txt"), "w"
        ) as f:
            for region in label["regions"]:
                empty_region = False
                x = region["shape_attributes"]["x"]
                y = region["shape_attributes"]["y"]
                width = region["shape_attributes"]["width"]
                height = region["shape_attributes"]["height"]

                x_center = (x + (x + width)) / 2.0
                y_center = (y + (y + height)) / 2.0

                x_center_norm = x_center / w
                y_center_norm = y_center / h
                width_norm = width / w
                height_norm = height / h

                f.write(
                    "{} {:.6f} {:.6f} {:.6f} {:.6f}\n".format(
                        category, x_center_norm, y_center_norm, width_norm, height_norm
                    )
                )
        if empty_region:
            cnt += 1
        tot += 1

    print("non_empty: {}; empty: {}; total: {}".format(tot - cnt, cnt, tot))

划分数据集

#!/usr/bin/env python
# -*- encoding: utf-8 -*-
import os
import shutil
from glob import glob

import yaml
from sklearn.model_selection import train_test_split
from tqdm import tqdm

img_dir = r"E:\ccc\origin_images"
label_dir = r"E:\ccc\origin_labels"
train_ratio = 0.8
val_ratio = 0.1
test_ratio = 0.1
root_path = r"E:\ccc"
img_train_path = r"images\train"
img_val_path = r"images\val"
img_test_path = r"images\test"
label_train_path = r"labels\train"
label_val_path = r"labels\val"
label_test_path = r"labels\test"
yaml_path = "config.yaml"


def copy_files(label_paths, img_dir, target_label_dir, origin_img_dir):
    os.makedirs(img_dir, exist_ok=True)
    os.makedirs(target_label_dir, exist_ok=True)
    for label in tqdm(label_paths):
        img = os.path.join(
            origin_img_dir, os.path.splitext(os.path.basename(label))[0] + ".jpg"
        )
        shutil.copy(img, img_dir)
        shutil.copy(label, target_label_dir)


if __name__ == "__main__":
    # img_paths = list(glob(os.path.join(img_dir, '*')))
    label_paths = list(glob(os.path.join(label_dir, "*")))
    # img_paths.sort()
    label_paths.sort()

    is_empty = []
    for path in label_paths:
        with open(path, "r") as f:
            data = f.read().strip()
            if not data:
                is_empty.append(True)
            else:
                is_empty.append(False)
    train_val_label, test_label, train_val_is_empty, test_is_empty = train_test_split(
        label_paths, is_empty, test_size=test_ratio, shuffle=True
    )
    train_label, val_label, train_is_empty, val_is_empty = train_test_split(
        train_val_label,
        train_val_is_empty,
        test_size=val_ratio / (train_ratio + val_ratio),
        shuffle=True,
    )

    train_img = [os.path.splitext(path)[0] + ".jpg" for path in train_label]
    val_img = [os.path.splitext(path)[0] + ".jpg" for path in val_label]
    test_img = [os.path.splitext(path)[0] + ".jpg" for path in test_label]

    copy_files(
        train_label,
        os.path.join(root_path, img_train_path),
        os.path.join(root_path, label_train_path),
        img_dir,
    )
    copy_files(
        val_label,
        os.path.join(root_path, img_val_path),
        os.path.join(root_path, label_val_path),
        img_dir,
    )
    copy_files(
        test_label,
        os.path.join(root_path, img_test_path),
        os.path.join(root_path, label_test_path),
        img_dir,
    )

    config = {
        "path": root_path,
        "train": img_train_path,
        "val": img_val_path,
        "test": img_test_path,
        "names": {0: "box"},
    }

    yaml_dir = os.path.dirname(yaml_path)
    if yaml_dir:
        os.makedirs(yaml_dir, exist_ok=True)

    with open(yaml_path, "w", encoding="utf-8") as f:
        yaml.dump(config, f)

会自动划分数据集和生成需要的yaml
参考：coco128.yaml

训练yolo

#!/usr/bin/env python
# -*- encoding: utf-8 -*-
from ultralytics import YOLO


if __name__ == '__main__':
    # Load a model
    # model = YOLO("yolov8n.yaml")  # build a new model from YAML
    model = YOLO("yolov8n.pt")  # load a pretrained model (recommended for training)
    # model = YOLO("yolov8n.yaml").load("yolov8n.pt")  # build from YAML and transfer weights

    # Train the model
    # results = model.train(data="config.yaml", epochs=100, imgsz=640, device="cpu", single_cls=True, save=True, workers=0)
    results = model.train(data="config.yaml", epochs=100, imgsz=640, device=[0], single_cls=True, save=True, workers=8)

测试yolo

#!/usr/bin/env python
# -*- encoding: utf-8 -*-
from ultralytics import YOLO


if __name__ == "__main__":
    model = YOLO("runs/detect/train/weights/best.pt")  # load a custom model
    # Validate the model
    metrics = model.val(
        data="config.yaml", split="val", plots=True
    )  # no arguments needed, dataset and settings remembered
    # metrics = model.val(
    #     data="config.yaml", split="test", plots=True
    # )  # no arguments needed, dataset and settings remembered
    # print(metrics)
    print("map50-95: ", metrics.box.map)  # map50-95
    print("map50: ", metrics.box.map50)  # map50
    print("map75: ", metrics.box.map75)  # map75
    # print('map50-95 of each category: ',metrics.box.maps)  # a list contains map50-95 of each category

预测yolo

#!/usr/bin/env python
# -*- encoding: utf-8 -*-
import os
from glob import glob

import cv2
from ultralytics import YOLO

model_path = "runs/detect/train/weights/best.pt"
imgs_dir = "/mnt/data/ccc/Eudemons_final"
gt_dir = "/mnt/data/ccc/Eudemons_v3/labels/test"
crop_dir = "/mnt/data/ccc/crop"
save_dir = "/mnt/data/ccc/overlay"
draw_predict = True
# draw_gt = True
draw_gt = False
pad = 0.035


def get_name_v1(filename, dir_name, idx):
    return os.path.join(
        dir_name, "{}{}.jpg".format(filename, f"_{idx}" if idx > 0 else "")
    )


def get_name_v2(path, dir_name, idx):
    filename, ext = os.path.splitext(os.path.basename(path))
    mac = os.path.basename(os.path.dirname(path))
    return os.path.join(
        dir_name, "{}_{}{}{}".format(mac, filename, f"_{idx}" if idx > 0 else "", ext)
    )


if __name__ == "__main__":
    model = YOLO(model_path)  # load a custom model
    paths = list(glob(os.path.join(imgs_dir, "*", "*.jpg")))
    draw_predict = draw_predict and save_dir
    draw_gt = draw_gt and gt_dir and os.path.exists(gt_dir) and save_dir

    if draw_predict:
        os.makedirs(save_dir, exist_ok=True)
    os.makedirs(crop_dir, exist_ok=True)

    # results = model(os.path.join(imgs_dir, "*.jpg"), stream=True)
    results = model(os.path.join(imgs_dir, "*", "*.jpg"), stream=True)

    for result in results:
        # result.save(os.path.basename(result.path))
        path = result.path
        filename = os.path.splitext(os.path.basename(path))[0]
        h, w = result.orig_shape
        orig_img = result.orig_img
        boxes = result.boxes
        boxes_num = boxes.shape[0]
        confidence = boxes.conf.tolist()
        xyxy = boxes.xyxy.tolist()

        for idx, (x1, y1, x2, y2) in enumerate(xyxy):
            x_pad = (x2 - x1) * pad
            y_pad = (y2 - y1) * pad
            x1 -= x_pad
            y1 -= y_pad
            x2 += x_pad
            y2 += y_pad

            x1 = max(0, int(x1))
            y1 = max(0, int(y1))
            x2 = min(w, int(x2))
            y2 = min(h, int(y2))

            cv2.imwrite(
                # get_name_v1(filename, crop_dir, idx),
                get_name_v2(path, crop_dir, idx),
                orig_img[y1:y2, x1:x2],
            )

        if draw_predict:
            for x1, y1, x2, y2 in xyxy:
                x1 = max(0, int(x1))
                y1 = max(0, int(y1))
                x2 = min(w - 1, int(x2))
                y2 = min(h - 1, int(y2))

                orig_img = cv2.rectangle(
                    orig_img,
                    (x1, y1),
                    (x2, y2),
                    color=(0, 0, 255),
                    lineType=cv2.LINE_AA,
                )

            if draw_gt:
                with open(os.path.join(gt_dir, filename + ".txt"), "r") as f:
                    for line in f:
                        line = line.strip()
                        if not line:
                            continue
                        _, x, y, width, height = line.split(" ")
                        x = float(x)
                        y = float(y)
                        width = float(width)
                        height = float(height)
                        orig_img = cv2.rectangle(
                            orig_img,
                            (int((x - width / 2) * w), int((y - height / 2) * h)),
                            (int((x + width / 2) * w), int((y + height / 2) * h)),
                            color=(0, 255, 0),
                            lineType=cv2.LINE_AA,
                        )
            cv2.imwrite(
                # os.path.join(save_dir, filename + ".jpg"),
                get_name_v2(path, save_dir, 0),
                orig_img,
            )