使用yolo-seg模型实现自定义自动动态抠图

yolov8导航

        如果大家想要了解关于yolov8的其他任务和相关内容可以点击这个链接，我这边整理了许多其他任务的说明博文，后续也会持续更新，包括yolov8模型优化、sam等等的相关内容。

YOLOv8（附带各种任务详细说明链接）

介绍

在这篇博客中，我们将探索如何使用YOLO-SEG模型来自动进行图像中对象的抠图。YOLO-SEG是一种结合了YOLO（You Only Look Once）快速目标检测能力和图像分割的深度学习模型，能够在单个网络中同时实现目标检测和像素级分割。

安装与设置：点击上面yolov8的导航跳转有参考

或者点击这里：超详细的YOLOv8安装与测试指南：让计算机视觉任务变得简单_yolo ultralytics 安装-CSDN博客

数据准备与训练：

看这里：

超详细YOLOv8实例分割全程概述：环境、训练、验证与预测详解_yolov8环境搭建-CSDN博客

抠图应用

训练完成后，你可以使用训练好的模型对新图像进行自动抠图。YOLO-SEG能够识别图像中的对象，并生成对应的分割掩码，从而实现精准抠图。

from ultralytics import YOLO
from matplotlib import pyplot as plt
import numpy as np
import cv2

# 加载sam模型，如果没有这个框架也会自动的为你下载
model = YOLO('yolov8n-seg.pt')

# 读取一张图片，
orig_img = cv2.imread('test.jpg')
# 这里特别注意，因为使用yolov8训练的时候默认会把图片resize成448*640的尺寸，所以这里也得改成你训练的尺寸
orig_img = cv2.resize(orig_img, (640, 448))  # 注意OpenCV中尺寸是先宽度后高度


# 使用模型进行推理， 后面save=True的参数可以输出测试分割的图片
results = model(orig_img,save=True)

# 这里是我获取mask编码部分的。
mask = results[0].masks.data[0].cpu().numpy().astype(np.bool_)


# 应用掩码到原始图像
# 如果你想要背景是黑色的
masked_image = np.zeros_like(orig_img)
masked_image[mask] = orig_img[mask]

# 如果你想要背景透明（假设原始图像是RGB格式）
# 创建一个RGBA图像，其中背景是透明的
alpha_channel = np.ones(mask.shape, dtype=orig_img.dtype) * 255  # 创建alpha通道
masked_image_rgba = np.dstack((masked_image, alpha_channel))  # 将alpha通道添加到RGB通道
masked_image_rgba[~mask] = (0, 0, 0, 0)  # 将非掩码区域设置为透明

# 保存图片，黑色背景
cv2.imwrite('masked_image_test.jpg', masked_image)

# 如果背景是透明的
cv2.imwrite('masked_image_test.png', masked_image_rgba)

原图：

正常输出结果：

抠图结果：

        到了这一步。你学废了吗。