计算机视觉自己训练数据集并进行测试

1.官网上下载yolo5

https://github.com/ultralytics/yolov5

2.制作数据集

2.1数据集进行标注

①下载lablme并打开labelm

②进行数据标注，并进行存放两个文件夹

③选择矩形进行框选，或者可以在Edit里面选择其他图像，另外可以在File里面选择文件自动保存

④将jason数据转为txt的yolo格式

import json
import os

name2id = {'mask': 0, 'nomask': 1}  # 标签名称


def convert(img_size, box):
    dw = 1.0 / img_size[0]
    dh = 1.0 / img_size[1]
    x = (box[0] + box[2]) / 2.0
    y = (box[1] + box[3]) / 2.0
    w = box[2] - box[0]
    h = box[3] - box[1]
    x *= dw
    w *= dw
    y *= dh
    h *= dh
    return x, y, w, h


def decode_json(json_folder_path, json_name):
    txt_name = os.path.join('D:/jetbrains/PycharmProjects/Computer_Vision/experiment/12109990929_ex5/jason/Annotations',
                            json_name.replace('.json', '.txt'))
    os.makedirs(os.path.dirname(txt_name), exist_ok=True)  # 确保目录存在

    try:
        with open(txt_name, 'w') as txt_file, open(os.path.join(json_folder_path, json_name), 'r', encoding='gb2312',
                                                   errors='ignore') as json_file:
            data = json.load(json_file)

            img_w = data['imageWidth']
            img_h = data['imageHeight']

            for shape in data['shapes']:
                label_name = shape['label']
                if shape['shape_type'] == 'rectangle':
                    x1, y1 = map(int, shape['points'][0])
                    x2, y2 = map(int, shape['points'][1])
                    bbox = convert((img_w, img_h), (x1, y1, x2, y2))
                    txt_file.write(f"{name2id[label_name]} {' '.join(map(str, bbox))}\n")

        # 成功转换后删除原始 JSON 文件
        os.remove(os.path.join(json_folder_path, json_name))

    except KeyError as e:
        print(f"键错误：{e} 在文件 {json_name} 中")
    except json.JSONDecodeError as e:
        print(f"JSON 解码错误：{e} 在文件 {json_name} 中")
    except Exception as e:
        print(f"意外错误：{e} 在文件 {json_name} 中")


if __name__ == "__main__":
    json_folder_path = 'D:/jetbrains/PycharmProjects/Computer_Vision/experiment/12109990929_ex5/jason/Annotations'
    json_names = [file for file in os.listdir(json_folder_path) if file.endswith('.json')]

    for json_name in json_names:
        decode_json(json_folder_path, json_name)

2.2数据集进行划分

yolo数据集划分(7:2:1)划分

import os, shutil, random
from tqdm import tqdm

"""
标注文件是yolo格式（txt文件）
训练集：验证集：测试集 （7：2：1） 
"""


