Update `convert_segment_masks_to_yolo_seg` to support custom datasets (#15176)

pull/15185/head^2
Muhammad Rizwan Munawar 3 months ago committed by GitHub
parent 3e4a581c35
commit 4d35458e7c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
  1. 2
      docs/en/usage/simple-utilities.md
  2. 6
      ultralytics/data/converter.py

@ -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)
```

@ -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)

Loading…
Cancel
Save