masks2segments function does not always work correctly #9784
Description
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:
- 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
- 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]
- 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!