def split_img(img_path, label_path, split_list):
    try:
        Data = './VOCdevkit/VOC2007/ImageSets'
        # Data是你要将要创建的文件夹路径（路径一定是相对于你当前的这个脚本而言的）
        # os.mkdir(Data)

        train_img_dir = Data + '/images/train'
        val_img_dir = Data + '/images/val'
        test_img_dir = Data + '/images/test'

        train_label_dir = Data + '/labels/train'
        val_label_dir = Data + '/labels/val'
        test_label_dir = Data + '/labels/test'

        # 创建文件夹
        os.makedirs(train_img_dir)
        os.makedirs(train_label_dir)
        os.makedirs(val_img_dir)
        os.makedirs(val_label_dir)
        os.makedirs(test_img_dir)
        os.makedirs(test_label_dir)

    except:
        print('文件目录已存在')

    train, val, test = split_list
    all_img = os.listdir(img_path)
    all_img_path = [os.path.join(img_path, img) for img in all_img]
    # all_label = os.listdir(label_path)
    # all_label_path = [os.path.join(label_path, label) for label in all_label]
    train_img = random.sample(all_img_path, int(train * len(all_img_path)))
    train_img_copy = [os.path.join(train_img_dir, img.split('\\')[-1]) for img in train_img]
    train_label = [toLabelPath(img, label_path) for img in train_img]
    train_label_copy = [os.path.join(train_label_dir, label.split('\\')[-1]) for label in train_label]
    for i in tqdm(range(len(train_img)), desc='train ', ncols=80, unit='img'):
        _copy(train_img[i], train_img_dir)
        _copy(train_label[i], train_label_dir)
        all_img_path.remove(train_img[i])
    val_img = random.sample(all_img_path, int(val / (val + test) * len(all_img_path)))
    val_label = [toLabelPath(img, label_path) for img in val_img]
    for i in tqdm(range(len(val_img)), desc='val ', ncols=80, unit='img'):
        _copy(val_img[i], val_img_dir)
        _copy(val_label[i], val_label_dir)
        all_img_path.remove(val_img[i])
    test_img = all_img_path
    test_label = [toLabelPath(img, label_path) for img in test_img]
    for i in tqdm(range(len(test_img)), desc='test ', ncols=80, unit='img'):
        _copy(test_img[i], test_img_dir)
        _copy(test_label[i], test_label_dir)


def _copy(from_path, to_path):
    shutil.copy(from_path, to_path)


def toLabelPath(img_path, label_path):
    img = img_path.split('\\')[-1]
    label = img.split('.jpg')[0] + '.txt'
    return os.path.join(label_path, label)


if __name__ == '__main__':
    img_path = 'JPEGImage'  # 你的图片存放的路径（路径一定是相对于你当前的这个脚本文件而言的）
    label_path = 'Annotations'  # 你的txt文件存放的路径（路径一定是相对于你当前的这个脚本文件而言的）
    split_list = [0.7, 0.2, 0.1]  # 数据集划分比例[train:val:test]
    split_img(img_path, label_path, split_list)

json转voc

import cv2
import json
import os
import os.path as osp
import shutil
import chardet
 
def get_encoding(path):
    f = open(path, 'rb')
    data = f.read()
    file_encoding = chardet.detect(data).get('encoding')
    f.close()
    return file_encoding
 
def is_pic(img_name):
    valid_suffix = ['JPEG', 'jpeg', 'JPG', 'jpg', 'BMP', 'bmp', 'PNG', 'png']
    suffix = img_name.split('.')[-1]
    if suffix not in valid_suffix:
        return False
    return True
 
 
class X2VOC(object):
    def __init__(self):
        pass
 
    def convert(self, image_dir, json_dir, dataset_save_dir):
        """转换。
        Args:
            image_dir (str): 图像文件存放的路径。
            json_dir (str): 与每张图像对应的json文件的存放路径。
            dataset_save_dir (str): 转换后数据集存放路径。
        """
        assert osp.exists(image_dir), "The image folder does not exist!"
        assert osp.exists(json_dir), "The json folder does not exist!"
        if not osp.exists(dataset_save_dir):
            os.makedirs(dataset_save_dir)
        # Convert the image files.
        new_image_dir = osp.join(dataset_save_dir, "JPEGImages")
        if osp.exists(new_image_dir):
            raise Exception(
                "The directory {} is already exist, please remove the directory first".
                format(new_image_dir))
        os.makedirs(new_image_dir)
        for img_name in os.listdir(image_dir):
            if is_pic(img_name):
                shutil.copyfile(
                    osp.join(image_dir, img_name),
                    osp.join(new_image_dir, img_name))
        # Convert the json files.
        xml_dir = osp.join(dataset_save_dir, "Annotations")
        if osp.exists(xml_dir):
            raise Exception(
                "The directory {} is already exist, please remove the directory first".
                format(xml_dir))
        os.makedirs(xml_dir)
        self.json2xml(new_image_dir, json_dir, xml_dir)
 
 
