yolov8.1训练自己的数据识别游戏人物的头部。

星楠_001

已于 2024-01-18 22:30:33 修改

阅读量937

点赞数 11

分类专栏： yolov8 文章标签： YOLO 游戏程序 python

于 2024-01-18 22:28:38 首次发布

本文链接：https://blog.csdn.net/agoodtimeo/article/details/135658560

版权

yolov8 专栏收录该内容

2 篇文章 1 订阅

订阅专栏

思路：通过官网：主页 - Ultralytics YOLOv8 文档了解了yolov8有命令行和python运行。通过yaml文件定位你的图片和标签。部署还是很简单的，基本上不会出错。

过程：

使用conda先新建一个python3.8的虚拟环境

conda create -n yolov8 python==3.8.0

然后安装官网上的的cuda+pytorch+ultralytics环境。

conda install -c pytorch -c nvidia -c conda-forge pytorch torchvision pytorch-cuda=11.8 ultralytics

然后环境就部署好了就是这么简单。

然后开始准备自己的文件夹.

在自己的工作文件夹下新建一个yolov8用来装数据

在yolov8新建

my_data.yaml内容：

# Ultralytics YOLO 🚀, AGPL-3.0 license
# COCO128 dataset https://www.kaggle.com/ultralytics/coco128 (first 128 images from COCO train2017) by Ultralytics
# Documentation: https://docs.ultralytics.com/datasets/detect/coco/
# Example usage: yolo train data=coco128.yaml
# parent
# ├── ultralytics
# └── datasets
#     └── coco128  ← downloads here (7 MB)


# 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: D://zyx-gitStudy//class//yolov8//ultralytics-main//datasets//yumi # dataset root dir
train: D://zyx-gitStudy//class//yolov8//ultralytics-main//datasets//yumi//images//train2017  # train images (relative to 'path') 128 images
val: D://zyx-gitStudy//class//yolov8//ultralytics-main//datasets//yumi//images//val2017  # val images (relative to 'path') 128 images
test:  # test images (optional)

# Classes
names:
  0: enemy
  1: head



 

# Download script/URL (optional)
# download: https://ultralytics.com/assets/coco128.zip

准备数据推荐PicPickhttps://picpick.app/zh/ 和 lableme

PicPick设置

保存img文件夹以便后续处理

使用快捷键更快。

PicPick我们使用固定区域截取640*640(你也可以使用别的尺寸)

在文件夹里的图片要重命名

import os

def rename_files(folder_path):
    # 获取文件夹中所有jpg文件的列表
    jpg_files = [file for file in os.listdir(folder_path) if file.endswith('.jpg')]

    # 按照0.1.2.3.4.5.6的顺序重命名文件
    for index, jpg_file in enumerate(jpg_files):
        new_name = f"{index}.jpg"
        old_path = os.path.join(folder_path, jpg_file)
        new_path = os.path.join(folder_path, new_name)
        os.rename(old_path, new_path)
        print(f"Renamed: {jpg_file} to {new_name}")

# 指定文件夹路径
folder_path = r"D:\zyx-gitStudy\class\yolov8\ultralytics-main\datasets\yumi\images\img"

# 调用函数进行重命名
rename_files(folder_path)

lableme环境下载，还是在conda里：

conda install labelme

使用：

labelme

打开图片文件夹然后就可以画框框

注意一定用0，1，2，后续会通过yaml来转换

画好了之后文件夹里就包含着图片和json。

记得备份！！！

由于yolov8读取不了json，这时候要json转txt。代码如下

import json
import os
import glob
from tqdm import tqdm


def convert_poly_to_rect(coordinateList):
    X = [int(coordinateList[2 * i]) for i in range(int(len(coordinateList) / 2))]
    Y = [int(coordinateList[2 * i + 1]) for i in range(int(len(coordinateList) / 2))]

    Xmax = max(X)
    Xmin = min(X)
    Ymax = max(Y)
    Ymin = min(Y)
    flag = False
    if (Xmax - Xmin) == 0 or (Ymax - Ymin) == 0:
        flag = True
    return [Xmin, Ymin, Xmax - Xmin, Ymax - Ymin], flag

def convert_labelme_json_to_txt(json_path,img_path,out_txt_path):
    json_list = glob.glob(json_path + '/*.json')
    num = len(json_list)
    for json_path in tqdm(json_list):
        with open(json_path, "r")as f_json:
            json_data = json.loads(f_json.read())
        infos = json_data['shapes']
        if len(infos) ==0:
            continue
        img_w = json_data['imageWidth']
        img_h = json_data['imageHeight']
        image_name = json_data['imagePath']
        image_path = os.path.join(img_path, image_name)
        if not os.path.exists(img_path):
            print(img_path, 'is None!')
            continue
        txt_name = os.path.basename(json_path).split('.')[0] + '.txt'
        txt_path = os.path.join(out_txt_path, txt_name)
        f = open(txt_path, 'w')
        for label in infos:
            points = label['points']
            if len(points) < 2:
                continue

            if len(points) == 2:
                x1 = points[0][0]
                y1 = points[0][1]
                x2 = points[1][0]
                y2 = points[1][1]
                points = [[x1, y1], [x2, y1], [x2, y2], [x1, y2]]
            else:
                if len(points) < 4:
                    continue

            segmentation = []
            for p in points:
                segmentation.append(int(p[0]))
                segmentation.append(int(p[1]))

            bbox, flag = convert_poly_to_rect(list(segmentation))
            x1, y1, w, h = bbox

            if flag:
                continue

            x_center = x1 + w/2
            y_center = y1 + h/2
            norm_x = x_center / img_w
            norm_y = y_center / img_h
            norm_w = w / img_w
            norm_h = h / img_h
            obj_cls = label['label']
            line = [obj_cls, norm_x, norm_y, norm_w, norm_h]
            line = [str(ll) for ll in line]
            line = ' '.join(line) + '\n'
            f.write(line)
        f.close()

