YOLOv5 segmentation model support

[AyushExel]

Usage:

# train
python segment/train.py --cfg yolov5(m/s/l)-seg 
# detect
python segment/predict.py --source

TODO:

• Resolve conflict
• Resolve pre-commit
• Enable overlap-mask by default
• Clean callback related residue code
• Write segmentation related readme
• torch.deterministic solution for segment?
• merge non_max_suppression_masks() with non_max_suppression()?
• remove edict dependency?
• DetectMultiBackend compatibility - PyTorch only
• export.py compatability all non-TF
• TensorBoard and W&B logging
• CI - Test segmentation: Requires YOLOv5n-seg.pt and COCO128-seg dataset
• export.py TF with TFDetectSegment()
• Fix csv logging
• cpu support
• implement SegmentationModel class (requires retraining all models) to help identify seg models easily on load.
• Fix Logging

🛠️ PR Summary

_{Made with ❤️ by Ultralytics Actions}

🌟 Summary

Extended YOLOv5 object detection with segmentation capabilities.

📊 Key Changes

• Implemented segmentation model, loss, and data loaders
• Included segmentation annotations in datasets
• Extended non-maximum suppression to handle masks
• Updated plotting utilities to visualize segmentations
• Added distinction between object and segmentation metrics

🎯 Purpose & Impact

• Purpose: Enhance YOLOv5 with instance segmentation, enabling the model to differentiate not just object boundaries but also pixel-level details within those boundaries.
• Impact: Enables YOLOv5 to be used in applications requiring detailed object delineation (like medical imaging or autonomous driving) and improves analysis by distinguishing individual object instances within classes. 👩‍⚕️🚗🔍

[glenn-jocher]

@AyushExel precommit shows utils.plots is duplicated (i.e. pasted twice). I'll fix so it matches master. I checked that it doesn't contain any new seg functions.

[AyushExel]

@glenn-jocher I have sent you the trained weights files on slack. Once you add them to assets, we can update predict.py to get the weights for inference

[glenn-jocher]

@AyushExel precommit is all green now! FYI PyCharm with Code -> Reformat will catch 90% of these issues before they get to precommit. precommit can push fixes for some problems but not all of them.

[AyushExel]

@glenn-jocher ohh btw ignore segments/agumentations.py . All models are trained without albumentations installed so we'll probably end up removing that part from this PR.. before that I want to confirm if @Laughing-q has already made some progress on augmentation.

[glenn-jocher]

@AyushExel looking good. I added a few questions to the TODO list above. Also replied in Slack regarding weights.

EDIT: Haven't really had time to go through, just added a few questions based on a quick glance while resolving precommit.

[AyushExel]

@glenn-jocher

torch.deterministic solution for segment?

I think having train.py accept --deterministic should be the easiest way forward? segmentation uses interpolate with bilinear mode which is not deterministic on cuda.

merge non_max_suppression_masks() with non_max_suppression()?

I separated them just to future-proof the repo (was thinking of semantic seg, pose etc) as I don't know how similar or different NMS for other tasks will be. The current functions should be easy to merge. Its your call

remove edict dependency?

Yeah we should remove that. I'll test with native dict and replace it

[glenn-jocher]

@glenn-jocher

torch.deterministic solution for segment?

I think having train.py accept --deterministic should be the easiest way forward? segmentation uses interpolate with bilinear mode which is not deterministic on cuda.

resolved in 8b70e64

merge non_max_suppression_masks() with non_max_suppression()?

I separated them just to future-proof the repo (was thinking of semantic seg, pose etc) as I don't know how similar or different NMS for other tasks will be. The current functions should be easy to merge. Its your call

yes we should definitely merge as much common code as we can to avoid future multiple-maintenance.

remove edict dependency?

Yeah we should remove that. I'll test with native dict and replace it

Perfect.

[Laughing-q]

@glenn-jocher I was testing the segment training with the latest code, but I met this issue when doing val:

After some tests, I found this seems to be related to the export arg in Detect. export equals to True when training which I don't know why. Maybe this is about some updates from master.

Some other code I think can be updated:

1. https://github.com/AyushExel/yolov5/blob/3c8415846bbd016decfa688d2e18bd9a59eb001c/utils/segment/loss.py#L86
   we should switch bilinear to nearest here, because bilinear could generate float number(like 0.5, 0.25...) but we need int number to separate the objects when overlap. This is also one of the reasons that mr=4 perform better than mr=1.
2. We can delete this upsample ops since we'll scale the masks in function plot_images_and_masks when the sizes are different.
   https://github.com/AyushExel/yolov5/blob/3c8415846bbd016decfa688d2e18bd9a59eb001c/segment/train.py#L353
   Btw, I have tested this and the plotting works fine after I deleted these two lines

[AyushExel]

I'm a little occupied today but I think I have a branch with image logging fix. I'll fix it by tomorrow

[glenn-jocher]

@glenn-jocher I was testing the segment training with the latest code, but I met this issue when doing val: After some tests, I found this seems to be related to the export arg in Detect. export equals to True when training which I don't know why. Maybe this is about some updates from master.

Some other code I think can be updated:

1. https://github.com/AyushExel/yolov5/blob/3c8415846bbd016decfa688d2e18bd9a59eb001c/utils/segment/loss.py#L86
   we should switch bilinear to nearest here, because bilinear could generate float number(like 0.5, 0.25...) but we need int number to separate the objects when overlap. This is also one of the reasons that mr=4 perform better than mr=1.
2. We can delete this upsample ops since we'll scale the masks in function plot_images_and_masks when the sizes are different.
   https://github.com/AyushExel/yolov5/blob/3c8415846bbd016decfa688d2e18bd9a59eb001c/segment/train.py#L353
   Btw, I have tested this and the plotting works fine after I deleted these two lines

@Laughing-q resolved now in latest PR. Bug was only on GPU so CI did not pick it up.

@AyushExel we need a GPU CI runner

[glenn-jocher]

@Laughing-q could you open a PR with your proposed changes?

@AyushExel the GPU CI runner won't be easy, was looking around a bit. Jiri did this with a cloud K80's for lightning. Let's not worry about it right now but maybe add to our proposed feature list.

[Laughing-q]

@Laughing-q could you open a PR with your proposed changes?

@glenn-jocher yeah sure, done in there.

[AyushExel]

@Laughing-q thanks! Merged.
@glenn-jocher the plotting issue should be fixed now. The problem here is that validation plot aslo served as plot_one_image function in the detection training. I've changed it so that the plots are now logged only at the last validation as in other epochs plots is set to False.

[glenn-jocher]

@AyushExel @Laughing-q PR is merged. Thank you for your contributions to YOLOv5 🚀 and Vision AI ⭐

[AyushExel]

@glenn-jocher thanks for all the guidance on this and @Laughing-q for building the the initial segmentation models. This is a big step forward.

[SanjayGhanagiri]

This is really great! Cant wait to explore new possibilities with this.