How to segment Objects with YOLOv8

Introduction

Segmentation is a key task in computer vision that has a wide range of uses in areas including medical imaging, robotics, and self-driving cars. YOLO (You Only Look Once) is a group of object detection models that are famous for their quickness and precision. To segment using YOLO, it is possible to expand a YOLO object detection model to anticipate pixel-wise masks for each object found in an image.

In this blog post, we will examine how to use YOLO for segmentation in practical applications. The subjects that we will discuss include:

• Choosing a pre-trained YOLO model
• Training the YOLO model
• Evaluating the model

Why use YOLO for segmentation?

There are various benefits to using YOLO for segmentation.
YOLO models are speedy and ideal for real-time usage:

• Speed: YOLO models are very fast, making them suitable for real-time applications.
• Accuracy: YOLO models are also very accurate, achieving state-of-the-art results on many segmentation benchmarks.
• Robustness: YOLO models are robust to noise and occlusions, making them suitable for challenging real-world environments.

Real-world applications of YOLO segmentation

YOLO segmentation has various practical applications, including medical imaging and robotics:

• Medical imaging: YOLO image analysis can identify tumors and other abnormalities in medical images. This data can aid doctors in diagnosing illnesses and creating treatment plans.
• Robotics: YOLO segmentation can be useful for separating objects in an area so that robots can interact with them safely and proficiently. For instance, robots can apply YOLO segmentation to select and position items in depots, or explore congested settings.
• Autonomous driving: YOLO segmentation can segment cars, people, and other objects on the road. This helps autonomous cars navigate safely.

How to use YOLO for images and videos

Step 1: Installing the necessary libraries

pip install opencv-python ultralytics numpy

Step 2: Importing libraries

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

Step 3: Choose your model

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

On this website, you can compare different models and weigh up their respective advantages and disadvantages. In this case we have chosen yolov8m-seg.pt.

Step 4: Segmenting Objects in Images with YOLOv8

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.5

results = model.predict(img, conf=conf)
colors = [random.choices(range(256), k=3) for _ in classes_ids]
print(results)
for result in results:
    for mask, box in zip(result.masks.xy, result.boxes):
        points = np.int32([mask])
        # cv2.polylines(img, points, True, (255, 0, 0), 1)
        color_number = classes_ids.index(int(box.cls[0]))
        cv2.fillPoly(img, points, colors[color_number])

Step 5: Save and Plot the result Image

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

cv2.imwrite("YourSavePath", img)

The whole code:

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

model = YOLO("yolov8m-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.5

results = model.predict(img, conf=conf)
colors = [random.choices(range(256), k=3) for _ in classes_ids]
print(results)
for result in results:
    for mask, box in zip(result.masks.xy, result.boxes):
        points = np.int32([mask])
        # cv2.polylines(img, points, True, (255, 0, 0), 1)
        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)

Conclusion

In this tutorial we have learned how to segments objects with YOLOv8 in images. If you found this code helpful, please clap your hands and comment on this post! I would also love for you to follow me to learn more about Data Science and other related topics. Thanks for reading!