# 推荐
pip install ultralytics

git clone https://github.com/ultralytics/ultralytics
pip install -qe ultralytics

import ultralytics

if __name__ == '__main__':
    ultralytics.checks()

# 输出结果
Ultralytics YOLOv8.0.13  Python-3.8.13 torch-1.11.0 CUDA:0 (NVIDIA GeForce GTX 1070 Ti, 8192MiB)
Setup complete  (16 CPUs, 31.9 GB RAM, 2715.9/3023.4 GB disk)

# 在命令行下运行
yolo task=detect mode=predict model=yolov8n.pt conf=0.25 source='bus.jpg'

# 若提示 export failure 导出失败，需要设置 opset=11
yolo mode=export model=yolov8n.pt format=onnx opset=11

# coco128 的数据集结构如下
└─coco128
   ├─ images
   │      └─ train2017
   └─ labels
          └─ train2017
          
# 需要转换成如下结构
└─datasets
    └─coco128
        ├─test    # 测试
        │  ├─images
        │  └─labels
        ├─train   # 训练
        │  ├─images
        │  └─labels
        └─valid   # 校验
            ├─images
            └─labels

CWD = os.getcwd()


def move_file(src, dest):
    if not os.path.isfile(src):
        print("%s not exist!" % src)
    else:
        fpath, fname = os.path.split(dest)  # 分离文件名和路径
        if not os.path.exists(fpath):
            os.makedirs(fpath)  # 创建路径
        shutil.move(src, dest)  # 移动文件
        # print("move %s -> %s" % (src, dest))


def copy_file(src, dest):
    if not os.path.isfile(src):
        print("%s not exist!" % src)
    else:
        fpath, fname = os.path.split(dest)  # 分离文件名和路径
        if not os.path.exists(fpath):
            os.makedirs(fpath)  # 创建路径
        shutil.copyfile(src, dest)  # 复制文件
        # print("copy %s -> %s" % (src, dest))


def get_valid_data():
    valid_data_list = []
    for image_name in os.listdir(image_dir):
        image_path = os.path.join(image_dir, image_name)
        label_name = os.path.splitext(image_name)[0] + '.txt'
        label_path = os.path.join(label_dir, label_name)
        if os.path.exists(label_path):
            valid_data_list.append((image_path, label_path))
    return valid_data_list


if __name__ == '__main__':
    # 获取当前路径
    CWD = os.getcwd()
    # 指定数据集 images 路径
    image_dir = F'{CWD}/coco128/images/train2017'
    # 指定数据集 labels 路径
    label_dir = F'{CWD}/coco128/labels/train2017'

    # 指定数据集存储路径
    out_dir = F'{CWD}/coco128-datasets'

    # 存入有效标注数据
    valid_data = get_valid_data()

    # 按 train：valid：test = 7:2:1 比例随机选取数据 移动到输出目录
    for i in tqdm(range(len(valid_data))):
        dest_dir = ''
        r = random.random()
        if r <= 0.1:
            # 移动到 test 文件夹
            dest_dir = F"{out_dir}/test"
        elif r > 0.2:
            # 移动到 train 文件夹
            dest_dir = F"{out_dir}/train"
        else:
            # 移动到 valid 文件夹
            dest_dir = F"{out_dir}/valid"

        move_file(valid_data[i][0], os.path.join(dest_dir, 'images', os.path.basename(valid_data[i][0])))
        move_file(valid_data[i][1], os.path.join(dest_dir, 'labels', os.path.basename(valid_data[i][1])))

    print('生成完毕！')

# 若出现 OSError: [WinError 1455] 页面文件太小，无法完成操作。需要 workers=0
yolo task=detect mode=train model=yolov8n.yaml data=coco128.yaml epochs=3 imgsz=640 resume=True workers=0

# 训练完成后 一般选用 best.pt