YOLO Object Segmentation with Python

Object segmentation goes a step further than object detection and involves identifying individual objects in an image and segmenting them from the rest of the image. The output of an instance segmentation model is a set of masks or contours that outline each object in the image, along with class labels and confidence scores for each object. Instance segmentation is useful when you need to know not only where objects are in an image, but also what their exact shape is.

YOLOv8 Segmentation models have -seg suffix and are pretrained on COCO dataset with the following Classes.

['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush']

YOLOv8 pretrained Segmentation models (nano, small, medium, large and extra large based on number of parameters) are shown in the table below:

Setup UltraLytics for YOLOv8

%pip install ultralytics
import ultralytics
ultralytics.checks()

Load YOLOv8 for Segmentation

from ultralytics import YOLO

# Load a model
# You can use different YOLOv8 variants (yolov8n-seg, yolov8s-seg, yolov8m-seg, yolov8l-seg, yolov8nx-seg)
model = YOLO('yolov8n-seg.pt')  # load a pretrained model

# Use the model
results = model('https://ultralytics.com/images/zidane.jpg')  # predict on an image
# Save the output image after segmentation in Google Colab 
results[0].save('/content/output_seg.jpg')

# Print the COCO dataset classes on which model is trained.
print(model.names.values())

Input Image

Display the Output Image after Segmentation

from IPython.display import Image

# Display the image in Google Colab 
Image(filename='/content/output_seg.jpg')

Output Predicted Image

Print Objects Masks

for r in results:
    print(r.masks)