class LabelMe2VOC(X2VOC):
    """将使用LabelMe标注的数据集转换为VOC数据集。
    """
    def json2xml(self, image_dir, json_dir, xml_dir):
        import xml.dom.minidom as minidom
        i = 0
        for img_name in os.listdir(image_dir):
            img_name_part = osp.splitext(img_name)[0]
            json_file = osp.join(json_dir, img_name_part + ".json")
            i += 1
            if not osp.exists(json_file):
                os.remove(osp.join(image_dir, img_name))
                continue
            xml_doc = minidom.Document()
            root = xml_doc.createElement("annotation")
            xml_doc.appendChild(root)
            node_folder = xml_doc.createElement("folder")
            node_folder.appendChild(xml_doc.createTextNode("JPEGImages"))
            root.appendChild(node_folder)
            node_filename = xml_doc.createElement("filename")
            node_filename.appendChild(xml_doc.createTextNode(img_name))
            root.appendChild(node_filename)
            with open(json_file, mode="r", \
                      encoding=get_encoding(json_file)) as j:
                json_info = json.load(j)
                if 'imageHeight' in json_info and 'imageWidth' in json_info:
                    h = json_info["imageHeight"]
                    w = json_info["imageWidth"]
                else:
                    img_file = osp.join(image_dir, img_name)
                    im_data = cv2.imread(img_file)
                    h, w, c = im_data.shape
                node_size = xml_doc.createElement("size")
                node_width = xml_doc.createElement("width")
                node_width.appendChild(xml_doc.createTextNode(str(w)))
                node_size.appendChild(node_width)
                node_height = xml_doc.createElement("height")
                node_height.appendChild(xml_doc.createTextNode(str(h)))
                node_size.appendChild(node_height)
                node_depth = xml_doc.createElement("depth")
                node_depth.appendChild(xml_doc.createTextNode(str(3)))
                node_size.appendChild(node_depth)
                root.appendChild(node_size)
                for shape in json_info["shapes"]:
                    if 'shape_type' in shape:
                        if shape["shape_type"] != "rectangle":
                            continue
                        (xmin, ymin), (xmax, ymax) = shape["points"]
                        xmin, xmax = sorted([xmin, xmax])
                        ymin, ymax = sorted([ymin, ymax])
                    else:
                        points = shape["points"]
                        points_num = len(points)
                        x = [points[i][0] for i in range(points_num)]
                        y = [points[i][1] for i in range(points_num)]
                        xmin = min(x)
                        xmax = max(x)
                        ymin = min(y)
                        ymax = max(y)
                    label = shape["label"]
                    node_obj = xml_doc.createElement("object")
                    node_name = xml_doc.createElement("name")
                    node_name.appendChild(xml_doc.createTextNode(label))
                    node_obj.appendChild(node_name)
                    node_diff = xml_doc.createElement("difficult")
                    node_diff.appendChild(xml_doc.createTextNode(str(0)))
                    node_obj.appendChild(node_diff)
                    node_box = xml_doc.createElement("bndbox")
                    node_xmin = xml_doc.createElement("xmin")
                    node_xmin.appendChild(xml_doc.createTextNode(str(xmin)))
                    node_box.appendChild(node_xmin)
                    node_ymin = xml_doc.createElement("ymin")
                    node_ymin.appendChild(xml_doc.createTextNode(str(ymin)))
                    node_box.appendChild(node_ymin)
                    node_xmax = xml_doc.createElement("xmax")
                    node_xmax.appendChild(xml_doc.createTextNode(str(xmax)))
                    node_box.appendChild(node_xmax)
                    node_ymax = xml_doc.createElement("ymax")
                    node_ymax.appendChild(xml_doc.createTextNode(str(ymax)))
                    node_box.appendChild(node_ymax)
                    node_obj.appendChild(node_box)
                    root.appendChild(node_obj)
            with open(osp.join(xml_dir, img_name_part + ".xml"), 'w') as fxml:
                xml_doc.writexml(
                    fxml,
                    indent='\t',
                    addindent='\t',
                    newl='\n',
                    encoding="utf-8")
 
 
