YOLO万能分割SAM自动标注，爬取图片，转换yolo格式，YOLOV8自定义数据集，训练模型及测试，智能垃圾分类

德炜_111

已于 2024-08-16 09:23:00 修改

阅读量859

点赞数 29

文章标签： YOLO 深度学习人工智能 python

于 2024-08-14 14:55:01 首次发布

本文链接：https://blog.csdn.net/weixin_42121606/article/details/141160217

版权

在学习如何使用YOLO进行目标检测学习时，我意识到虽然可以轻易在网上找到预构建的数据集，但我想要创建自己的数据集来训练模型。然而，我发现手动为图片添加标注非常耗时。因此，我在寻找是否有自动化的方法来标注图片，以简化这一过程。

简易流程图

代码地址：garbage_classify: 垃圾分类

1.爬取图片

首先需要解决的是图片的问题。理想情况下，训练用的图片应该由行业专家提供，但作为初学者，往往只能从网上获取。YOLO的训练需要大量图片，虽然可以通过手动下载并保存图片，但这种方式非常耗时。

使用代码进行图片爬取可以显著提高效率。在这部分中，我们使用谷歌浏览器的 ChromeDriver 进行自动化操作。安装好相关的 Python 包后，运行文件 `01下载图片.py` 即可根据指定的图片类别下载图片。代码会根据类别名称创建相应的文件夹，并将对应的图片保存在其中。

进一步优化：可以将搜索网站提取出来，做成一个配置文件，这样就能从多个网站进行图片查询和下载。这不仅能提高图片收集的效率，还能增加数据来源的多样性。

def fetch_image_urls(driver, search_query, num_images):
    """
    根据搜索查询获取指定数量的图片URL。
    该函数使用提供的WebDriver对象来执行网页搜索，并提取图片链接。
    它会滚动浏览页面，以便加载更多图片，然后从页面源代码中提取图片URL
    参数:
    - driver: WebDriver对象，用于执行网页操作。
    - search_query: str类型，要搜索的图片关键词。
    - num_images: int类型，需要获取的图片数量。
    """
    logging.info(f'正在打开百度图片搜索页面：https://image.baidu.com/search/index?tn=baiduimage&word={+search_query}')
    driver.get(f'https://image.baidu.com/search/index?tn=baiduimage&word={+search_query}')
    logging.info('等待页面加载完成...')
    time.sleep(2)  # 增加等待时间，确保页面完全加载

    logging.info('模拟滚动浏览器窗口，加载更多图片...')
    body = driver.find_element(By.TAG_NAME, 'body')
    scroll_count = 0
    num_images_max = num_images * 1.1
    image_urls = set()  # 使用 set 以避免重复

    while len(image_urls) < num_images_max:
        body.send_keys(Keys.PAGE_DOWN)
        time.sleep(1)
        scroll_count += 1
        logging.info(f'已滚动 {scroll_count} 次，当前已加载 {len(image_urls)} 张图片')
        page_source = driver.page_source
        soup = BeautifulSoup(page_source, 'html.parser')
        img_tags = soup.find_all('img', {'class': 'main_img img-hover'})
        for img_tag in img_tags:
            image_url = img_tag.get('src') or img_tag.get('data-src')
            if image_url and image_url.startswith('//'):
                image_url = 'https:' + image_url
            if image_url:
                image_urls.add(image_url)
            if len(image_urls) >= num_images_max:
                break
        if len(image_urls) >= num_images_max:
            break

    return list(image_urls)[10:]  # 返回列表，并跳过前 10 个


def download_images(driver, search_query, num_images, output_folder, image_name_field):
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    session = setup_requests_session()
    image_urls = fetch_image_urls(driver, search_query, num_images)
    logging.info(f'开始下载图片，共需下载 {num_images} 张...')

    # 使用 ThreadPoolExecutor 来并行下载图片
    success_count = 0
    with ThreadPoolExecutor(max_workers=10) as executor:# 设置10个线程下载图片
        futures = [executor.submit(download_image, url, output_folder, image_name_field, i, session) for i, url in
                   enumerate(image_urls[:num_images])]
        for i, future in enumerate(as_completed(futures)):
            result = future.result()
            if 'Success' in result:
                success_count += 1
            if (i + 1) % 10 == 0 or i == len(image_urls[:num_images]) - 1:
                logging.info(f'{i + 1}/{num_images} 下载完成: {result}')

    logging.info('所有图片下载完成。')
    return output_folder