if __name__ == "__main__":
#图片路径和json路径和txt路径，可以是一样的.
    img_path = r''
    json_path = r''
    out_txt_path = r''

    if not os.path.exists(out_txt_path):
        os.makedirs(out_txt_path)
    
    convert_labelme_json_to_txt(json_path,img_path,out_txt_path)

把json放到别的文件夹，或者直接删除

转换为txt后，就可以随机分配训练集和测试集

注意我把图片和txt放到一起了,你们可以不用这样，最终效果就是，txt和img一对的按比例分成训练集和测试集

#随机分配img
import os
import shutil
import random

def move_jpg_files_by_random_ratio(source_folder, destination_folders, ratios):
    # 确保目标文件夹存在，如果不存在则创建
    for folder in destination_folders:
        if not os.path.exists(folder):
            os.makedirs(folder)

    # 获取源文件夹中的所有文件
    files = [file for file in os.listdir(source_folder) if file.lower().endswith('.jpg')]

    for file in files:
        # 构建源文件的完整路径
        source_path = os.path.join(source_folder, file)

        # 根据比例随机选择目标文件夹
        destination_folder = random.choices(destination_folders, weights=ratios)[0]

        # 构建目标文件夹的完整路径
        destination_path = os.path.join(destination_folder, file)

        # 移动文件
        shutil.move(source_path, destination_path)
        print(f"Moved {file} to {destination_path}")


# 示例用法
source_folder = r"D:\zyx-gitStudy\class\yolov8\ultralytics-main\datasets\yumi\images\img"
destination_folders = [r"D:\zyx-gitStudy\class\yolov8\ultralytics-main\datasets\yumi\images\train2017", r"D:\zyx-gitStudy\class\yolov8\ultralytics-main\datasets\yumi\images\val2017"]
ratios = [0.7, 0.3]  # 70%的概率放到第一个文件夹，30%的概率放到第二个文件夹

move_jpg_files_by_random_ratio(source_folder, destination_folders, ratios)

现在根据名字再把txt拉过来

import os
import shutil

def move_files_by_name(src_folder, dest_folder, third_folder):
    # 获取原始文件夹中的所有文件
    src_files = [f for f in os.listdir(src_folder) if os.path.isfile(os.path.join(src_folder, f))]

    # 获取目标文件夹中的所有文件
    dest_files = [f for f in os.listdir(dest_folder) if os.path.isfile(os.path.join(dest_folder, f))]

    # 获取第三个文件夹中的所有文件
    third_files = [f for f in os.listdir(third_folder) if os.path.isfile(os.path.join(third_folder, f))]

    for src_file in src_files:

        # 获取原始文件的名称和后缀
        src_name, src_ext = os.path.splitext(src_file)

        # 在目标文件夹中搜索相同名称的文件
        matching_files = [f for f in dest_files if f.startswith(src_name) and f != src_file]
        
        #如果找到匹配的文件，移动到第三个文件夹
        if matching_files:
            dest_file = matching_files[0]
            dest_path = os.path.join(dest_folder, dest_file)
            third_path = os.path.join(third_folder, f"{src_name}.txt")
            shutil.move(dest_path, third_path)
            print(f"Moved {dest_file} to {third_path} in the third folder.")

# 训练集
source_folder = r"D:\zyx-gitStudy\class\yolov8\ultralytics-main\datasets\yumi\images\train2017"
destination_folder = r"D:\zyx-gitStudy\class\yolov8\ultralytics-main\datasets\yumi\images\img"
third_folder = r"D:\zyx-gitStudy\class\yolov8\ultralytics-main\datasets\yumi\images\train2017"
# 调用函数移动文件
move_files_by_name(source_folder, destination_folder, third_folder)
print()
# 测试集
source_folder = r"D:\zyx-gitStudy\class\yolov8\ultralytics-main\datasets\yumi\images\val2017"
destination_folder = r"D:\zyx-gitStudy\class\yolov8\ultralytics-main\datasets\yumi\images\img"
third_folder = r"D:\zyx-gitStudy\class\yolov8\ultralytics-main\datasets\yumi\images\val2017"
# 调用函数移动文件
move_files_by_name(source_folder, destination_folder, third_folder)

ok了，下一步开始训练

if __name__ == '__main__':
    # Your training code here
    from ultralytics import YOLO


    # 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='D:\\zyx-gitStudy\\class\\yolov8\\ultralytics-main\\my_data.yaml', epochs=100, imgsz=640)

注意用conda环境下的python

运行，注意你的my_data.yaml的路径和yaml里面的路径是否正确.

开始训练了，训练的结果在根目录下的runs

best.pt就是最好的模型，last.pt是最后一次训练时的模型。

测试:

from pathlib import Path
from ultralytics import YOLO  # 确保正确导入模型类
import cv2
# 加载模型
weights_path = 'D:\\zyx-gitStudy\\class\\yolov8\\runs\detect\\train13\\weights\\best.pt'  # 替换为您的权重文件路径
model = YOLO(weights_path)
# from ndarray
im2 = cv2.imread(r"D:\zyx-gitStudy\class\yolov8\ultralytics-main\datasets\yumi\images\img\a.jpg")
results = model.predict(source=im2, save=True, save_txt=True)  # save predictions as labels