Update convert_segment_masks_to_yolo_seg to support custom datasets (…

…ultralytics#15176)

docs/en/usage/simple-utilities.md

@@ -63,7 +63,7 @@ The converted masks will be saved in the specified output directory.
 ```python
 from ultralytics.data.converter import convert_segment_masks_to_yolo_seg
 
-# For COCO dataset we have 80 classes
+# The classes here is the total classes in the dataset, for COCO dataset we have 80 classes
 convert_segment_masks_to_yolo_seg(masks_dir="path/to/masks_dir", output_dir="path/to/output_dir", classes=80)
 ```
 

ultralytics/data/converter.py

@@ -344,13 +344,13 @@ def convert_segment_masks_to_yolo_seg(masks_dir, output_dir, classes):
  Args:
  masks_dir (str): The path to the directory where all mask images (png, jpg) are stored.
  output_dir (str): The path to the directory where the converted YOLO segmentation masks will be stored.
- classes (int): Total classes in the dataset i.e for COCO classes=80
+ classes (int): Total classes in the dataset i.e. for COCO classes=80
 
  Example:
  ```python
  from ultralytics.data.converter import convert_segment_masks_to_yolo_seg
 
- # for coco dataset, we have 80 classes
+ # The classes here is the total classes in the dataset, for COCO dataset we have 80 classes
  convert_segment_masks_to_yolo_seg('path/to/masks_directory', 'path/to/output/directory', classes=80)
  ```
 
@@ -373,7 +373,7 @@ def convert_segment_masks_to_yolo_seg(masks_dir, output_dir, classes):
  """
  import os
 
- pixel_to_class_mapping = {i + 1: i for i in range(80)}
+ pixel_to_class_mapping = {i + 1: i for i in range(classes)}
  for mask_filename in os.listdir(masks_dir):
  if mask_filename.endswith(".png"):
  mask_path = os.path.join(masks_dir, mask_filename)