在下载图片之前，我们需要创建一个名为 `classify.txt` 的文件，用于定义图片类别。这个文件将作为后续逻辑执行的基础。`classify.txt` 文件中每一行代表一个类别，其中第一列是中文名称，第二列是对应的英文名称。

代码会根据每个类别名称在百度图片中进行搜索。我们可以在代码中设置每个类别所要下载的图片数量。为了演示这个流程，我们以垃圾分类为背景创建一个 demo。

手机 Mobilephones
塑料瓶 Plastic
纸箱 Cardboard
易拉罐 Zip-topcan
电灯泡 LightBulb
电池 Batteries
香蕉皮 BananaSkin
玉米核 CornKernel

下载的图片会保存在 `data` 文件夹中。然而，下载的图片可能与我们期望的不完全一致，对于这些不符合要求的图片，我们可以直接删除。因为如果将这些不符合类别的图片用于模型训练，可能会引入噪音，影响模型的性能。

2.万能分割

解决了图片问题后，接下来就是图像标注。传统的标注通常需要人工手动完成，但通过学习发现，在 YOLO 中有一个万能分割模型，可以使用 ultralytics 库中的 `auto_annotate` 函数来自动为图像数据进行标注。

SAM (Segment Anything Model) - Ultralytics YOLO Docs

万能分割代码

import os
from ultralytics.data.annotator import auto_annotate # type: ignore

def main():
    classify_file = 'classify.txt'
    base_folder = '../data/01下载图片'
    det_model = "yolov8m.pt"
    sam_model = 'sam_b.pt'

    # 读取 classify.txt 文件并处理每一行
    with open(classify_file, 'r', encoding='utf-8') as file:
        for line in file:
            parts = line.strip().split()
            if len(parts) < 1:
                continue

            name = parts[0]
            data_folder = os.path.join(base_folder, name)

            if not os.path.exists(data_folder):
                print(f"目录不存在: {data_folder}")
                continue

            print(f'处理数据目录: {data_folder}')
            try:
                auto_annotate(data=data_folder, det_model=det_model, sam_model="sam_b.pt")
                print(f'成功处理目录: {data_folder}')
            except Exception as e:
                print(f'处理目录 {data_folder} 时发生错误: {e}')

if __name__ == "__main__":
    main()

在使用 YOLO 模型进行分割任务时，通常是指利用在 COCO 数据集上训练的模型。COCO 数据集包含 80 种常见物体，这些模型经过大规模数据训练，能够在图像中快速识别和定位这些类别的物体，并生成相应的分割边界或区域。

执行02文件，为每张照片生成相应的分割边界或区域，并且生成标注文件

image 43/45 D:\PycharmProjects\garbage_classify\src\..\data\01下载图片\塑料瓶\塑料瓶_7.jpg: 640x384 2 bottles, 19.3ms
image 44/45 D:\PycharmProjects\garbage_classify\src\..\data\01下载图片\塑料瓶\塑料瓶_8.jpg: 640x640 1 bottle, 27.5ms
image 45/45 D:\PycharmProjects\garbage_classify\src\..\data\01下载图片\塑料瓶\塑料瓶_9.jpg: 640x640 4 bottles, 1 vase, 28.1ms
Speed: 2.1ms preprocess, 27.6ms inference, 1.5ms postprocess per image at shape (1, 3, 640, 640)
成功处理目录: ../data/01下载图片\塑料瓶
处理数据目录: ../data/01下载图片\纸箱

