实现效果
原始图片
使用YOLOv10检测与MobileSAM分割后的结果如下:
引言
本文基于前沿的YOLOv10目标检测模型
与轻量化分割一切的MobileSAM模型
实现物体的目标检测与分割效果。本文给出了完整的实现步骤与代码详解,供小伙伴们学习。**所有相关文件、模型及相关资料均已打包,文末可免费获取。**有任何问题的话,欢迎在评论区留言讨论。
实现步骤
首先我们需要下载YOLOv10源码与MobileSAM源码并解压,地址如下:
环境配置
我们进入到yolov10
源码的项目目录,然后使用conda进行项目虚拟环境配置:
conda create -n yolov10 python=3.9 conda activate yolov10 pip install -r requirements.txt pip install -e .
YOLOv10进行目标检测
这里我们直接使用官方提供的yolov10n.pt
进行检测,并获取检测框的位置信息,用于后续MobileSAM进行分割,代码如下:
# 读取图像 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模型:
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)
使用模型对目标进行分割:
# 矩形框提示 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同样可以与其他目标检测模型以这种方式进行结合使用,达到分割目标的效果。
好了,这篇文章就介绍到这里,感谢点赞关注,更多精彩内容持续更新中~