Image Segmentation With YOLOv11

Image Segmentation With YOLOv11

Short example of doing image segmentation with YOLOv11 (Ultralytics)

from ultralytics import YOLO
import random
import cv2
import numpy as np

model = YOLO("yolo11x-seg.pt") 

img = cv2.imread("YourImagePath")

# if you want all classes
yolo_classes = list(model.names.values())
classes_ids = [yolo_classes.index(clas) for clas in yolo_classes]

conf = 0.2

results = model.predict(img, conf=conf)
colors = [random.choices(range(256), k=3) for _ in classes_ids]

for result in results:
    for mask, box in zip(result.masks.xy, result.boxes):
        points = np.int32([mask])
        color_number = classes_ids.index(int(box.cls[0]))
        cv2.fillPoly(img, points, colors[color_number])

cv2.imshow("Image", img)
cv2.waitKey(0)

cv2.imwrite("YourSavePath", img)

Sumber: