# YOLOv5 Image Segmentation: A Complete Training and Implementation Guide

> Learn how to use YOLOv5 for image segmentation. This guide covers training and implementation of segmentation models with pixel-wise object masks and fast inference.

- Repository: [Ultralytics/yolov5](https://github.com/ultralytics/yolov5)
- Tags: tutorial
- Published: 2026-03-06

---

**YOLOv5 supports instance segmentation through the `SegmentationModel` class, which extends standard detection architectures with a mask-head prototype branch, enabling pixel-wise object masks while maintaining the same fast inference pipeline used for object detection.**

The ultralytics/yolov5 repository includes a dedicated instance segmentation pipeline that extends the framework's detection capabilities. YOLOv5 image segmentation leverages a prototype-based mask generation system to produce high-quality pixel-wise object boundaries without sacrificing the real-time performance that makes YOLO models popular for production deployments.

## How YOLOv5 Image Segmentation Works

### Architecture Overview

The segmentation implementation centers on the `SegmentationModel` class defined in [`models/yolo.py`](https://github.com/ultralytics/yolov5/blob/main/models/yolo.py) (lines 36-44). This class inherits from `DetectionModel` and reuses the same backbone, neck, and detection head architecture, then appends a prototype mask branch that generates per-object masks through a coefficient-based approach.

The architecture produces masks by combining a low-resolution **prototype tensor** (`protos`) with per-object **coefficient vectors** (`masks_in`). This design, implemented in [`utils/segment/general.py`](https://github.com/ultralytics/yolov5/blob/main/utils/segment/general.py) (lines 24-41), multiplies the up-sampled prototype tensor by object-specific coefficients to generate full-resolution masks, which are then cropped to each detected bounding box.

### Mask Post-Processing Pipeline

After inference, the `crop_mask`, `process_mask_*`, and `scale_image` functions in [`utils/segment/general.py`](https://github.com/ultralytics/yolov5/blob/main/utils/segment/general.py) (lines 9-23) handle the conversion from model outputs to final mask predictions. These utilities manage up-sampling, resizing to original image dimensions, and cropping masks to fit precise object boundaries.

### Loss Function Implementation

The training loss combines standard detection objectives (box, objectness, classification) with a dedicated mask loss. The `ComputeLoss` class in [`utils/segment/loss.py`](https://github.com/ultralytics/yolov5/blob/main/utils/segment/loss.py) implements this using **binary cross-entropy** to optimize the prototype-coefficient formulation, ensuring accurate mask boundaries alongside detection metrics.

## Training a YOLOv5 Segmentation Model

### Dataset Preparation

Training requires images paired with binary segmentation masks. Organize your data with separate `images/` and `masks/` directories, where mask filenames match their corresponding images. Each mask should be a binary image where object pixels are marked for segmentation targets.

Create a data configuration YAML file specifying paths, class count (`nc`), and class names. The [`utils/segment/dataloaders.py`](https://github.com/ultralytics/yolov5/blob/main/utils/segment/dataloaders.py) module loads these image-mask pairs, applies augmentations, and returns tensors in the format `(imgs, targets, paths, _, masks)`.

### Training Script Configuration

Use [`segment/train.py`](https://github.com/ultralytics/yolov5/blob/main/segment/train.py) as the dedicated training entry point. This script mirrors the standard detection trainer but adds segmentation-specific arguments including `--mask-ratio` (controls ground-truth mask down-sampling, default 4) and `--no-overlap` (disables overlapping masks for faster training).

```bash
python segment/train.py \
  --data my_seg.yaml \
  --weights yolov5s-seg.pt \
  --cfg yolov5s-seg.yaml \
  --batch-size 16 \
  --imgsz 640 \
  --epochs 100 \
  --device 0

```

Model configurations reside in `models/segment/*.yaml` (such as [`yolov5s-seg.yaml`](https://github.com/ultralytics/yolov5/blob/main/yolov5s-seg.yaml)), defining the architecture backbone and head specifically for segmentation tasks.

## Running Inference with Segmentation

For prediction, load pretrained segmentation weights and enable mask output:

```python
import torch

# Load model

model = torch.hub.load('ultralytics/yolov5', 'custom', path='yolov5s-seg.pt')

# Run inference with masks enabled

results = model('image.jpg', size=640, masks=True)
results.show()  # Display with overlaid masks

```

To convert output masks to polygon format for COCO-style evaluation, use the `masks2segments` utility:

```python
from utils.segment.general import masks2segments

# masks shape: (N, H, W)

segments = masks2segments(masks, strategy='largest')

# Returns list of (M_i, 2) numpy arrays

```

## Key Files and Implementation Details

- **[`segment/train.py`](https://github.com/ultralytics/yolov5/blob/main/segment/train.py)**: Training entry point with segmentation-specific data loading and loss computation
- **[`segment/predict.py`](https://github.com/ultralytics/yolov5/blob/main/segment/predict.py)**: CLI inference tool for mask generation
- **[`models/yolo.py`](https://github.com/ultralytics/yolov5/blob/main/models/yolo.py)**: Contains `SegmentationModel` class definition extending detection architecture
- **[`utils/segment/dataloaders.py`](https://github.com/ultralytics/yolov5/blob/main/utils/segment/dataloaders.py)**: Handles image-mask pair loading and augmentation
- **[`utils/segment/loss.py`](https://github.com/ultralytics/yolov5/blob/main/utils/segment/loss.py)**: Implements combined detection and mask loss functions
- **[`utils/segment/general.py`](https://github.com/ultralytics/yolov5/blob/main/utils/segment/general.py)**: Post-processing utilities including `crop_mask` and prototype generation
- **`models/segment/*.yaml`**: Architecture configurations (e.g., [`yolov5s-seg.yaml`](https://github.com/ultralytics/yolov5/blob/main/yolov5s-seg.yaml))

## Summary

- YOLOv5 image segmentation extends the detection architecture via `SegmentationModel`, adding a prototype-based mask head to [`models/yolo.py`](https://github.com/ultralytics/yolov5/blob/main/models/yolo.py)
- Training uses [`segment/train.py`](https://github.com/ultralytics/yolov5/blob/main/segment/train.py) with binary mask annotations stored alongside images, loaded by [`utils/segment/dataloaders.py`](https://github.com/ultralytics/yolov5/blob/main/utils/segment/dataloaders.py)
- The loss function in [`utils/segment/loss.py`](https://github.com/ultralytics/yolov5/blob/main/utils/segment/loss.py) combines detection metrics with binary cross-entropy mask optimization
- Post-processing utilities in [`utils/segment/general.py`](https://github.com/ultralytics/yolov5/blob/main/utils/segment/general.py) handle mask up-sampling, cropping, and scaling to original image dimensions
- Pretrained weights (e.g., `yolov5s-seg.pt`) and configuration files in `models/segment/` enable immediate training or fine-tuning

## Frequently Asked Questions

### What is the difference between YOLOv5 detection and segmentation models?

Detection models predict bounding boxes and class labels, while segmentation models add pixel-wise mask predictions. The `SegmentationModel` class in [`models/yolo.py`](https://github.com/ultralytics/yolov5/blob/main/models/yolo.py) inherits from `DetectionModel` and appends a prototype mask branch, enabling instance segmentation that identifies exact object boundaries rather than just rectangular regions.

### How do I prepare mask annotations for YOLOv5 training?

Store binary masks as PNG files in a `masks/` directory parallel to your `images/` folder, using identical filenames to link masks with images. Each mask should contain binary pixel values indicating object presence. The dataloader in [`utils/segment/dataloaders.py`](https://github.com/ultralytics/yolov5/blob/main/utils/segment/dataloaders.py) automatically pairs these files during training batch generation.

### Can I use pretrained detection weights for segmentation training?

Yes, you can initialize segmentation training with detection backbone weights, though using dedicated segmentation pretrained weights (e.g., `yolov5s-seg.pt`) typically yields better results. The backbone architecture is shared between detection and segmentation models, allowing knowledge transfer from detection datasets to segmentation tasks.

### What hardware requirements are needed for training segmentation models?

Segmentation training requires more GPU memory than detection due to the additional mask processing and prototype generation overhead. A modern NVIDIA GPU with at least 8GB VRAM is recommended for training `yolov5s-seg` at 640px resolution with batch size 16. Reduce `--batch-size` or increase `--mask-ratio` (default 4) if encountering memory constraints.