Flat edges on segmentation masks?

[joedegroot]

I've noticed that YOLOv8 segmentation masks have flat edges where the mask is adjacent to the bounding box of the detected instance. This is most noticeable on the top edge of the mask, but it occurs on all four edges of the mask. These flat edges result in an incorrect segmentation mask, and lead to incorrect measurements of object size and shape, especially for small objects.

The issue exists both with the YOLOv8 pretrained models, as well as models I've trained using transfer learning with one of the pretrained models as a backbone.

What is the cause of this segmentation mask error? Is there any way to fix or improve it?

Example (note the flat tops of all the puppies heads):

Things I've tried:

• Larger model (i.e. yolov8x-seg.py instead of yolov8m-seg.py). This did not improve the mask significantly.
• More data - I've increased the number of images/labels during training, but this did not help. Also, the issue exists with the yolov8 pretrained models, so I assume it's inherent to the model in some way.
• Larger imgsz: I've gone as far as to run prediction on an image the same size as the model, so that no resizing takes place. Issue still persists.
• Enable retina_masks=True during prediction. This does not seem to affect the output at all.
• I've manually processed the box data returned by the model, to see if adjusting the confidence threshold for the mask helps at all. This certainly allows me to change the appearance of the mask (make it smaller or larger), but no matter what, the flat edges persist.

This issue can be easily reproduced with the CLI like this:
yolo predict task=segment model=yolov8m-seg.pt imgsz=640 source=https://ultralytics.com/images/zidane.jpg show_boxes=False
Output (cropped, to show the issue):

[pderrenger]

Hi there! 👋

Thanks for bringing this to our attention. The flat edges you're observing on segmentation masks are likely due to the limitations of the model's architecture in capturing the finer details at the boundaries of objects, especially for smaller objects.

One potential way to mitigate this is to experiment with different segmentation head architectures or loss functions that may be more sensitive to boundary details. Additionally, post-processing techniques like CRF (Conditional Random Fields) can sometimes help refine the edges of segmentation masks.

Here's a quick example of how you might apply CRF post-processing:

from ultralytics import YOLO
import pydensecrf.densecrf as dcrf
from pydensecrf.utils import unary_from_softmax

# Load your model
model = YOLO('yolov8m-seg.pt')

# Perform prediction
results = model('path/to/image.jpg')

# Assuming you have a single image and mask
predicted_mask = results[0].masks.data[0]  # Get the first predicted mask

# Prepare the unary potentials and run CRF
softmax = predicted_mask.softmax(dim=0).cpu().numpy()  # Convert to softmax
unary = unary_from_softmax(softmax)
d = dcrf.DenseCRF2D(predicted_mask.shape[1], predicted_mask.shape[0], 2)  # width, height, n_classes
d.setUnaryEnergy(unary)

# You can set your own hyperparameters for CRF here
d.addPairwiseGaussian(sxy=3, compat=3)
d.addPairwiseBilateral(sxy=80, srgb=13, rgbim=image, compat=10)

# Run inference
q = d.inference(5)
refined_mask = np.argmax(np.array(q), axis=0).reshape((predicted_mask.shape[1], predicted_mask.shape[0]))

# Now `refined_mask` is your CRF-refined mask

Please note that the above code is just a starting point and you'll need to install pydensecrf and adjust the parameters to fit your specific use case.

We're always working on improving our models, and your feedback is valuable in this process. If you have any more insights or questions, feel free to share! 🚀

[pderrenger]

Hi Ron,

Thank you for reaching out. The proposed solution can help mitigate the issue, but it may not completely eliminate it, especially for small, round objects. I recommend checking if the issue persists with the latest versions of the YOLO package, as updates often include improvements and bug fixes. If the problem continues, further adjustments or custom post-processing techniques might be necessary.

[jorgemasgomez]

Hi! I've been dealing with this issue of flat edges. My solution has been to set a higher imgsz in the training function than the images in the dataset. For example, my dataset was at 320 px, and I set imgsize=640. My question is, why does this happen? I've experienced it with several datasets — if you increase the imgsize, these flat edges disappear.

[glenn-jocher]

Hi! Increasing imgsz during training allows the model to capture more detail, which can help reduce flat edges in segmentation masks. This happens because higher resolution provides more information for the model to learn finer details.

[jorgemasgomez]

Hi again!

Sorry for insisting, but I think this topic deserves some kind of post because there is no information available, and it completely changes the segmentation results paradigm. For small and round objects, the results are much better when removing these flat edges.

Also, today I’ve been dealing with SAHI, and I was having issues because I set the slices to 320 px and the imgsize in the AutoDetectionModel.from_pretrained function to 640 px. This caused problems because the different slices weren’t being joined correctly, but to me, the issue didn’t make sense: if I’m training my YOLO model with 320 px images and resizing them to 640 px, why was SAHI failing if I was doing the same thing?

The point is that after trying different things, I decided to set the slices to 640 px and
the imgsize in AutoDetectionModel.from_pretrained to 1280 px. And it worked! The segmentation is much better.

I think there is a lack of documentation on this topic to properly justify in scientific publications how this resizing affects results. Any idea where I can find more information? If you could look into this and make a post, that would be amazing.

Thanks again for answering all our questions; you’re doing great work!

Without resizing

Resizing

[glenn-jocher]

Thank you for your insights. We appreciate your feedback and will consider it for future documentation updates. For now, exploring academic papers on image resizing and segmentation might provide additional context. If you have further questions, feel free to ask here.

[TrinhNC]

Have you solved this problem?

[glenn-jocher]

@TrinhNC the flat edges on segmentation masks are a known limitation due to the mask resolution. Improving mask quality may require model architecture adjustments or post-processing techniques. For further insights, consider exploring the YOLO community discussions.

[Ron-Ventura]

no ):
We just made manipulation on the data post processing, added some layers of filters

[glenn-jocher]

Thanks for sharing your approach. If you have further questions or need more assistance, feel free to ask!