image 1/43 D:\PycharmProjects\garbage_classify\src\..\data\01下载图片\纸箱\纸箱_1.jpg: 608x640 1 refrigerator, 50.4ms
image 2/43 D:\PycharmProjects\garbage_classify\src\..\data\01下载图片\纸箱\纸箱_11.jpg: 416x640 (no detections), 20.2ms
image 3/43 D:\PycharmProjects\garbage_classify\src\..\data\01下载图片\纸箱\纸箱_12.jpg: 640x640 (no detections), 27.1ms
image 4/43 D:\PycharmProjects\garbage_classify\src\..\data\01下载图片\纸箱\纸箱_13.jpg: 480x640 (no detections), 25.5ms

3.标注转换

YOLO 分割出的数据是基于 COCO 数据集的标准，该数据集涵盖了 80 类常见物体。然而，我们需要训练的物体类别不仅限于这 80 类。因此，需要对分割出的文件进行一些调整。我们在前面下载图片时，根据 `classify.txt` 文件为每个类别创建了文件夹。现在，我们可以将分割出的物体类别修改为与这些文件夹名称对应的类别，以匹配我们自定义的分类需求。

    def process_label_line(self, line, image, width, height):
        """
        根据标签行在图像上绘制矩形框和类别ID。
    
        参数:
        - line: 包含标签信息的字符串，格式为"class_id x_center y_center box_width box_height"。
        - image: 待绘制的图像。
        - width: 图像的宽度。
        - height: 图像的高度。
    
        返回:
        - image: 绘制完成的图像。
        """
        # 移除字符串首尾的空白字符并按空格分割，获取标签信息
        line = line.strip().split()
        # 获取类别ID
        class_id = line[0]
        # 解析标签行中的坐标和尺寸信息
        x_center = float(line[1])
        y_center = float(line[2])
        box_width = float(line[3])
        box_height = float(line[4])
    
        # 计算矩形框的左上角和右下角坐标
        x_min = int((x_center - box_width / 2) * width)
        y_min = int((y_center - box_height / 2) * height)
        x_max = int((x_center + box_width / 2) * width)
        y_max = int((y_center + box_height / 2) * height)
    
        # 如果类别ID不在颜色映射中，则随机生成一个颜色
        if class_id not in self.color_map:
            self.color_map[class_id] = tuple(np.random.randint(0, 256, 3).tolist())
    
        # 获取对应类别ID的颜色
        color = self.color_map[class_id]
        # 设置矩形框的边框宽度
        thickness = 2
        # 在图像上绘制矩形框
        cv2.rectangle(image, (x_min, y_min), (x_max, y_max), color, thickness)
        # 在图像上绘制类别ID
        cv2.putText(image, f'{class_id}', (x_min, y_min - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, thickness)
    
        return image

自动生成的标注文件是按照每个像素点来标注的，也就是在物体的边缘处标记点。然而，这种格式不适合 YOLO 的训练需求。因此，我们需要将这些原始标注文件转换成 YOLO 可接受的格式，即中心点坐标加上宽高的形式。

例如：原始类别标注需要转换为 "class_id x_center y_center box_width box_height" 的格式。

还有一点，有部分图片可能包含其他目标物体，这些物体不是我们当前图片类下的目标，对于这些图片也是需要删除的。为了看出使用分割后又做了标注文件转换，在代码里我把标注后的框也画出来。生成了另一个文件夹，我们可以在生成新的文件中查看图片，图片是根据标注文件画的框，我们可以删除不符合的图片。

像上面那张图片，YOLO识别到了手和手机，标了两个框，但程序修改后两个框都是手机，那么这样的数据进入模型就是噪声需要删除。

这样的数据是符合要求的。

yolo文件下是不带框的图片，data下的是带框的图片，方便我们筛选不合格的图片。

4.同步图片

删除后，我们执行代码04同步数据代码，就可以同步yolo文件夹下的图片和筛选后的文件下所有文件保持一致，包括标注文件。

D:\Users\user\AppData\Local\Programs\Python\Python312\python.exe D:\PycharmProjects\garbage_classify\src\04同步数据.py 
删除文件: ../data\yolo\手机_yolo_训练图片\手机_17.jpg
删除文件: ../data\yolo\手机_yolo_训练图片\手机_28.jpg
删除文件: ../data\yolo\手机_yolo_训练图片\手机_40.jpg
删除文件: ../data\yolo\手机_yolo_训练图片\手机_41.jpg
删除文件: ../data\yolo\手机_yolo_训练图片\手机_30.jpg
删除文件: ../data\yolo\手机_yolo_训练图片\手机_13.jpg
删除文件: ../data\yolo\手机_yolo_训练图片\手机_45.jpg
删除文件: ../data\yolo\手机_yolo_训练图片\手机_8.jpg
删除文件: ../data\yolo\手机_yolo_训练图片\手机_29.jpg
文件夹同步完成。耗时：0.00 秒
删除文件: ../data\yolo\手机_yolo_训练标签\手机_45.txt
删除文件: ../data\yolo\手机_yolo_训练标签\手机_17.txt
删除文件: ../data\yolo\手机_yolo_训练标签\手机_40.txt
删除文件: ../data\yolo\手机_yolo_训练标签\手机_13.txt
删除文件: ../data\yolo\手机_yolo_训练标签\手机_41.txt
删除文件: ../data\yolo\手机_yolo_训练标签\手机_29.txt
删除文件: ../data\yolo\手机_yolo_训练标签\手机_28.txt
删除文件: ../data\yolo\手机_yolo_训练标签\手机_8.txt
删除文件: ../data\yolo\手机_yolo_训练标签\手机_30.txt

删除前45张图片

删除前36张图片

PS：如果发现删除后部分类别数据特别少，可以使用单独的 streamlit 页面下载图片重复上面步骤，代码在下载图片app.py

5.数据拆分

筛选后执行05数据拆分文件后将yolo文件夹下的图片复制到images目录，将txt文件移动到labels目录。复制完后会在yolo文件夹下创建，train，val，test 三个文件夹分别是训练集、验证集、测试集，这三个文件夹是最终YOLO训练使用的数据集，每个数据集文件夹下有两个文件夹images、和labels，图片和标注txt的名字都是一一对应的。

    # 验证集比例
    val_ratio = 0.2
    # 测试集比例
    test_ratio = 0.2

    if val_ratio + test_ratio >= 0.8:
        print("验证集和测试集比例之和必须小于 0.8。")
        return

    print("开始拆分数据集...")
    split_dataset(images_dir, labels_dir, base_dir, val_ratio=val_ratio, test_ratio=test_ratio)
    print("数据集拆分完成。")

D:\Users\user\AppData\Local\Programs\Python\Python312\python.exe D:\PycharmProjects\garbage_classify\src\05数据拆分.py 
开始创建数据集目录结构...
创建目录: ..\data\yolo\train\images
创建目录: ..\data\yolo\train\labels
创建目录: ..\data\yolo\val\images
创建目录: ..\data\yolo\val\labels
创建目录: ..\data\yolo\test\images
创建目录: ..\data\yolo\test\labels
开始拆分数据集...
数据集拆分和文件复制完成，耗时 0.34 秒
数据集拆分完成。

6.类别改成数字

执行06将txt文件中的英文类别根据classify.txt文件修改成数字，这也是YOLO要求的格式。

7.图片增强

07后自动增强数据，会根据训练集中的图片进行随机地翻转、调整亮度对比度、旋转、调整尺寸、位移缩放旋转、调整颜色、随机裁剪、透视变换、运动模糊和图像压缩等，生成新的图片和新的标注文件。可以控制增强轮次，循环一轮会增加一倍的训练数据。

def get_random_transform():
    """
    返回一个随机图像变换序列。这些变换用于数据增强，能够随机地翻转、调整亮度对比度、旋转、调整尺寸、
    位移缩放旋转、调整颜色、随机裁剪、透视变换、运动模糊和图像压缩。通过这些变换，可以增加模型训练的
    泛化能力，提高模型对各种情况的适应性。

    返回:
        A.Compose对象，包含一系列随机变换。
    """
    return A.Compose([
        # 随机水平翻转图像，翻转的概率是0到1之间的随机数。
        A.HorizontalFlip(p=np.random.uniform(0.0, 1.0)),
        # 随机调整图像的亮度和对比度，操作的概率是0到1之间的随机数。
        A.RandomBrightnessContrast(p=np.random.uniform(0.0, 1.0)),
        # 随机旋转图像90度的倍数，旋转的概率是0到1之间的随机数。
        A.RandomRotate90(p=np.random.uniform(0.0, 1.0)),
        # 将图像调整为640x640的尺寸。
        A.Resize(height=640, width=640),
        # 随机进行位移、缩放和旋转操作，各项参数都是随机决定的。
        A.ShiftScaleRotate(
            shift_limit=np.random.uniform(0.0, 0.1),
            scale_limit=np.random.uniform(0.0, 0.2),
            rotate_limit=np.random.uniform(0, 90),
            p=np.random.uniform(0.0, 1.0)
        ),
        # 使用CLAHE方法随机调整图像的对比度，操作的概率是0到1之间的随机数。
        A.CLAHE(clip_limit=np.random.uniform(1.0, 10.0), p=np.random.uniform(0.0, 1.0)),
        # 随机调整图像的Gamma值，操作的概率是0到1之间的随机数。
        A.RandomGamma(gamma_limit=(np.random.uniform(50, 150), np.random.uniform(50, 150)),
                      p=np.random.uniform(0.0, 1.0)),
        # 随机调整图像的色相、饱和度和明度，操作的概率是0到1之间的随机数。
        A.HueSaturationValue(
            hue_shift_limit=np.random.uniform(-30, 30),
            sat_shift_limit=np.random.uniform(-50, 50),
            val_shift_limit=np.random.uniform(-50, 50),
            p=np.random.uniform(0.0, 1.0)
        ),
        # 随机裁剪图像到480x480的尺寸，操作的概率是0到1之间的随机数。
        A.RandomCrop(height=480, width=480, p=np.random.uniform(0.0, 1.0)),
        # 随机对图像进行透视变换，操作的概率是0到1之间的随机数。
        A.Perspective(p=np.random.uniform(0.0, 1.0)),
        # 随机对图像进行运动模糊，操作的概率是0到1之间的随机数。
        A.MotionBlur(blur_limit=(3, 7), p=np.random.uniform(0.0, 1.0)),
        # 随机调整图像的压缩质量，操作的概率是0到1之间的随机数。
        A.ImageCompression(quality_lower=30, quality_upper=100, p=np.random.uniform(0.0, 1.0))
    ], bbox_params=A.BboxParams(format='yolo', label_fields=['class_labels']))

8.训练模型

类别名称根据图片类别的目录 classify.txt 转换执行生成classify文件.py可以生成names

# 模型训练时使用的yaml配置文件。该文件说明了数据的地址和待训练的类别
# 配置三个路径，分别对应训练数据、验证数据、测试数据的地址

# 训练数据用于模型的训练
train: D:/PycharmProjects/garbage_classify/data/yolo/train
# 验证数据用于模型训练过程中的评估和参数调试
val: D:/PycharmProjects/garbage_classify/data/yolo/val
# 测试数据用于模型完成训练后的测试
test: D:/PycharmProjects/garbage_classify/data/yolo/test

# 待训练的类别数量 number of classes
nc: 8

augmentation:  # 数据增强配置
  mosaic: true
  mixup: true
  hsv_h: 0.015  # 色调增强
  hsv_s: 0.7  # 饱和度增强
  hsv_v: 0.4  # 亮度增强
  degrees: 0.0  # 旋转角度
  translate: 0.1  # 平移
  scale: 0.5  # 缩放
  shear: 0.0  # 剪切
  perspective: 0.0  # 透视
  flipud: 0.0  # 垂直翻转
  fliplr: 0.5  # 水平翻转
  mosaic_scale: 1.0  # Mosaic 缩放
  mixup_scale: 0.0  # MixUp 缩放
  copy_paste: 0.0  # 复制粘贴

# 类别名称
names: [
'Mobilephones', 'Plastic', 'Cardboard', 'Zip-topcan', 'LightBulb', 'Batteries', 'BananaSkin', 'CornKernel'
]

def train_model():
    # 检查是否有可用的GPU
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")

    # 清空显存缓存
    torch.cuda.empty_cache()

    model = YOLO("yolov8m.pt")

    # 在训练前再次清空缓存
    torch.cuda.empty_cache()

    model.train(
        data='./data.yaml',
        epochs=10,
        batch=3,  # 减少批处理大小
        lr0=0.0001,  # 更低的初始学习率
        amp=False,  # 启用自动混合精度训练
        device=device,  # 确保使用GPU
        mosaic=True,  # 启用 Mosaic 增强
        mixup=True,  # 启用 MixUp 增强
        augment=True  # 启用其他数据增强策略
    )

    # 训练后再次清空缓存
    torch.cuda.empty_cache()

训练中

9.测试模型

def speak_predictions(predictions, class_mapping):
    """将预测结果转换成语音播报
    """
    for prediction in predictions:
        class_id = prediction['class_id']
        class_name = prediction['class_name']
        confidence = prediction['confidence']

        chinese_name = class_mapping.get(class_name)  # 查找中文名称
        # 根据类别决定垃圾种类
        if class_name in ["Mobilephones", "Plastic", "Cardboard", "Zip-topcan"]:
            garbage_type = "可回收垃圾"
        elif class_name in ["LightBulb", "Batteries"]:
            garbage_type = "有害垃圾"
        elif class_name in ["BananaSkin", "CornKernel"]:
            garbage_type = "湿垃圾"
        else:
            garbage_type = "未知类别"

        text = f"预测类别: {class_id}, 预测类别名称: {class_name}, 中文名称: {chinese_name}, 置信度: {confidence:.2f}"
        print(text)

        textSay = f"识别到的物体是: {chinese_name}, 属于: {garbage_type}"
        print(textSay)
        speak_text(textSay)
if __name__ == "__main__":
    model_path = 'D:/yolo_materials/ultralytics/runs/detect/train80/weights/best.pt'  # 模型路径
    model = load_model(model_path)
    image_filename = 'merged2222.png'  # 图片文件名
    image_path = os.path.join('D:/Users/user/Desktop/深度学习项目/手机/', image_filename)

    # 获取预测结果
    predict_results = predict_image(model, image_path, save=True, save_txt=True)
    # 提取预测结果
    predictions = extract_predictions(predict_results)
    # 加载类别映射文件
    class_mapping = load_class_mapping('classify.txt')

    # 显示图片
    display_image(os.path.join(predict_results[0].save_dir, image_filename))
    # # 语音播报预测结果
    speak_predictions(predictions, class_mapping)
    video_source = 'D:/Users/user/Desktop/深度学习项目/手机/df82253f0201a27c745b5ac518bc6ae7.mp4'

D:\Users\user\AppData\Local\Programs\Python\Python312\python.exe D:\PycharmProjects\garbage_classify\src\09测试.py 

image 1/1 D:\Users\user\Desktop\深度学习项目\手机\img_17465.jpg: 640x512 1 Plastic, 56.5ms
Speed: 2.5ms preprocess, 56.5ms inference, 84.9ms postprocess per image at shape (1, 3, 640, 512)
Results saved to D:\yolo_materials\ultralytics\runs\detect\predict160
1 label saved to D:\yolo_materials\ultralytics\runs\detect\predict160\labels

测试结果

代码地址

garbage_classify: 垃圾分类

原本计划把这些都做成页面，傻瓜式操作，因为太懒了没做。代码还有很多优化的地方，欢迎交流。

我们详细介绍了如何从零开始创建自定义数据集，并使用YOLO模型进行目标检测的全流程。通过自动化图片下载、自动标注、标注格式转换、数据增强等步骤，我们能够快速构建高质量的数据集，大幅度减少手动标注的时间和精力投入。这不仅提高了数据处理的效率，还确保了模型训练数据的多样性和准确性。希望通过这篇文章，给你带来收获。