def convert(pics,anns,save_dir):
    """
    将使用labelme标注的数据转换为VOC格式
    请将labelme标注的文件中，所有img文件保存到pics文件夹中，所有xml文件保存到anns文件夹中，结构如下：
            --labelmedata
            ---pics
            ----img0.jpg
            ----img1.jpg
            ----......
            ---anns
            ----img0.mxl
            ----img1.xml
            ----......
    :param pics: img文件所在文件夹的路径
    :param anns: xml文件所在文件夹的路径
    :param save_dir: 输出VOC格式数据的保存路径
    :return:
    """
    labelme2voc = LabelMe2VOC().convert
    labelme2voc(pics, anns, save_dir)
 
if __name__=="__main__":
    convert(pics=r"JPEGImages",  # 修改图片路径
            anns=r"json",   # 修改json格式文件路径
            save_dir=r"VOC")  # 保存VOC格式的路径

VOC转YOLO格式

import xml.etree.ElementTree as ET
import pickle
import os
from os import listdir, getcwd
from os.path import join
 
 
def convert(size, box):
    x_center = (box[0] + box[1]) / 2.0
    y_center = (box[2] + box[3]) / 2.0
    x = x_center / size[0]
    y = y_center / size[1]
    w = (box[1] - box[0]) / size[0]
    h = (box[3] - box[2]) / size[1]
    return (x, y, w, h)
 
 
def convert_annotation(xml_files_path, save_txt_files_path, classes):
    xml_files = os.listdir(xml_files_path)
    print(xml_files)
    for xml_name in xml_files:
        print(xml_name)
        xml_file = os.path.join(xml_files_path, xml_name)
        out_txt_path = os.path.join(save_txt_files_path, xml_name.split('.')[0] + '.txt')
        out_txt_f = open(out_txt_path, 'w')
        tree = ET.parse(xml_file)
        root = tree.getroot()
        size = root.find('size')
        w = int(size.find('width').text)
        h = int(size.find('height').text)
 
        for obj in root.iter('object'):
            difficult = obj.find('difficult').text
            cls = obj.find('name').text
            if cls not in classes or int(difficult) == 1:
                continue
            cls_id = classes.index(cls)
            xmlbox = obj.find('bndbox')
            b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text),
                 float(xmlbox.find('ymax').text))
            # b=(xmin, xmax, ymin, ymax)
            print(w, h, b)
            bb = convert((w, h), b)
            out_txt_f.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')
 
 
if __name__ == "__main__":
    # 需要转换的类别，需要一一对应
    classes1 = ['boat', 'cat']
    # 2、voc格式的xml标签文件路径
    xml_files1 = r'C:\Users\86159\Desktop\VOC2007\Annotations'
    # 3、转化为yolo格式的txt标签文件存储路径
    save_txt_files1 = r'C:\Users\86159\Desktop\VOC2007\label'
 
    convert_annotation(xml_files1, save_txt_files1, classes1)

最后的文件划分如下

3.进行yolo5测试，在yolo5-master官方包下

3.1修改自己的data.yaml用于数据集（在仿照data/coco.yaml）

train: ../Maskdata/train/images#训练集
val: ../Maskdata/val/images#测试集
test: ../Maskdata/test/images#验证集
nc: 2#类别
names:
  0: mask
  1: nomask

3.2复制models/yolov5s.yaml重命名为my_yolov5s.yaml（修改类别）

3.3目录结构

4.训练yolo5，运行train.py

4.1修改参数，找到def parse_opt(known=False):

训练结果

5.用训练出来的模型进行预测，修改detect.py

5.1找到def parse_opt():修改参数

5.2运行结果