ultralytics yolo segmentation 分割示例：加载官方segmentation 模型进行推理

news2024/10/9 16:36:27

Ultralytics YOLO 是计算机视觉和 ML 领域专业人士的高效工具。

安装 ultralytics 库：

pip install ultralytics

该示例的模型地址：ultralyticsyolosegmentation分割模型资源-CSDN文库

大家也可以在官方连接下载，因为有时下载容易断开链接，这里做了备份。

示例具体实现代码如下：

#-*-coding:utf-8-*-
# date:2024-10
# Author: DataBall XIAN
# function: 实例分割
import cv2
import numpy as np
from ultralytics import YOLO
import random
# Load a model
ckpt_dir = "./ckpt/" # 模型缓存地址
m = YOLO(ckpt_dir + "yolo11m-seg.pt")
path_img = "1.jpg"
res = m.predict(path_img)

# 结果遍历可视化
for r in res:
    img = np.copy(r.orig_img)
    img_o = img.copy()
    # 绘制分割轮廓信息
    for ci, c in enumerate(r):
        # label
        label = c.names[c.boxes.cls.tolist().pop()]
        print("label:",label)
        # bbox
        x1, y1, x2, y2 = c.boxes.xyxy.cpu().numpy().squeeze().astype(np.int32) # 边界框
        # mask
        b_mask = np.zeros(img.shape[:2], np.uint8)
        # Create contour mask
        contour = c.masks.xy.pop().astype(np.int32).reshape(-1, 1, 2)
        _ = cv2.drawContours(b_mask, [contour], -1, (255,255,255), cv2.FILLED)
        mask3ch = cv2.cvtColor(b_mask, cv2.COLOR_GRAY2BGR)
        # 融合分割信息
        RGB_ = np.array((random.randint(100,255),random.randint(80,240),random.randint(90,233))) # 获取随机颜色
        idx_= np.where(b_mask==255)
        img[idx_] = (img[idx_].astype(np.float32)*0.2 + RGB_.astype(np.float32) *0.8).astype(np.uint8) # 分割颜色融合

    stk_ = np.hstack((img_o,img))
    cv2.namedWindow("img",0)
    cv2.imshow("img",stk_)
    if cv2.waitKey(0) == 27:
        break

示例可视化效果：