masks2segments function does not always work correctly

[vladoossss]

Search before asking

• I have searched the YOLOv5 issues and found no similar bug report.

YOLOv5 Component

No response

Bug

When i use segment/predict.py --save-txt to save segmentation results, sometimes i had an error.

The reason for this was an empty mask (only zeros), which was passed to the input of the masks2segments function. In this regard, c = cv2.findContours(x, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[0] produced an empty tuple(due to previous transformations over the mask) and then there was an error.

I made the following changes in 3 places to avoid this:

1. general.py

def masks2segments(masks, strategy='largest'):
    # Convert masks(n,160,160) into segments(n,xy)
    segments = []
    for x in masks.int().cpu().numpy().astype('uint8'):
        c = cv2.findContours(x, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[0]
        if c:
            if strategy == 'concat':  # concatenate all segments
                c = np.concatenate([x.reshape(-1, 2) for x in c])
            elif strategy == 'largest':  # select largest segment
                c = np.array(c[np.array([len(x) for x in c]).argmax()]).reshape(-1, 2)
            segments.append(c.astype('float32'))
        else:
            segments.append(0)
    return segments

2. segment/predict.py

# Segments
if save_txt:
    segments = reversed(masks2segments(masks))
    segments = [scale_segments(im.shape[2:], x, im0.shape).round() if not isinstance(x, int) else 0 for x in segments]

3. segment/predict.py

# Write results
for j, (*xyxy, conf, cls) in enumerate(reversed(det[:, :6])):
    if save_txt:  # Write to file
        if not isinstance(segments[j], int):
            segj = segments[j].reshape(-1)  # (n,2) to (n*2)
            line = (cls, *segj, conf) if save_conf else (cls, *segj)  # label format
            with open(f'{txt_path}.txt', 'a') as f:
                f.write(('%g ' * len(line)).rstrip() % line + '\n')
        else:
            line = (cls, segments[j], conf) if save_conf else (cls, segments[j])  # label format
            with open(f'{txt_path}.txt', 'a') as f:
                f.write(('%g ' * len(line)).rstrip() % line + '\n')

It helps me to avoid an error.
I understand that my changes are not perfect. The main goal was to point out the problem itself.

Environment

No response

Minimal Reproducible Example

No response

Additional

No response

Are you willing to submit a PR?

• Yes I'd like to help by submitting a PR!

[glenn-jocher]

@vladoossss thanks! Can you please submit a PR with your proposed changes?

Change #1 doesn't suffice by itself?

[vladoossss]

@vladoossss thanks! Can you please submit a PR with your proposed changes?

Change #1 doesn't suffice by itself?

Yes, i've submitted new PR, based on this issue.
All 3 changes together worked for me.

[glenn-jocher]

@vladoossss good news 😃! Your original issue may now be fixed ✅ in PR #9785.

To receive this update:

• Git – git pull from within your yolov5/ directory or git clone https://github.com/ultralytics/yolov5 again
• PyTorch Hub – Force-reload model = torch.hub.load('ultralytics/yolov5', 'yolov5s', force_reload=True)
• Notebooks – View updated notebooks
• Docker – sudo docker pull ultralytics/yolov5:latest to update your image

Thank you for spotting this issue and informing us of the problem. Please let us know if this update resolves the issue for you, and feel free to inform us of any other issues you discover or feature requests that come to mind. Happy trainings with YOLOv5 🚀!