《博主简介》
小伙伴们好,我是阿旭。专注于人工智能、AIGC、python、计算机视觉相关分享研究。
✌更多学习资源,可关注公-仲-hao:【阿旭算法与机器学习】,共同学习交流~
👍感谢小伙伴们点赞、关注!
《------往期经典推荐------》
二、机器学习实战专栏【链接】 ,已更新31期,欢迎关注,持续更新中~~
三、深度学习【Pytorch】专栏【链接】
四、【Stable Diffusion绘画系列】专栏【链接】
五、YOLOv8改进专栏【链接】,持续更新中~~
六、YOLO性能对比专栏【链接】,持续更新中~
《------正文------》
实现效果
原始图片
使用YOLOv10检测与MobileSAM分割后的结果如下:
引言
本文基于前沿的YOLOv10目标检测模型
与轻量化分割一切的MobileSAM模型
实现物体的目标检测与分割效果。本文给出了完整的实现步骤与代码详解,供小伙伴们学习。**所有相关文件、模型及相关资料均已打包,文末可免费获取。**有任何问题的话,欢迎在评论区留言讨论。
实现步骤
首先我们需要下载YOLOv10源码与MobileSAM源码并解压,地址如下:
环境配置
我们进入到yolov10
源码的项目目录,然后使用conda进行项目虚拟环境配置:
python
conda create -n yolov10 python=3.9
conda activate yolov10
pip install -r requirements.txt
pip install -e .
YOLOv10进行目标检测
这里我们直接使用官方提供的yolov10n.pt
进行检测,并获取检测框的位置信息,用于后续MobileSAM进行分割,代码如下:
python
# 读取图像
image = cv2.imread('test1.jpg')
# 加载YOLOv10模型
model = YOLOv10("yolov10_main/yolov10n.pt")
results = model.predict(image)
# 绘制检测结果
results[0].show()
# 获取结果
res = results[0]
# 获取检测框位置信息
box_locations = res.boxes.xyxy.cpu().numpy().astype(int)
print(box_locations)
检测结果如下所示:
MobilseSAM进行分割
我们将上述检测框信息作为提示信息,传入MobileSAM对目标进行分割,具体代码如下:
加载MobileSAM模型:
python
from mobile_encoder.setup_mobile_sam import setup_model
from segment_anything import SamPredictor
# 加载分割模型
checkpoint = torch.load('MobileSAM-master/weights/mobile_sam.pt',map_location=torch.device('cpu'))
mobile_sam = setup_model()
mobile_sam.load_state_dict(checkpoint,strict=True)
device = "cpu"
mobile_sam.to(device=device)
mobile_sam.eval()
predictor = SamPredictor(mobile_sam)
predictor.set_image(image)
使用模型对目标进行分割:
python
# 矩形框提示
for box in box_locations:
input_box = np.array([box[0], box[1], box[2], box[3]])
masks, _, _ = predictor.predict(
point_coords=None,
point_labels=None,
box=input_box[None, :],
multimask_output=False,
)
# masks为布尔值
# 显示mask
image = add_mask(image, masks[0])
# 显示检测框
image = show_box(image, box)
cv2.imshow("res", image)
cv2.imwrite("res.jpg", image)
cv2.waitKey(0)
显示结果如下:
可以看到,MobileSAM可以很好的对YOLOv10检测后的结果进行分割。MobileSAM同样可以与其他目标检测模型以这种方式进行结合使用,达到分割目标的效果。
好了,这篇文章就介绍到这里,感谢点赞关注,更多精彩内容持续更新中~
资料获取
关于本文的相关代码及数据集资料都已打包好,供需要的小伙伴们学习,获取方式如下: